[flang] Fix optimization of array assignments after #146408 (#147371)

Host associated variables were not being handled properly.
For array references, get the fixed shape extents from the value
type instead, that works correctly in all cases.

Author: luporl

flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp

@@ -805,13 +805,12 @@ llvm::LogicalResult BroadcastAssignBufferization::matchAndRewrite(
                                                shape, /*slice=*/mlir::Value{});
     } else {
       // Array references must have fixed shape, when used in assignments.
+      auto refTy = mlir::cast<fir::ReferenceType>(lhs.getType());
+      auto seqTy = mlir::cast<fir::SequenceType>(refTy.getElementType());
+      llvm::ArrayRef<int64_t> fixedShape = seqTy.getShape();
       int64_t flatExtent = 1;
-      for (const mlir::Value &extent : extents) {
-        mlir::Operation *op = extent.getDefiningOp();
-        assert(op && "no defining operation for constant array extent");
-        flatExtent *= fir::toInt(mlir::cast<mlir::arith::ConstantOp>(*op));
-      }
-
+      for (int64_t extent : fixedShape)
+        flatExtent *= extent;
       flatArrayType =
           fir::ReferenceType::get(fir::SequenceType::get({flatExtent}, eleTy));
       flatArray = builder.createConvert(loc, flatArrayType, flatArray);

flang/test/HLFIR/opt-scalar-assign.fir

@@ -155,3 +155,32 @@ func.func @_QPtest6(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>> {f
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
+
+func.func @_QQmain() {
+  return
+}
+
+func.func private @_QFPtest7(%arg0: !fir.ref<tuple<!fir.box<!fir.array<2x2xi32>>>> {fir.host_assoc}) attributes {fir.host_symbol = @_QQmain, llvm.linkage = #llvm.linkage<internal>} {
+  %0 = fir.dummy_scope : !fir.dscope
+  %c0_i32 = arith.constant 0 : i32
+  %1 = fir.coordinate_of %arg0, %c0_i32 : (!fir.ref<tuple<!fir.box<!fir.array<2x2xi32>>>>, i32) -> !fir.ref<!fir.box<!fir.array<2x2xi32>>>
+  %2 = fir.load %1 : !fir.ref<!fir.box<!fir.array<2x2xi32>>>
+  %3 = fir.box_addr %2 : (!fir.box<!fir.array<2x2xi32>>) -> !fir.ref<!fir.array<2x2xi32>>
+  %c0 = arith.constant 0 : index
+  %4:3 = fir.box_dims %2, %c0 : (!fir.box<!fir.array<2x2xi32>>, index) -> (index, index, index)
+  %c1 = arith.constant 1 : index
+  %5:3 = fir.box_dims %2, %c1 : (!fir.box<!fir.array<2x2xi32>>, index) -> (index, index, index)
+  %6 = fir.shape %4#1, %5#1 : (index, index) -> !fir.shape<2>
+  %7:2 = hlfir.declare %3(%6) {fortran_attrs = #fir.var_attrs<host_assoc>, uniq_name = "_QFEa"} : (!fir.ref<!fir.array<2x2xi32>>, !fir.shape<2>) -> (!fir.ref<!fir.array<2x2xi32>>, !fir.ref<!fir.array<2x2xi32>>)
+  %c0_i32_0 = arith.constant 0 : i32
+  hlfir.assign %c0_i32_0 to %7#0 : i32, !fir.ref<!fir.array<2x2xi32>>
+  return
+}
+
+// CHECK-LABEL:   func.func private @_QFPtest7({{.*}}) {{.*}} {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 0 : i32
+// CHECK:           fir.do_loop %[[VAL_1:.*]] = %{{.*}} to %{{.*}} step %{{.*}} unordered {
+// CHECK:             %[[VAL_2:.*]] = hlfir.designate %{{.*}} (%[[VAL_1]]) : (!fir.ref<!fir.array<4xi32>>, index) -> !fir.ref<i32>
+// CHECK:             hlfir.assign %[[VAL_0]] to %[[VAL_2]] : i32, !fir.ref<i32>
+// CHECK:           }
+// CHECK:         }