[Mlir-commits] [mlir] [mlir][SCF] Fix region branch op interfaces for `scf.forall` and its terminator (PR #174221)

[Matthias Springer]

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

>From fec4bc1f6fc8b1e1eac3b7519ea83412fa8a4b8f 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 1/2] Implement RegionBranchTerminatorOpInterface for
 scf.forall.in_parallel

---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td |  8 +-----
 mlir/lib/Dialect/SCF/IR/SCF.cpp            | 31 +++++++++++++---------
 2 files changed, 19 insertions(+), 20 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 8803a6d136f7a..78dc27d6ec746 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2013,21 +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(getOperation(), getOperation()->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.emplace_back(getOperation(),
+                         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.emplace_back(getOperation(),
+                         ResultRange{getResults().end(), getResults().end()});
+  }
 }
 
 //===----------------------------------------------------------------------===//

>From 0ab6ffd60ed551bff3278d75fc0e45593a5a3d88 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Mon, 12 Jan 2026 17:52:42 +0000
Subject: [PATCH 2/2] add test case

---
 mlir/lib/Dialect/SCF/IR/SCF.cpp                    |  2 +-
 .../Analysis/DataFlow/test-dead-code-analysis.mlir | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 78dc27d6ec746..91799f48351c4 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2019,7 +2019,7 @@ void ForallOp::getSuccessorRegions(RegionBranchPoint point,
   // 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
+    // When first entering the forall 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
diff --git a/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir b/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir
index 1ae6855351506..98f1602b1effc 100644
--- a/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir
+++ b/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir
@@ -269,3 +269,17 @@ func.func @test_dca_doesnt_crash() -> () {
 func.func @test_dca_doesnt_crash_2() -> () attributes {symbol = @notexistant} {
    return
 }
+
+func.func @test_forall_op_control_flow() {
+  // CHECK: test_forall_op_control_flow:
+  // CHECK:  region #0
+  // CHECK:   ^bb0 = live
+  // CHECK: region_preds: (all) predecessors:
+  // CHECK:   scf.forall (%{{.*}}) in (4) {...} {tag = "test_forall_op_control_flow"}
+  // CHECK: op_preds: (all) predecessors:
+  // CHECK:   scf.forall (%{{.*}}) in (4) {...} {tag = "test_forall_op_control_flow"}
+  // CHECK:   scf.forall.in_parallel {...}
+  scf.forall (%arg0) in (4) {
+  } {tag = "test_forall_op_control_flow"}
+  return
+}