[Mlir-commits] [mlir] 3c02e3a - [mlir][scf] Sink ops into multi-block scf.execute_region (#216993)

[llvmlistbot at llvm.org]

Author: Kunal Dubey
Date: 2026-08-18T13:55:24+02:00
New Revision: 3c02e3ab3c29df7ff3a6751191c5c5240dbfb51b

URL: https://github.com/llvm/llvm-project/commit/3c02e3ab3c29df7ff3a6751191c5c5240dbfb51b
DIFF: https://github.com/llvm/llvm-project/commit/3c02e3ab3c29df7ff3a6751191c5c5240dbfb51b.diff

LOG: [mlir][scf] Sink ops into multi-block scf.execute_region (#216993)

Teach scf.execute_region to report its body is invoked exactly once via
RegionBranchOpInterface. This make the execute_region eligible for
sinking.

Also walk all blocks in a region instead of just the entry block so uses
in non-entry blocks can trigger sinking.

Fixes #216733

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
    mlir/lib/Dialect/SCF/IR/SCF.cpp
    mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
    mlir/test/Dialect/SCF/control-flow-sink.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 1e7bba2cdad35..2fe8046a92463 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -79,7 +79,8 @@ def ConditionOp : SCF_Op<"condition", [
 //===----------------------------------------------------------------------===//
 
 def ExecuteRegionOp : SCF_Op<"execute_region", [
-    DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getSuccessorInputs"]>,
+    DeclareOpInterfaceMethods<RegionBranchOpInterface, [
+      "getRegionInvocationBounds", "getSuccessorInputs"]>,
     DeclareOpInterfaceMethods<PromotableRegionOpInterface>,
     RecursiveMemoryEffects]> {
   let summary = "operation that executes its region exactly once";

diff  --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 36990db4b7dc4..3388fa490f996 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -277,6 +277,11 @@ void ExecuteRegionOp::getSuccessorRegions(
   regions.push_back(RegionSuccessor(getOperation()));
 }
 
+void ExecuteRegionOp::getRegionInvocationBounds(
+    ArrayRef<Attribute>, SmallVectorImpl<InvocationBounds> &bounds) {
+  bounds.emplace_back(/*lb=*/1, /*ub=*/1);
+}
+
 ValueRange ExecuteRegionOp::getSuccessorInputs(RegionSuccessor successor) {
   return successor.isOperation() ? ValueRange(getOperation()->getResults())
                                  : ValueRange();

diff  --git a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
index 19cf4646bddd4..b546d3adef08f 100644
--- a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
@@ -110,8 +110,10 @@ void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
 void Sinker::sinkRegion(Region *region) {
   // Initialize the work queue with all the ops in the region.
   std::vector<Operation *> stack;
-  for (Operation &op : region->getOps())
-    stack.push_back(&op);
+  for (Block &block : *region)
+    // Seed from all block so that uses in non-entry blocks and trigger sinking
+    for (Operation &op : block)
+      stack.push_back(&op);
 
   // Process all the ops depth-first. This ensures that nodes of subgraphs are
   // sunk in the correct order.

diff  --git a/mlir/test/Dialect/SCF/control-flow-sink.mlir b/mlir/test/Dialect/SCF/control-flow-sink.mlir
index 2eb888b378ba1..6a3e2b651ee8a 100644
--- a/mlir/test/Dialect/SCF/control-flow-sink.mlir
+++ b/mlir/test/Dialect/SCF/control-flow-sink.mlir
@@ -58,3 +58,25 @@ func.func @test_scf_if_double_sink(%arg0: i1, %arg1: i32) {
   }
   return
 }
+
+// -----
+
+func.func private @consume(i32) -> ()
+
+// CHECK-LABEL: @test_scf_execute_region_multiblock_sink
+// CHECK-SAME:  (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
+// CHECK: scf.execute_region
+// CHECK-NEXT:   %[[V0:.*]] = arith.muli %[[ARG0]], %[[ARG1]]
+// CHECK-NEXT:   cf.br ^bb1
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT:   call @consume(%[[V0]])
+func.func @test_scf_execute_region_multiblock_sink(%arg0: i32, %arg1: i32) {
+  %0 = arith.muli %arg0, %arg1 : i32
+  scf.execute_region {
+    cf.br ^bb1
+  ^bb1:
+    func.call @consume(%0) : (i32) -> ()
+    scf.yield
+  }
+  return
+}