[flang-commits] [flang] [flang][hlfir] Support box in user defined assignments (PR #77578)

[via flang-commits]

https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/77578

When dealing with overlaps in user defined assignments, some entities with descriptors (fir.box) may be saved without descriptors. The current code was replacing the original box entity with the "raw" copy with a simple cast instead of creating a box for the copy. This patch ensures a fir.embox is emitted instead.

>From 8dbdaf698b909a412167ea054e96c188a695f62c Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Wed, 10 Jan 2024 02:35:31 -0800
Subject: [PATCH] [flang][hlfir] Support box in user defined assignments

When dealing with overlaps in user defined assignments, some entities with
descriptors (fir.box) may be saved without descriptors.
The current code was replacing the original box entity with the "raw"
copy with a simple cast instead of creating a box for the copy.
This patch ensures a fir.embox is emitted instead.
---
 .../LowerHLFIROrderedAssignments.cpp          |  8 +++++++
 .../user-defined-assignment.fir               | 23 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp
index c4aee7d39e4a3c..84101353a740a0 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp
@@ -421,6 +421,14 @@ convertToMoldType(mlir::Location loc, fir::FirOpBuilder &builder,
   }
   // Variable to Variable mismatch (e.g., fir.heap<T> vs fir.ref<T>), or value
   // to Value mismatch (e.g. i1 vs fir.logical<4>).
+  if (mlir::isa<fir::BaseBoxType>(mold.getType()) &&
+      !mlir::isa<fir::BaseBoxType>(input.getType())) {
+    // An entity may have have been saved without descriptor while the original
+    // value had a descriptor (e.g., it was not contiguous).
+    auto emboxed = hlfir::convertToBox(loc, builder, input, mold.getType());
+    assert(!emboxed.second && "temp should already be in memory");
+    input = hlfir::Entity{fir::getBase(emboxed.first)};
+  }
   return hlfir::Entity{builder.createConvert(loc, mold.getType(), input)};
 }
 
diff --git a/flang/test/HLFIR/order_assignments/user-defined-assignment.fir b/flang/test/HLFIR/order_assignments/user-defined-assignment.fir
index 521288d0cfa12b..61836b8bcc57eb 100644
--- a/flang/test/HLFIR/order_assignments/user-defined-assignment.fir
+++ b/flang/test/HLFIR/order_assignments/user-defined-assignment.fir
@@ -180,3 +180,26 @@ func.func @test_scalar_forall_overlap(%i: !fir.ref<!fir.array<10xi32>>) {
 // CHECK:    fir.call @logical_value_to_numeric(%[[VAL_32]], %[[VAL_33]]) : (!fir.ref<i32>, !fir.logical<4>) -> ()
 // CHECK:  }
 // CHECK:  fir.freemem %[[VAL_15]] : !fir.heap<!fir.array<?xi1>>
+
+func.func @test_saved_scalar_box(%arg0: !fir.box<!fir.type<sometype>>, %arg1: !fir.class<!fir.type<sometype>>) {
+  hlfir.region_assign {
+    hlfir.yield %arg0 : !fir.box<!fir.type<sometype>>
+  } to {
+    hlfir.yield %arg1 : !fir.class<!fir.type<sometype>>
+  } user_defined_assign  (%arg2: !fir.box<!fir.type<sometype>>) to (%arg3: !fir.class<!fir.type<sometype>>) {
+    fir.call @user_assign_box(%arg3, %arg2) : (!fir.class<!fir.type<sometype>>, !fir.box<!fir.type<sometype>>) -> ()
+  }
+  return
+}
+func.func private @user_assign_box(!fir.class<!fir.type<sometype>>, !fir.box<!fir.type<sometype>>) -> ()
+
+// CHECK-LABEL:   func.func @test_saved_scalar_box(
+// CHECK-SAME:                    %[[VAL_0:.*]]: !fir.box<!fir.type<sometype>>,
+// CHECK-SAME:                    %[[VAL_1:.*]]: !fir.class<!fir.type<sometype>>) {
+// CHECK:           %[[VAL_2:.*]] = hlfir.as_expr %[[VAL_0]] : (!fir.box<!fir.type<sometype>>) -> !hlfir.expr<!fir.type<sometype>>
+// CHECK:           %[[VAL_3:.*]]:3 = hlfir.associate %[[VAL_2]]
+// CHECK:           %[[VAL_4:.*]] = fir.embox %[[VAL_3]]#1 : (!fir.ref<!fir.type<sometype>>) -> !fir.box<!fir.type<sometype>>
+// CHECK:           fir.call @user_assign_box(%[[VAL_1]], %[[VAL_4]]) : (!fir.class<!fir.type<sometype>>, !fir.box<!fir.type<sometype>>) -> ()
+// CHECK:           hlfir.end_associate %[[VAL_3]]#1, %[[VAL_3]]#2 : !fir.ref<!fir.type<sometype>>, i1
+// CHECK:           return
+// CHECK:         }