[Mlir-commits] [mlir] [mlir] Fix RemoveDeadRegionBranchOpSuccessorInputs producing invalid scf.for (PR #216620)

[Maksim Levental]

https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/216620

## Summary

`RemoveDeadRegionBranchOpSuccessorInputs` (a `RegionBranchOpInterface` canonicalization) can leave an `scf.for` structurally invalid, aborting under `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS`:

```
'scf.for' op mismatch in number of loop-carried values and defined values
LLVM ERROR: IR failed to verify after pattern application
```

### Root cause

The pattern builds its tied-value sets from `getSuccessorOperandInputMapping`, which is derived from `getSuccessorRegions`. For an `scf.for` with a **statically-known trip count of 1**, `getSuccessorRegions` drops the `region -> region` back edge (the loop provably never iterates back). That back edge is what forwards a `scf.yield` operand to the region `iter_args`, so without it the `iter_arg` and its corresponding op result are no longer tied through a shared operand. The pattern then removes a dead `iter_arg` (and its init/yield operands) **without** the structurally-required result, yielding a loop with `0` iter_args but `1` result/yield.

The greedy driver repairs this on a later iteration (`InlineRegionBranchOp` folds the single-trip loop), so it is only observed with expensive pattern-API checks, which verify after **every** pattern application.

This is the dominant cause of the Linalg tiling / pack-unpack / convolution and `scf` loop-canonicalization failures reported in #163599 under `-DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON`.

### Fix

Reconstruct the structural `iter_arg <-> result` tie directly from the region terminator's successor operands — querying **both** the self/back edge (`RegionSuccessor(&region)`) and the exit edge (`RegionSuccessor(op)`) explicitly, independent of the trip-count-refined `getSuccessorRegions`. A live result then blocks removal of its dead `iter_arg`, so the pattern no longer produces an invalid intermediate.

The self/back-edge query is guarded to single-region ops (`getNumRegions() == 1`, i.e. `scf.for`/`scf.forall`); multi-region ops (`scf.while`/`scf.if`) have no `region -> same-region` edge. Crucially, `getSuccessorRegions`' trip-count precision is left untouched, so `InlineRegionBranchOp` still folds single-trip loops (the optimization is preserved).

### Status / notes (draft)

- Verified locally with a Debug + `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON` build: the previously-crashing `Dialect/Linalg/{decompose-pack-tile,decompose-unpack-tile,continuous-tiling-full,transform-op-tile-pack-unpack,transform-tile-and-winograd-rewrite,transform-op-peel-and-vectorize-conv}.mlir` and `Dialect/SCF/{for-loop-peeling,canonicalize}.mlir` now pass, single-trip loops still fold, and a 602-test sweep shows no regressions.
- These tests already exercise the fix when built with the flag; a dedicated regression test can be added.
- Does not yet address the analogous coupling for multi-region loops (`scf.while`). A cleaner production form might derive the tie from `LoopLikeOpInterface` (`getRegionIterArgs()`/`getLoopResults()`), currently blocked by a library layering cycle (`MLIRLoopLikeInterface` depends on `MLIRControlFlowInterfaces`).

Part of #163599.

>From 4dc3ef8e1bb53c9912e6bbe6424a39ffb5fc2ba4 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Sun, 16 Aug 2026 19:35:43 -0700
Subject: [PATCH] [mlir] Fix RemoveDeadRegionBranchOpSuccessorInputs producing
 invalid scf.for

RemoveDeadRegionBranchOpSuccessorInputs builds its tied-value sets from
RegionBranchOpInterface::getSuccessorOperandInputMapping, which is derived from
getSuccessorRegions. For an scf.for with a statically-known trip count of 1,
getSuccessorRegions drops the region->region back edge (the loop provably never
iterates back). That back edge is what forwards a yield operand to the region
iter_args, so without it an iter_arg and its corresponding op result are no
longer tied through a shared operand. The pattern then removes a dead iter_arg
without its (structurally required) result, producing an scf.for with
mismatched loop-carried counts:

    'scf.for' op mismatch in number of loop-carried values and defined values

The greedy driver repairs this on a later iteration, so it is only observed
with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS (which verifies the IR after
every pattern application). It shows up across Linalg tiling / pack-unpack /
convolution lowering and scf loop canonicalization; see issue #163599.

Reconstruct the structural iter_arg<->result tie directly from the region
terminator's successor operands (querying both the self/back edge and the exit
edge explicitly), independent of the trip-count-refined getSuccessorRegions, so
a live result blocks removal of its dead iter_arg. The trip-count precision in
getSuccessorRegions is left intact so InlineRegionBranchOp still folds
single-trip loops.
---
 mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 65 +++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 475969a7c0d09..bd17a8382d581 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -838,6 +838,65 @@ static llvm::EquivalenceClasses<Value> computeTiedSuccessorInputs(
 /// }
 /// There are two sets: {{%r1}, {%r2}}. Each set has one value, so there each
 /// value can be removed independently of the other values.
+/// Augment `tied` with *structural* ties between successor inputs that are
+/// forwarded from the same region-terminator operand.
+///
+/// `getSuccessorOperandInputMapping` is built from the op's
+/// `getSuccessorRegions`, which may be refined by control-flow precision. For
+/// example, `scf.for` with a statically-known trip count of 1 drops its
+/// region->region back edge (the loop provably never iterates back). That back
+/// edge is what forwards a region terminator operand to the region iter_args;
+/// without it, the iter_arg and its corresponding op result are no longer tied
+/// through a shared operand, and `RemoveDeadRegionBranchOpSuccessorInputs` may
+/// remove a dead iter_arg without its (structurally required) result, producing
+/// a structurally invalid op.
+///
+/// The structural tie exists regardless of trip count: for a single-region
+/// self-looping op (e.g. `scf.for`, `scf.forall`), the region terminator
+/// forwards each operand both to a region iter_arg (self/back edge) and to the
+/// corresponding op result (exit edge), so those two inputs must be added or
+/// removed together. Reconstruct that coupling by querying the terminator's
+/// operands for both edges explicitly, independent of the (possibly refined)
+/// `getSuccessorRegions`.
+static void
+augmentWithStructuralTerminatorTies(RegionBranchOpInterface branchOp,
+                                    llvm::EquivalenceClasses<Value> &tied) {
+  Operation *op = branchOp.getOperation();
+  // Only single-region ops have a region->same-region back edge. For
+  // multi-region ops (e.g. scf.while/scf.if) the terminator does not branch to
+  // its own region, and querying that edge is invalid.
+  if (op->getNumRegions() != 1)
+    return;
+  Region &region = op->getRegion(0);
+  if (region.empty())
+    return;
+  auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(
+      region.back().getTerminator());
+  if (!terminator)
+    return;
+
+  DenseMap<OpOperand *, Value> operandToInput;
+  auto record = [&](OperandRange ops, ValueRange inputs) {
+    if (ops.size() != inputs.size())
+      return;
+    for (auto [operand, input] :
+         llvm::zip_equal(operandsToOpOperands(ops), inputs)) {
+      auto [it, inserted] = operandToInput.try_emplace(&operand, input);
+      tied.insert(input);
+      if (!inserted)
+        tied.unionSets(it->second, input);
+    }
+  };
+  // Exit edge: terminator operands -> op results.
+  RegionSuccessor exitEdge(op);
+  record(terminator.getSuccessorOperands(exitEdge),
+         branchOp.getSuccessorInputs(exitEdge));
+  // Self/back edge: terminator operands -> region iter_args.
+  RegionSuccessor backEdge(&region);
+  record(terminator.getSuccessorOperands(backEdge),
+         branchOp.getSuccessorInputs(backEdge));
+}
+
 struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
   RemoveDeadRegionBranchOpSuccessorInputs(MLIRContext *context, StringRef name,
                                           PatternBenefit benefit = 1)
@@ -856,6 +915,12 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
     regionBranchOp.getSuccessorOperandInputMapping(operandToInputs);
     llvm::EquivalenceClasses<Value> tiedSuccessorInputs =
         computeTiedSuccessorInputs(operandToInputs);
+    // The successor-operand/input mapping above is derived from
+    // `getSuccessorRegions`, which may drop structurally-required edges for
+    // control-flow precision (e.g. the back edge of a statically-single-trip
+    // loop). Add back the structural ties so that a region iter_arg is never
+    // removed without its corresponding op result.
+    augmentWithStructuralTerminatorTies(regionBranchOp, tiedSuccessorInputs);
 
     // Determine which values to remove and group them by block and operation.
     SmallVector<Value> valuesToRemove;