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

[Kunal Dubey]

https://github.com/xakep8 created https://github.com/llvm/llvm-project/pull/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.

Added test for the same.

Fixes #216733

>From 5304c1b740edab5f009a9e052a42a8f609c103e7 Mon Sep 17 00:00:00 2001
From: Kunal Dubey <xakep8 at protonmail.com>
Date: Tue, 18 Aug 2026 12:48:01 +0530
Subject: [PATCH] [mlir][scf] Sink ops into multi-block scf.execute_region

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.

Added test for the same.
---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td    |  3 ++-
 mlir/lib/Dialect/SCF/IR/SCF.cpp               |  5 +++++
 .../Transforms/Utils/ControlFlowSinkUtils.cpp |  6 +++--
 mlir/test/Dialect/SCF/control-flow-sink.mlir  | 22 +++++++++++++++++++
 4 files changed, 33 insertions(+), 3 deletions(-)

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
+}