From 86e6a8c33a410c9219806ee87acff2ea011de979 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier@nvidia.com>
Date: Fri, 27 Jan 2023 12:49:55 +0100
Subject: [PATCH] [flang][hlfir] Handle scalar to array in hlfir.assign
 codegen.

The scalar must be placed in memory before creating descriptors and
calling the runtime assignment API.

Differential Revision: https://reviews.llvm.org/D142698
---
 .../Optimizer/HLFIR/Transforms/ConvertToFIR.cpp    | 16 +++++++++++-
 flang/test/HLFIR/assign-codegen.fir                | 29 ++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
index a71463d..51cca5c 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
@@ -92,12 +92,26 @@ public:
            "variable to fir::ExtendedValue must not require cleanup");
 
     if (lhs.isArray()) {
+      const bool rhsIsValue = fir::isa_trivial(fir::getBase(rhsExv).getType());
+      if (rhsIsValue) {
+        // createBox can only be called for fir::ExtendedValue that are
+        // already in memory. Place the integer/real/complex/logical scalar
+        // in memory (convert to the LHS type so that i1 are allocated in
+        // a proper Fortran logical storage).
+        mlir::Type lhsValueType = lhs.getFortranElementType();
+        mlir::Value rhsVal =
+            builder.createConvert(loc, lhsValueType, fir::getBase(rhsExv));
+        mlir::Value temp = builder.create<fir::AllocaOp>(loc, lhsValueType);
+        builder.create<fir::StoreOp>(loc, rhsVal, temp);
+        rhsExv = temp;
+      }
+
       // Use the runtime for simplicity. An optimization pass will be added to
       // inline array assignment when profitable.
       auto to = fir::getBase(builder.createBox(loc, lhsExv));
       auto from = fir::getBase(builder.createBox(loc, rhsExv));
       bool cleanUpTemp = false;
-      if (mayAlias(rhs, lhs))
+      if (!rhsIsValue && mayAlias(rhs, lhs))
         std::tie(from, cleanUpTemp) = genTempFromSourceBox(loc, builder, from);
 
       auto toMutableBox = builder.createTemporary(loc, to.getType());
diff --git a/flang/test/HLFIR/assign-codegen.fir b/flang/test/HLFIR/assign-codegen.fir
index b7d0447..91ca33d 100644
--- a/flang/test/HLFIR/assign-codegen.fir
+++ b/flang/test/HLFIR/assign-codegen.fir
@@ -155,3 +155,32 @@ func.func @array(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.ref<!fir.array<
 // CHECK:  %[[VAL_29:.*]] = fir.call @_FortranAAssign(%[[VAL_26]], %[[VAL_27]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.ref<i8>, i32) -> none
 // CHECK:  %[[VAL_30:.*]] = fir.box_addr %[[VAL_22]] : (!fir.box<!fir.heap<!fir.array<100xi32>>>) -> !fir.heap<!fir.array<100xi32>>
 // CHECK:  fir.freemem %[[VAL_30]] : !fir.heap<!fir.array<100xi32>>
+
+
+func.func @test_scalar_to_array(%lhs: !fir.box<!fir.array<?xi32>>, %rhs: i32) {
+  hlfir.assign %rhs to %lhs : i32, !fir.box<!fir.array<?xi32>>
+  return
+}
+// CHECK-LABEL:   func.func @test_scalar_to_array(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: i32) {
+// CHECK:  %[[VAL_2:.*]] = fir.alloca !fir.box<!fir.array<?xi32>>
+// CHECK:  %[[VAL_5:.*]] = fir.alloca i32
+// CHECK:  fir.store %[[VAL_1]] to %[[VAL_5]] : !fir.ref<i32>
+// CHECK:  %[[VAL_6:.*]] = fir.embox %[[VAL_5]] : (!fir.ref<i32>) -> !fir.box<i32>
+// CHECK:  fir.store %[[VAL_0]] to %[[VAL_2]] : !fir.ref<!fir.box<!fir.array<?xi32>>>
+// CHECK:  %[[VAL_10:.*]] = fir.convert %[[VAL_2]] : (!fir.ref<!fir.box<!fir.array<?xi32>>>) -> !fir.ref<!fir.box<none>>
+// CHECK:  %[[VAL_11:.*]] = fir.convert %[[VAL_6]] : (!fir.box<i32>) -> !fir.box<none>
+// CHECK:  %[[VAL_13:.*]] = fir.call @_FortranAAssign(%[[VAL_10]], %[[VAL_11]]
+
+
+func.func @test_i1_scalar_to_array(%lhs: !fir.box<!fir.array<?x!fir.logical<4>>>, %rhs: i1) {
+  hlfir.assign %rhs to %lhs : i1, !fir.box<!fir.array<?x!fir.logical<4>>>
+  return
+}
+// CHECK-LABEL:   func.func @test_i1_scalar_to_array(
+// CHECK:  %[[VAL_5:.*]] = fir.convert %{{.*}} : (i1) -> !fir.logical<4>
+// CHECK:  %[[VAL_6:.*]] = fir.alloca !fir.logical<4>
+// CHECK:  %[[VAL_7:.*]] = fir.embox %[[VAL_6]] : (!fir.ref<!fir.logical<4>>) -> !fir.box<!fir.logical<4>>
+// CHECK:  %[[VAL_12:.*]] = fir.convert %[[VAL_7]] : (!fir.box<!fir.logical<4>>) -> !fir.box<none>
+// CHECK:  %[[VAL_14:.*]] = fir.call @_FortranAAssign(%{{.*}}, %[[VAL_12]]
-- 
2.7.4