[flang-commits] [flang] [Flang] - Add optional inlining of allocatable assignments with hlfir.expr RHS (PR #186880)

[Eugene Epshteyn via flang-commits]

================
@@ -125,9 +136,158 @@ class InlineHLFIRAssignConversion
   }
 };
 
+/// Expand hlfir.assign of hlfir.expr RHS to allocatable LHS.
+/// When RHS is an hlfir.expr (e.g., from hlfir.elemental), there is no
+/// aliasing concern because expressions don't represent memory locations.
+/// This allows us to inline the assignment even for allocatables.
+///
+/// The generated code:
+/// 1. Gets the shape from the RHS expression
+/// 2. Uses genReallocIfNeeded to handle allocation/reallocation properly
+/// 3. Generates a loop nest to assign elements (via storage handler callback)
+/// 4. Finalizes the reallocation
+///
+/// Example transformation for: allocatable_array = elemental_expr
+///   hlfir.assign %expr to %alloc realloc : !hlfir.expr<?xf64>,
+///                                          !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+/// into:
+///   // Check allocation status and reallocate if needed
+///   // ... (genReallocIfNeeded handles this) ...
+///   // Loop over elements
+///   fir.do_loop %i = %c1 to %extent step %c1 unordered {
+///     %rhs_val = hlfir.apply %expr, %i : ...
+///     %lhs_elem = hlfir.designate %lhs_box (%i) : ...
+///     hlfir.assign %rhs_val to %lhs_elem : f64, !fir.ref<f64>
+///   }
+class InlineAllocatableExprAssignConversion
+    : public mlir::OpRewritePattern<hlfir::AssignOp> {
+public:
+  using mlir::OpRewritePattern<hlfir::AssignOp>::OpRewritePattern;
+
+  llvm::LogicalResult
+  matchAndRewrite(hlfir::AssignOp assign,
+                  mlir::PatternRewriter &rewriter) const override {
+    // This pattern only handles allocatable assignments
+    if (!assign.isAllocatableAssignment())
+      return rewriter.notifyMatchFailure(
+          assign, "AssignOp is not an allocatable assignment");
+
+    hlfir::Entity rhs{assign.getRhs()};
+    hlfir::Entity lhs{assign.getLhs()};
+
+    // RHS must be an hlfir.expr (this is the key condition - no aliasing)
+    if (!mlir::isa<hlfir::ExprType>(rhs.getType()))
+      return rewriter.notifyMatchFailure(
+          assign,
+          "RHS is not an hlfir.expr - cannot inline allocatable assign");
+
+    // RHS must be an array
+    if (!rhs.isArray())
+      return rewriter.notifyMatchFailure(assign,
+                                         "AssignOp's RHS is not an array");
+
+    // Check element types are trivial and match
+    mlir::Type rhsEleTy = rhs.getFortranElementType();
+    if (!fir::isa_trivial(rhsEleTy))
+      return rewriter.notifyMatchFailure(
+          assign, "AssignOp's RHS data type is not trivial");
+
+    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");
+
+    // LHS must be a reference to a box (allocatable)
+    mlir::Type lhsType = lhs.getType();
+    if (!fir::isBoxAddress(lhsType))
+      return rewriter.notifyMatchFailure(assign,
+                                         "LHS is not a reference to a box");
+
+    LLVM_DEBUG(llvm::dbgs()
+               << "InlineHLFIRAssign: inlining allocatable expr assignment\n");
+
+    mlir::Location loc = assign->getLoc();
+    fir::FirOpBuilder builder(rewriter, assign.getOperation());
+    builder.setInsertionPoint(assign);
+
+    // Get the shape of the RHS expression
+    mlir::Value rhsShape = hlfir::genShape(loc, builder, rhs);
+    llvm::SmallVector<mlir::Value> rhsExtents =
+        hlfir::getIndexExtents(loc, builder, rhsShape);
+
+    // Create a MutableBoxValue for the LHS allocatable
+    mlir::Value lhsBoxRef = lhs.getFirBase();
+
+    // Create MutableBoxValue - for trivial types, no length params needed
+    fir::MutableBoxValue mutableBox(lhsBoxRef, /*lenParamaters=*/{},
----------------
eugeneepshteyn wrote:

Small typo fix

https://github.com/llvm/llvm-project/pull/186880