[Mlir-commits] [mlir] [mlir] Fix RemoveDeadRegionBranchOpSuccessorInputs producing invalid scf.for (LoopLikeOpInterface tie) (PR #216627)
Maksim Levental
llvmlistbot at llvm.org
Sun Aug 16 20:01:16 PDT 2026
https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/216627
## 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`, 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 op result are no longer tied through a shared operand. The pattern 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.
### Fix
Add an optional `RegionBranchStructuralTieFn` callback to `populateRegionBranchOpInterfaceCanonicalizationPatterns` (mirroring the existing `replBuilderFn` hook on the inlining pattern). It lets a dialect report pairs of successor inputs that are structurally coupled and must be added/removed together. `RemoveDeadRegionBranchOpSuccessorInputs` unions those pairs into its tied-value sets.
SCF supplies a callback derived from `LoopLikeOpInterface` (`getRegionIterArgs()` / `getLoopResults()`), tying each loop `iter_arg` to its result. With the tie restored, a live result blocks removal of its dead `iter_arg`, so the pattern no longer produces an invalid intermediate. `getSuccessorRegions`' trip-count precision is untouched, so `InlineRegionBranchOp` still folds single-trip loops (optimization preserved). Because the coupling is read from `LoopLikeOpInterface`, `scf.while` is handled correctly too, not just single-region `scf.for`/`scf.forall`.
The callback lives in the loop dialect rather than the interface library because `MLIRLoopLikeInterface` depends on `MLIRControlFlowInterfaces`; deriving the tie in SCF avoids that layering cycle.
### Validation (Debug + `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON`)
- Previously-crashing tests pass: `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`; single-trip loops still fold.
- No regressions in a 602-test sweep over SCF/ControlFlow/Affine/Linalg/OpenACC/EmitC/X86Vector/SparseTensor/Transforms/SCFTo*.
### Notes (draft)
- Existing tests exercise the fix under the expensive-checks build; a dedicated regression test can be added.
- Alternative to the interface-library approach in #216620 (which reconstructed the tie from region terminator operands and was restricted to single-region loops). This version derives the tie from `LoopLikeOpInterface` directly and is cleaner.
Part of #163599.
>From 340ac93ab8b9ba1e56594153a931aa8d8e945a50 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Sun, 16 Aug 2026 20:00:38 -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 (InlineRegionBranchOp folds
the single-trip loop), 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.
Add an optional structural-tie callback to
populateRegionBranchOpInterfaceCanonicalizationPatterns (mirroring the existing
replBuilderFn hook on the inlining pattern) so a dialect can report pairs of
successor inputs that are structurally coupled and must be added/removed
together. SCF supplies one derived from LoopLikeOpInterface
(getRegionIterArgs()/getLoopResults()), tying each loop iter_arg to its result.
With the tie restored, a live result blocks removal of its dead iter_arg, so the
pattern no longer emits an invalid intermediate; the trip-count precision in
getSuccessorRegions is left intact so InlineRegionBranchOp still folds
single-trip loops.
The callback lives in the loop dialect rather than in the interface library
because MLIRLoopLikeInterface depends on MLIRControlFlowInterfaces; deriving the
tie in SCF avoids that layering cycle.
---
.../mlir/Interfaces/ControlFlowInterfaces.h | 19 ++++++++--
mlir/lib/Dialect/SCF/IR/SCF.cpp | 26 ++++++++++++--
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 36 +++++++++++++++----
3 files changed, 70 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
index 48b80e2059e11..8657e50eb5a9c 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
@@ -330,11 +330,26 @@ Region *getEnclosingRepetitiveRegion(Operation *op);
/// exists.
Region *getEnclosingRepetitiveRegion(Value value);
+/// Callback that reports pairs of successor-input values (e.g. a region
+/// iter_arg and its corresponding op result) that are structurally coupled and
+/// must be added or removed together, even if they are not linked through a
+/// shared successor operand in `getSuccessorOperandInputMapping`. This is needed
+/// when `getSuccessorRegions` refines the control-flow graph and drops an edge
+/// that would otherwise tie them (e.g. a statically-single-trip `scf.for` drops
+/// its region->region back edge, so its iter_args and results are no longer tied
+/// through the yield operand). Each reported pair is unioned into the set of
+/// tied successor inputs. Loop-like ops can implement this using
+/// `getRegionIterArgs()` / `getLoopResults()`.
+using RegionBranchStructuralTieFn =
+ std::function<void(Operation *, SmallVectorImpl<std::pair<Value, Value>> &)>;
+
/// Populate canonicalization patterns that simplify successor operands/inputs
/// of region branch operations. Only operations with the given name are
-/// matched.
+/// matched. `structuralTieFn` is optional; see `RegionBranchStructuralTieFn`.
void populateRegionBranchOpInterfaceCanonicalizationPatterns(
- RewritePatternSet &patterns, StringRef opName, PatternBenefit benefit = 1);
+ RewritePatternSet &patterns, StringRef opName,
+ RegionBranchStructuralTieFn structuralTieFn = nullptr,
+ PatternBenefit benefit = 1);
/// Helper function for the region branch op inlining pattern that builds
/// replacement values for non-successor-input values.
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 54d28783ddc4a..3b93f784e3694 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -1004,11 +1004,33 @@ struct ForOpTensorCastFolder : public OpRewritePattern<ForOp> {
};
} // namespace
+/// Structural-tie callback for `RegionBranchOpInterface` canonicalization: a
+/// loop-like op's region iter_args and their corresponding op results are
+/// structurally coupled (same loop-carried slot, enforced by the op verifier)
+/// and must be added/removed together. The successor-operand/input mapping may
+/// not link them when `getSuccessorRegions` drops an edge for control-flow
+/// precision (e.g. a statically-single-trip `scf.for` drops its back edge), so
+/// report the coupling explicitly via `LoopLikeOpInterface`.
+static void addLoopLikeStructuralTies(
+ Operation *op, SmallVectorImpl<std::pair<Value, Value>> &ties) {
+ auto loopOp = dyn_cast<LoopLikeOpInterface>(op);
+ if (!loopOp)
+ return;
+ std::optional<ResultRange> results = loopOp.getLoopResults();
+ if (!results)
+ return;
+ Block::BlockArgListType iterArgs = loopOp.getRegionIterArgs();
+ if (iterArgs.size() != results->size())
+ return;
+ for (auto [iterArg, result] : llvm::zip_equal(iterArgs, *results))
+ ties.emplace_back(iterArg, result);
+}
+
void ForOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<ForOpTensorCastFolder>(context);
populateRegionBranchOpInterfaceCanonicalizationPatterns(
- results, ForOp::getOperationName());
+ results, ForOp::getOperationName(), addLoopLikeStructuralTies);
populateRegionBranchOpInterfaceInliningPattern(
results, ForOp::getOperationName(),
/*replBuilderFn=*/[](OpBuilder &builder, Location loc, Value value) {
@@ -3733,7 +3755,7 @@ void WhileOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<WhileConditionTruth, WhileCmpCond, WhileOpAlignBeforeArgs,
WhileMoveIfDown>(context);
populateRegionBranchOpInterfaceCanonicalizationPatterns(
- results, WhileOp::getOperationName());
+ results, WhileOp::getOperationName(), addLoopLikeStructuralTies);
populateRegionBranchOpInterfaceInliningPattern(results,
WhileOp::getOperationName());
}
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 475969a7c0d09..4006631f57a9b 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -839,9 +839,12 @@ 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.
struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
- RemoveDeadRegionBranchOpSuccessorInputs(MLIRContext *context, StringRef name,
- PatternBenefit benefit = 1)
- : RewritePattern(name, benefit, context) {}
+ RemoveDeadRegionBranchOpSuccessorInputs(
+ MLIRContext *context, StringRef name,
+ RegionBranchStructuralTieFn structuralTieFn = nullptr,
+ PatternBenefit benefit = 1)
+ : RewritePattern(name, benefit, context),
+ structuralTieFn(std::move(structuralTieFn)) {}
LogicalResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override {
@@ -857,6 +860,21 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
llvm::EquivalenceClasses<Value> tiedSuccessorInputs =
computeTiedSuccessorInputs(operandToInputs);
+ // The mapping above is derived from `getSuccessorRegions`, which may refine
+ // the control-flow graph and drop structurally-required edges (e.g. the
+ // region->region back edge of a statically-single-trip loop). Union any
+ // structural ties reported by the caller so that, e.g., a loop iter_arg is
+ // never removed without its corresponding op result.
+ if (structuralTieFn) {
+ SmallVector<std::pair<Value, Value>> ties;
+ structuralTieFn(op, ties);
+ for (auto &[a, b] : ties) {
+ tiedSuccessorInputs.insert(a);
+ tiedSuccessorInputs.insert(b);
+ tiedSuccessorInputs.unionSets(a, b);
+ }
+ }
+
// Determine which values to remove and group them by block and operation.
SmallVector<Value> valuesToRemove;
DenseMap<Block *, BitVector> blockArgsToRemove;
@@ -940,6 +958,8 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
return success();
}
+
+ RegionBranchStructuralTieFn structuralTieFn;
};
/// Return the "owner" of a value: the parent block for block arguments, the
@@ -1300,11 +1320,13 @@ struct InlineRegionBranchOp : public RewritePattern {
} // namespace
void mlir::populateRegionBranchOpInterfaceCanonicalizationPatterns(
- RewritePatternSet &patterns, StringRef opName, PatternBenefit benefit) {
+ RewritePatternSet &patterns, StringRef opName,
+ RegionBranchStructuralTieFn structuralTieFn, PatternBenefit benefit) {
patterns.add<MakeRegionBranchOpSuccessorInputsDead,
- RemoveDuplicateSuccessorInputUses,
- RemoveDeadRegionBranchOpSuccessorInputs>(patterns.getContext(),
- opName, benefit);
+ RemoveDuplicateSuccessorInputUses>(patterns.getContext(), opName,
+ benefit);
+ patterns.add<RemoveDeadRegionBranchOpSuccessorInputs>(
+ patterns.getContext(), opName, std::move(structuralTieFn), benefit);
}
void mlir::populateRegionBranchOpInterfaceInliningPattern(
More information about the Mlir-commits
mailing list