[llvm-branch-commits] [mlir] [mlir][SCF] Remove `RegionBranchOpInterface` from `scf.forall` (PR #174221)

[Matthias Springer via llvm-branch-commits]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/174221

>From 9cad651911cba021a434b0fa8b712b37ee1e592f Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Thu, 8 Jan 2026 12:58:02 +0000
Subject: [PATCH] Implement RegionBranchTerminatorOpInterface for
 scf.forall.in_parallel

---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td |  8 +-----
 mlir/lib/Dialect/SCF/IR/SCF.cpp            | 30 +++++++++++++---------
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 8bdf3e0b566ef..3efa403287d51 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -649,13 +649,6 @@ def ForallOp : SCF_Op<"forall", [
 
     /// Returns true if the mapping specified for this forall op is linear.
     bool usesLinearMapping();
-
-    /// RegionBranchOpInterface
-
-    OperandRange getEntrySuccessorOperands(RegionSuccessor successor) {
-      return getInits();
-    }
-
   }];
 }
 
@@ -667,6 +660,7 @@ def InParallelOp : SCF_Op<"forall.in_parallel", [
        Pure,
        Terminator,
        DeclareOpInterfaceMethods<InParallelOpInterface>,
+       ReturnLike,
        HasParent<"ForallOp">,
       ] # GraphRegionNoTerminator.traits> {
   let summary = "terminates a `forall` block";
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index a25b369d540ac..a4497f7f8464a 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2013,20 +2013,26 @@ void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
               ForallOpReplaceConstantInductionVar>(context);
 }
 
-/// Given the region at `index`, or the parent operation if `index` is None,
-/// return the successor regions. These are the regions that may be selected
-/// during the flow of control. `operands` is a set of optional attributes that
-/// correspond to a constant value for each operand, or null if that operand is
-/// not a constant.
 void ForallOp::getSuccessorRegions(RegionBranchPoint point,
                                    SmallVectorImpl<RegionSuccessor> &regions) {
-  // In accordance with the semantics of forall, its body is executed in
-  // parallel by multiple threads. We should not expect to branch back into
-  // the forall body after the region's execution is complete.
-  if (point.isParent())
-    regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
-  else
-    regions.push_back(RegionSuccessor::parent(getResults()));
+  // There are two region branch points:
+  // 1. "parent": entering the forall op for the first time.
+  // 2. scf.in_parallel terminator
+  if (point.isParent()) {
+    // When first entering the forallc op, the control flow typically branches
+    // into the forall body. (In parallel for multiple threads.)
+    regions.push_back(RegionSuccessor(&getRegion()));
+    // However, when there are 0 threads, the control flow may branch back to
+    // the parent immediately.
+    regions.push_back(RegionSuccessor::parent(
+        ResultRange{getResults().end(), getResults().end()}));
+  } else {
+    // In accordance with the semantics of forall, its body is executed in
+    // parallel by multiple threads. We should not expect to branch back into
+    // the forall body after the region's execution is complete.
+    regions.push_back(RegionSuccessor::parent(
+        ResultRange{getResults().end(), getResults().end()}));
+  }
 }
 
 //===----------------------------------------------------------------------===//