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

[via flang-commits]

================
@@ -125,9 +136,174 @@ 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();
+    mlir::Type boxType = fir::unwrapRefType(lhsType);
+    auto boxBaseType = mlir::cast<fir::BaseBoxType>(boxType);
+    mlir::Type baseTy = boxBaseType.getEleTy();
+    if (auto heapTy = mlir::dyn_cast<fir::HeapType>(baseTy))
+      baseTy = heapTy.getEleTy();
+
+    // Create MutableBoxValue - for trivial types, no length params needed
+    fir::MutableBoxValue mutableBox(lhsBoxRef, /*nonDeferredParams=*/{},
+                                    /*mutableProperties=*/{});
+
+    // Generate ones for lower bounds (Fortran default)
+    mlir::Value one =
+        builder.createIntegerConstant(loc, builder.getIndexType(), 1);
+    llvm::SmallVector<mlir::Value> lbounds;
+    for (size_t i = 0; i < rhsExtents.size(); ++i)
+      lbounds.push_back(one);
+
+    // Use genReallocIfNeeded to handle allocation/reallocation properly.
+    // This implements Fortran 10.2.1.3 point 3:
+    // - If not allocated, allocate with RHS shape
+    // - If allocated with same shape, keep existing allocation
+    // - If allocated with different shape, reallocate
+    //
+    // The storage handler callback performs the actual assignment loop.
+    bool useWorkshare = flangomp::shouldUseWorkshareLowering(assign);
+    auto storageHandler = [&](fir::ExtendedValue storage) {
+      // Create an hlfir.declare for the storage to get a proper Entity.
+      // This is necessary because hlfir::Entity requires a value from an
+      // HLFIR operation, not a raw pointer.
+      auto declare = hlfir::genDeclare(loc, builder, storage, ".tmp.assign",
----------------
jeanPerier wrote:

"tmp.assign" may be a bit misleading since this may be the original LHS storage if it is already allocated properly.
It is in that regards a bit tricky because it should likely carry the attributes of the original variable to avoid issues.

For instance, if the LHS has the target attribute and the RHS expression is using a POINTER, later optimization may assume the hlfir.expr is safe to inline because the hlfir.declare affected by the assign of the hlfir.expr does not have the TARGET attribute (it is likely current optimizations would not do such inlining currently, but this is being discussed [in another PR](https://github.com/llvm/llvm-project/pull/186916#issuecomment-4071827283)).

You could also drop the hlfir.declare altogether to avoid these kind of issues, a fir.box is always a valid hlfir::Entity, so you could just call `hlfir::Entity lhsEntity{fir::getBase(fir::factory::createBoxValue(builder, loc, storage));` instead.

That way, the alias analysis code will have to go up the chain to reason about this memory (and will currently fail and return may alias to anything).




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