@@ -42,6 +42,18 @@ namespace hlfir {
4242
4343#define DEBUG_TYPE " opt-bufferization"
4444
45+ // An engineering option to allow treating hlfir.assign with realloc
46+ // attribute as never requiring (re)allocation of the LHS.
47+ // Setting it to true may result in incorrect code, so it is present
48+ // just for quicker benchmarking of apps that may benefit
49+ // from multiversioning such hlfir.assign operations under
50+ // the dynamic checks of the type/shape/allocation status.
51+ static llvm::cl::opt<bool > assumeNoLhsReallocation (
52+ " flang-assume-no-lhs-reallocation" ,
53+ llvm::cl::desc (" Assume that hlfir.assign never (re)allocates the LHS, i.e. "
54+ " that the LHS and RHS are always conformant" ),
55+ llvm::cl::init(false ));
56+
4557namespace {
4658
4759// / This transformation should match in place modification of arrays.
@@ -462,16 +474,20 @@ ElementalAssignBufferization::findMatch(hlfir::ElementalOp elemental) {
462474 // the incoming expression
463475 match.array = match.assign .getLhs ();
464476 mlir::Type arrayType = mlir::dyn_cast<fir::SequenceType>(
465- fir::unwrapPassByRefType (match.array .getType ()));
466- if (!arrayType)
477+ hlfir::getFortranElementOrSequenceType (match.array .getType ()));
478+ if (!arrayType) {
479+ LLVM_DEBUG (llvm::dbgs () << " AssignOp's result is not an array\n " );
467480 return std::nullopt ;
481+ }
468482
469483 // require that the array elements are trivial
470484 // TODO: this is just to make the pass easier to think about. Not an inherent
471485 // limitation
472486 mlir::Type eleTy = hlfir::getFortranElementType (arrayType);
473- if (!fir::isa_trivial (eleTy))
487+ if (!fir::isa_trivial (eleTy)) {
488+ LLVM_DEBUG (llvm::dbgs () << " AssignOp's data type is not trivial\n " );
474489 return std::nullopt ;
490+ }
475491
476492 // The array must have the same shape as the elemental.
477493 //
@@ -485,8 +501,10 @@ ElementalAssignBufferization::findMatch(hlfir::ElementalOp elemental) {
485501 // there is no reallocation of the lhs due to the assignment.
486502 // We can probably try generating multiple versions of the code
487503 // with checking for the shape match, length parameters match, etc.
488- if (match.assign .getRealloc ())
504+ if (match.assign .isAllocatableAssignment () && !assumeNoLhsReallocation) {
505+ LLVM_DEBUG (llvm::dbgs () << " AssignOp may involve (re)allocation of LHS\n " );
489506 return std::nullopt ;
507+ }
490508
491509 // the transformation wants to apply the elemental in a do-loop at the
492510 // hlfir.assign, check there are no effects which make this unsafe
@@ -606,6 +624,8 @@ llvm::LogicalResult ElementalAssignBufferization::matchAndRewrite(
606624
607625 // create the loop at the assignment
608626 builder.setInsertionPoint (match->assign );
627+ hlfir::Entity arrayObj = hlfir::derefPointersAndAllocatables (
628+ loc, builder, hlfir::Entity{match->array });
609629
610630 // Generate a loop nest looping around the hlfir.elemental shape and clone
611631 // hlfir.elemental region inside the inner loop
@@ -619,8 +639,8 @@ llvm::LogicalResult ElementalAssignBufferization::matchAndRewrite(
619639 rewriter.eraseOp (yield);
620640
621641 // Assign the element value to the array element for this iteration.
622- auto arrayElement = hlfir::getElementAt (
623- loc, builder, hlfir::Entity{match-> array } , loopNest.oneBasedIndices );
642+ auto arrayElement =
643+ hlfir::getElementAt ( loc, builder, arrayObj , loopNest.oneBasedIndices );
624644 builder.create <hlfir::AssignOp>(
625645 loc, elementValue, arrayElement, /* realloc=*/ false ,
626646 /* keep_lhs_length_if_realloc=*/ false , match->assign .getTemporaryLhs ());
0 commit comments