[MLIR] Fix control-flow sinking through nested regions

Project: llvm/llvm-project | Language: C++ | Concepts: algorithmic-optimization, data-structures | Difficulty: mid
Link: https://github.com/llvm/llvm-project/pull/217168

The problem

The MLIR control-flow sinking utility did not correctly handle operations used inside nested regions (e.g., an scf.for inside an scf.if). Specifically, it would not sink operations into an enclosing region if their only uses were in nested subregions, missing optimization opportunities and causing incorrect code motion in certain cases.

Why it's non-trivial

A naive fix might only check if an operation's uses are within the immediate region, missing that nested regions (regions within regions) should also be considered. Properly handling this requires understanding and traversing the region hierarchy, ensuring that operations are only sunk when it is safe and correct to do so, even in the presence of complex nesting.

The approach

1. Correctly identify users in nested regions as eligible for sinking

Previously, the code only allowed sinking if all users were dominated by the region's entry block. The fix adds a check to also allow users in descendant (nested) regions:

-    // The user is dominated by the region if its containing block is dominated
-    // by the region's entry block.
-    return domInfo.dominates(&region->front(), user->getBlock());
+    // by the region's entry block. Additionally, allow users that are in
+    // descendant regions of the region (e.g., nested loops) since those
+    // should still permit sinking into the outer region.
+    return domInfo.dominates(&region->front(), user->getBlock()) ||
+           region->isAncestor(user->getParentRegion());

This change ensures that operations used in nested regions are considered for sinking, not just those in the immediate region.

2. Properly filter which operations to consider for sinking

The code originally ignored operations already inside the target region, but did not handle nested regions correctly. The fix updates the check:

-    // Ignore block arguments and ops that are already inside the region.
-    if (!op || op->getParentRegion() == region)
+    // Ignore block arguments and ops already contained in the target region,
+    // including ops in nested regions. Only consider defs outside the target
+    // region.
+    if (!op || region->isAncestor(op->getParentRegion()))
       continue;

Now, any operation already inside the target region or its descendants is skipped, preventing redundant or unsafe moves.

3. Traverse all operations in the region, including nested regions

The original work queue only seeded operations from the immediate blocks of the region. The fix uses a recursive walk to include all nested operations:

-  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);
+  region->walk([&](Operation *op) { stack.push_back(op); });

This ensures that all relevant operations, even those in nested regions, are considered for sinking.

4. Regression tests for nested region sinking

New tests are added to verify that the sinking now works for cases where an operation is used inside a nested scf.for within an scf.if:

+func.func @test_scf_if_sink_through_loop(
+    %arg0: i1, %arg1: index, %arg2: i32, %arg3: i32) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = arith.muli %arg2, %arg3 : i32
+  scf.if %arg0 {
+    scf.for %arg4 = %c0 to %arg1 step %c1 {
+      func.call @sink_i32(%0) : (i32) -> ()
+    }
+  }
+  return
+}

This test checks that %0 = arith.muli ... is sunk into the scf.if region, even though its use is in a nested scf.for.

From the review thread

No notable review discussion on this PR.

Takeaways

• When optimizing code with nested data structures (like regions or blocks), always consider the full hierarchy; direct containment checks may miss valid cases.
• When moving or transforming code, ensure that all users—including those in nested or descendant structures—are accounted for to avoid missed optimizations or incorrect transformations.
• Automated tests for regressions in complex control-flow or nesting scenarios are essential to prevent subtle bugs from reappearing.