[flang][hlfir] Expand array hlfir.assign's.

Expand hlfir.assign with in-memory array RHS and LHS into
a loop nest with element-by-element assignments.
For small arrays this may result in further loop nest unrolling
enabling more value propagation and redundancy elimination.

Note the change in flang/test/HLFIR/opt-bufferization.fir:
the hlfir.assign inside hlfir.elemental gets expanded by the new
pattern.

Depends on D159151

Reviewed By: tblah

Differential Revision: https://reviews.llvm.org/D159246

Author: vzakhari

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

@@ -417,6 +417,102 @@ mlir::LogicalResult BroadcastAssignBufferization::matchAndRewrite(
   return mlir::success();
 }
 
+/// Expand hlfir.assign of array RHS to array LHS into a loop nest
+/// of element-by-element assignments:
+///   hlfir.assign %4 to %5 : !fir.ref<!fir.array<3x3xf32>>,
+///                           !fir.ref<!fir.array<3x3xf32>>
+/// into:
+///   fir.do_loop %arg1 = %c1 to %c3 step %c1 unordered {
+///     fir.do_loop %arg2 = %c1 to %c3 step %c1 unordered {
+///       %6 = hlfir.designate %4 (%arg2, %arg1)  :
+///           (!fir.ref<!fir.array<3x3xf32>>, index, index) -> !fir.ref<f32>
+///       %7 = fir.load %6 : !fir.ref<f32>
+///       %8 = hlfir.designate %5 (%arg2, %arg1)  :
+///           (!fir.ref<!fir.array<3x3xf32>>, index, index) -> !fir.ref<f32>
+///       hlfir.assign %7 to %8 : f32, !fir.ref<f32>
+///     }
+///   }
+///
+/// The transformation is correct only when LHS and RHS do not alias.
+/// This transformation does not support runtime checking for
+/// non-conforming LHS/RHS arrays' shapes currently.
+class VariableAssignBufferization
+    : public mlir::OpRewritePattern<hlfir::AssignOp> {
+private:
+public:
+  using mlir::OpRewritePattern<hlfir::AssignOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(hlfir::AssignOp assign,
+                  mlir::PatternRewriter &rewriter) const override;
+};
+
+mlir::LogicalResult VariableAssignBufferization::matchAndRewrite(
+    hlfir::AssignOp assign, mlir::PatternRewriter &rewriter) const {
+  if (assign.isAllocatableAssignment())
+    return rewriter.notifyMatchFailure(assign, "AssignOp may imply allocation");
+
+  hlfir::Entity rhs{assign.getRhs()};
+  // TODO: ExprType check is here to avoid conflicts with
+  // ElementalAssignBufferization pattern. We need to combine
+  // these matchers into a single one that applies to AssignOp.
+  if (rhs.getType().isa<hlfir::ExprType>())
+    return rewriter.notifyMatchFailure(assign, "RHS is not in memory");
+
+  if (!rhs.isArray())
+    return rewriter.notifyMatchFailure(assign,
+                                       "AssignOp's RHS is not an array");
+
+  mlir::Type rhsEleTy = rhs.getFortranElementType();
+  if (!fir::isa_trivial(rhsEleTy))
+    return rewriter.notifyMatchFailure(
+        assign, "AssignOp's RHS data type is not trivial");
+
+  hlfir::Entity lhs{assign.getLhs()};
+  if (!lhs.isArray())
+    return rewriter.notifyMatchFailure(assign,
+                                       "AssignOp's LHS is not an array");
+
+  mlir::Type lhsEleTy = lhs.getFortranElementType();
+  if (!fir::isa_trivial(lhsEleTy))
+    return rewriter.notifyMatchFailure(
+        assign, "AssignOp's LHS data type is not trivial");
+
+  if (lhsEleTy != rhsEleTy)
+    return rewriter.notifyMatchFailure(assign,
+                                       "RHS/LHS element types mismatch");
+
+  fir::AliasAnalysis aliasAnalysis;
+  mlir::AliasResult aliasRes = aliasAnalysis.alias(lhs, rhs);
+  if (!aliasRes.isNo()) {
+    LLVM_DEBUG(llvm::dbgs() << "VariableAssignBufferization:\n"
+                            << "\tLHS: " << lhs << "\n"
+                            << "\tRHS: " << rhs << "\n"
+                            << "\tALIAS: " << aliasRes << "\n");
+    return rewriter.notifyMatchFailure(assign, "RHS/LHS may alias");
+  }
+
+  mlir::Location loc = assign->getLoc();
+  fir::FirOpBuilder builder(rewriter, assign.getOperation());
+  builder.setInsertionPoint(assign);
+  rhs = hlfir::derefPointersAndAllocatables(loc, builder, rhs);
+  lhs = hlfir::derefPointersAndAllocatables(loc, builder, lhs);
+  mlir::Value shape = hlfir::genShape(loc, builder, lhs);
+  llvm::SmallVector<mlir::Value> extents =
+      hlfir::getIndexExtents(loc, builder, shape);
+  hlfir::LoopNest loopNest =
+      hlfir::genLoopNest(loc, builder, extents, /*isUnordered=*/true);
+  builder.setInsertionPointToStart(loopNest.innerLoop.getBody());
+  auto rhsArrayElement =
+      hlfir::getElementAt(loc, builder, rhs, loopNest.oneBasedIndices);
+  rhsArrayElement = hlfir::loadTrivialScalar(loc, builder, rhsArrayElement);
+  auto lhsArrayElement =
+      hlfir::getElementAt(loc, builder, lhs, loopNest.oneBasedIndices);
+  builder.create<hlfir::AssignOp>(loc, rhsArrayElement, lhsArrayElement);
+  rewriter.eraseOp(assign);
+  return mlir::success();
+}
+
 class OptimizedBufferizationPass
     : public hlfir::impl::OptimizedBufferizationBase<
           OptimizedBufferizationPass> {
@@ -438,6 +534,7 @@ class OptimizedBufferizationPass
     // This requires small code reordering in ElementalAssignBufferization.
     patterns.insert<ElementalAssignBufferization>(context);
     patterns.insert<BroadcastAssignBufferization>(context);
+    patterns.insert<VariableAssignBufferization>(context);
 
     if (mlir::failed(mlir::applyPatternsAndFoldGreedily(
             func, std::move(patterns), config))) {

flang/test/HLFIR/opt-bufferization.fir

@@ -251,7 +251,7 @@ func.func @write(%arg: !fir.ref<!fir.array<42xi32>>, %arg1: !fir.ref<!fir.array<
 // CHECK:           %[[VAL_5B:.*]]:2 = hlfir.declare %[[ARG_1]](%[[VAL_4]]) {uniq_name = "array2"} : (!fir.ref<!fir.array<42xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<42xi32>>, !fir.ref<!fir.array<42xi32>>)
 // CHECK:           %[[VAL_6:.*]] = hlfir.elemental %[[VAL_4]] unordered : (!fir.shape<1>) -> !hlfir.expr<42xi32> {
 // CHECK:           ^bb0(%[[VAL_7:.*]]: index):
-// CHECK:             hlfir.assign %[[VAL_5B]]#0 to %[[VAL_5]]#0 : !fir.ref<!fir.array<42xi32>>, !fir.ref<!fir.array<42xi32>>
+// CHECK:             hlfir.assign
 // CHECK:             %[[VAL_8:.*]] = hlfir.designate %[[VAL_5]]#0 (%[[VAL_7]])  : (!fir.ref<!fir.array<42xi32>>, index) -> !fir.ref<i32>
 // CHECK:             %[[VAL_9:.*]] = fir.load %[[VAL_8]] : !fir.ref<i32>
 // CHECK:             %[[VAL_10:.*]] = arith.subi %[[VAL_9]], %[[VAL_1]] : i32