[Flang] Legalize the hlfir.exactly_once operation

[ceseo]

When reshape is used with allocatable components in derived types within a where construct, a hlfir.exactly_once operation is generated. This commit implements the conversion pattern for it in the HLFIR bufferization pass.

Fixes #130532

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Carlos Seo (ceseo)

Changes

When reshape is used with allocatable components in derived types within a where construct, a hlfir.exactly_once operation is generated. This commit implements the conversion pattern for it in the HLFIR bufferization pass.

Fixes #130532

---

Full diff: https://github.com/llvm/llvm-project/pull/152202.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp (+42-3)
• (added) flang/test/HLFIR/bufferize-exactly_once.f90 (+20)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
index 9109f2b331567..327e385a11eec 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
@@ -29,6 +29,7 @@
 #include "flang/Optimizer/OpenMP/Passes.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
@@ -880,6 +881,42 @@ struct EvaluateInMemoryOpConversion
   }
 };
 
+struct ExactlyOnceOpConversion
+    : public mlir::OpConversionPattern<hlfir::ExactlyOnceOp> {
+  using mlir::OpConversionPattern<hlfir::ExactlyOnceOp>::OpConversionPattern;
+  explicit ExactlyOnceOpConversion(mlir::MLIRContext *ctx)
+      : mlir::OpConversionPattern<hlfir::ExactlyOnceOp>{ctx} {}
+  llvm::LogicalResult
+  matchAndRewrite(hlfir::ExactlyOnceOp exactlyOnce, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    fir::FirOpBuilder builder(rewriter, exactlyOnce.getOperation());
+
+    // The body of the exactly_once op may contain operations that will require
+    // to be translated
+    HLFIRListener listener{builder, rewriter};
+    builder.setListener(&listener);
+
+    // Clone the exactly_once region body and convert it inline
+    mlir::Region &region = exactlyOnce.getBody();
+    mlir::Block &block = region.front();
+    mlir::IRMapping mapper;
+
+    // Clone all operations except the terminator
+    for (auto &op : block.without_terminator()) {
+      rewriter.clone(op, mapper);
+    }
+
+    // Get the yielded value and replace the exactly_once operation with it.
+    // Later we will register the pattern and mark the ExactlyOnceOp as
+    // illegal (see below).
+    auto yield = mlir::cast<hlfir::YieldOp>(block.getTerminator());
+    mlir::Value yieldedValue = mapper.lookupOrDefault(yield.getEntity());
+
+    rewriter.replaceOp(exactlyOnce, yieldedValue);
+    return mlir::success();
+  }
+};
+
 class BufferizeHLFIR : public hlfir::impl::BufferizeHLFIRBase<BufferizeHLFIR> {
 public:
   using BufferizeHLFIRBase<BufferizeHLFIR>::BufferizeHLFIRBase;
@@ -899,8 +936,9 @@ class BufferizeHLFIR : public hlfir::impl::BufferizeHLFIRBase<BufferizeHLFIR> {
                     AssociateOpConversion, CharExtremumOpConversion,
                     ConcatOpConversion, DestroyOpConversion,
                     EndAssociateOpConversion, EvaluateInMemoryOpConversion,
-                    NoReassocOpConversion, SetLengthOpConversion,
-                    ShapeOfOpConversion, GetLengthOpConversion>(context);
+                    ExactlyOnceOpConversion, NoReassocOpConversion,
+                    SetLengthOpConversion, ShapeOfOpConversion,
+                    GetLengthOpConversion>(context);
     patterns.insert<ElementalOpConversion>(context, optimizeEmptyElementals);
     mlir::ConversionTarget target(*context);
     // Note that YieldElementOp is not marked as an illegal operation.
@@ -909,7 +947,8 @@ class BufferizeHLFIR : public hlfir::impl::BufferizeHLFIRBase<BufferizeHLFIR> {
     // survives this pass, the verifier will detect it because it has to be
     // a child of ElementalOp and ElementalOp's are explicitly illegal.
     target.addIllegalOp<hlfir::ApplyOp, hlfir::AssociateOp, hlfir::ElementalOp,
-                        hlfir::EndAssociateOp, hlfir::SetLengthOp>();
+                        hlfir::EndAssociateOp, hlfir::ExactlyOnceOp,
+                        hlfir::SetLengthOp>();
 
     target.markUnknownOpDynamicallyLegal([](mlir::Operation *op) {
       return llvm::all_of(op->getResultTypes(),
diff --git a/flang/test/HLFIR/bufferize-exactly_once.f90 b/flang/test/HLFIR/bufferize-exactly_once.f90
new file mode 100644
index 0000000000000..be97a70113653
--- /dev/null
+++ b/flang/test/HLFIR/bufferize-exactly_once.f90
@@ -0,0 +1,20 @@
+! RUN: bbc -emit-hlfir %s -o - | FileCheck %s
+
+program main
+  call test06()
+  print *,'pass'
+end program main
+
+subroutine test06()
+  type ty1
+     integer ,allocatable :: a(:,:,:)
+  end type ty1
+  type(ty1) :: str(1)
+  integer ,allocatable :: b(:,:,:)
+  allocate(str(1)%a(1,1,1),b(1,1,1))
+  b=1
+  write(6,*) "b                                 = ", b
+  write(6,*) "reshape((/(b,jj=1,1)/),(/1,1,1/)) = ", reshape((/(b,jj=1,1)/),(/1,1,1/))
+  where ((/.true./)) str=(/(ty1(reshape((/(b,jj=1,1)/),(/1,1,1/))),ii=1,1)/)
+  ! CHECK: hlfir.exactly_once : !hlfir.expr<1x1x1xi32>
+end subroutine test06

[vzakhari]

Thank you for the changes, but I think this is not the right place to fix this. I believe all hlfir.exactly_once operations are expected to disappear after LowerHLFIROrderedAssignments pass. Let me look deeper into it.

[vzakhari]

I posted a simpler reproducer at #130532 (comment). I think the nesting of hlfir.exactly_once confusing the ordered-assignment pass is the culprit here.

[ceseo]

I'll close this and create another one for the proper fix.