[Mlir-commits] [mlir] [OpenACC] Replace terminators with scf.yield in wrapMultiBlockRegionWithSCFExecuteRegion (PR #183758)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-openacc

@llvm/pr-subscribers-mlir

Author: Delaram Talaashrafi (delaram-talaashrafi)

<details>
<summary>Changes</summary>

When wrapping a multi-block region in `scf.execute_region`, replace both `func::ReturnOp` and `acc::YieldOp` in all the blocks with `scf.yield` so the region has a valid SCF terminator.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp (+17-12) 
- (modified) mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp (+76) 


``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
index b76d3f338924c..65471d2f40867 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
@@ -14,6 +14,7 @@
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
@@ -157,21 +158,25 @@ wrapMultiBlockRegionWithSCFExecuteRegion(Region &region, IRMapping &mapping,
   rewriter.cloneRegionBefore(region, exeRegionOp.getRegion(),
                              exeRegionOp.getRegion().end(), mapping);
 
-  // Find and replace the ACC terminator with scf.yield
-  Operation *terminator = exeRegionOp.getRegion().back().getTerminator();
-  if (auto yieldOp = dyn_cast<acc::YieldOp>(terminator)) {
-    if (yieldOp.getNumOperands() > 0) {
-      region.getParentOp()->emitError(
-          "acc.loop with results not yet supported");
-      return nullptr;
+  // Find and replace the ACC terminators with scf.yield in each block
+  for (Block &block : exeRegionOp.getRegion().getBlocks()) {
+    if (block.empty()) {
+      continue;
     }
-  } else if (!isa<acc::TerminatorOp>(terminator)) {
-    llvm_unreachable("unexpected terminator in ACC region");
+    Operation *blockTerminator = block.getTerminator();
+    if (isa<func::ReturnOp>(*blockTerminator) ||
+        isa<acc::YieldOp>(*blockTerminator)) {
+      if (blockTerminator->getNumOperands()) {
+        region.getParentOp()->emitError(
+            "acc.loop with results not yet supported");
+        return nullptr;
+      }
+      rewriter.setInsertionPointToEnd(&block);
+      (void)scf::YieldOp::create(rewriter, blockTerminator->getLoc());
+      blockTerminator->erase();
+    } 
   }
 
-  rewriter.eraseOp(terminator);
-  rewriter.setInsertionPointToEnd(&exeRegionOp.getRegion().back());
-  scf::YieldOp::create(rewriter, loc);
   return exeRegionOp;
 }
 
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
index 83c162e32dedc..69b3dd11e47db 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
@@ -649,6 +649,82 @@ TEST_F(OpenACCUtilsLoopTest,
   EXPECT_EQ(mulCount, 1u);
 }
 
+TEST_F(OpenACCUtilsLoopTest,
+       WrapMultiBlockRegionWithSCFExecuteRegionIssue2257LoopPattern) {
+  auto [module, funcOp] = createModuleWithFunc();
+
+  OwningOpRef<acc::ParallelOp> parallelOp =
+      acc::ParallelOp::create(b, loc, TypeRange{}, ValueRange{});
+  Region &region = parallelOp->getRegion();
+  // Block order as in all.mlir: ^bb0 entry, ^bb1 header, ^bb2 exit, ^bb3 body
+  Block *entry = b.createBlock(&region, region.begin());
+  Block *header = b.createBlock(&region, region.end());
+  Block *exitBlock = b.createBlock(&region, region.end());
+  Block *bodyBlock = b.createBlock(&region, region.end());
+
+  // ^bb0 (entry): setup then cf.br ^bb1
+  b.setInsertionPointToEnd(entry);
+  Value c1 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
+  arith::AddIOp::create(b, loc, c1, c1); // some op like fir.store in all.mlir
+  cf::BranchOp::create(b, loc, header);
+
+  // ^bb1 (header): 2 preds (entry, body). cond_br to exit or body
+  b.setInsertionPointToEnd(header);
+  Value cond =
+      arith::ConstantOp::create(b, loc, b.getI1Type(), b.getBoolAttr(false));
+  cf::CondBranchOp::create(b, loc, cond, exitBlock, bodyBlock);
+
+  // ^bb2 (exit): return — use acc.yield (replaced with scf.yield by the util)
+  b.setInsertionPointToEnd(exitBlock);
+  acc::YieldOp::create(b, loc);
+
+  // ^bb3 (body): body ops then cf.br ^bb1
+  b.setInsertionPointToEnd(bodyBlock);
+  Value c2 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(2));
+  Value c3 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(3));
+  arith::MulIOp::create(b, loc, c2, c3);
+  cf::BranchOp::create(b, loc, header);
+
+  EXPECT_EQ(region.getBlocks().size(), 4u);
+
+  b.setInsertionPointAfter(parallelOp.get());
+  IRMapping mapping;
+  scf::ExecuteRegionOp exeRegionOp =
+      wrapMultiBlockRegionWithSCFExecuteRegion(region, mapping, loc, b);
+
+  ASSERT_TRUE(exeRegionOp);
+
+  EXPECT_EQ(exeRegionOp.getRegion().getBlocks().size(), 4u);
+
+  // First block (entry): branch to header
+  Block &exeEntry = exeRegionOp.getRegion().front();
+  EXPECT_TRUE(isa<cf::BranchOp>(exeEntry.getTerminator()));
+
+  // Second block (header): cond_br, 2 predecessors (entry and body)
+  Block &exeHeader = *std::next(exeRegionOp.getRegion().begin());
+  EXPECT_TRUE(isa<cf::CondBranchOp>(exeHeader.getTerminator()));
+
+  // Third block (exit): scf.yield
+  Block &exeExit = *std::next(exeRegionOp.getRegion().begin(), 2);
+  EXPECT_TRUE(isa<scf::YieldOp>(exeExit.getTerminator()));
+
+  // Fourth block (body): branch back to header
+  Block &exeBody = exeRegionOp.getRegion().back();
+  EXPECT_TRUE(isa<cf::BranchOp>(exeBody.getTerminator()));
+  EXPECT_EQ(exeBody.getTerminator()->getSuccessor(0), &exeHeader);
+
+  // Body ops preserved: one addi (entry), one muli (body)
+  unsigned addCount = 0;
+  unsigned mulCount = 0;
+  exeRegionOp.getRegion().walk([&](arith::AddIOp) { ++addCount; });
+  exeRegionOp.getRegion().walk([&](arith::MulIOp) { ++mulCount; });
+  EXPECT_EQ(addCount, 1u);
+  EXPECT_EQ(mulCount, 1u);
+}
+
 //===----------------------------------------------------------------------===//
 // Error Case Tests
 //===----------------------------------------------------------------------===//

``````````

</details>


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