[Mlir-commits] [mlir] [mlir][affine] Fix SSA result handling and dominance in sibling loop fusion (PR #232841)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 16 09:53:40 PDT 2026
https://github.com/dashygo097 created https://github.com/llvm/llvm-project/pull/232841
### Summary
This PR fixes the handling of SSA results when performing sibling fusion on result-producing affine.for loops, particularly reduction loops.
Previously, sibling fusion could fail to preserve the results of the original source loop, access an invalidated destination loop operation, or produce SSA dominance violations when fusing dependent loops.
This change improves reduction promotion, SSA result remapping, and sibling fusion legality checks to ensure that the resulting IR remains valid.
### Changes
This patch improves SSA dependency checking in `findSiblingNodeToFuse()` to prevent invalid sibling fusion. The `canFuseWithSibNode()`lambda uses `getForwardSlice()` to detect direct and transitive SSA dependencies between candidate loops and skips fusion when a loop result is used by the other loop. It also checks for intervening SSA result users to avoid fusion that would violate dominance.
In `LoopFusionUtils.cpp`, `fuseLoops()` clones the source loop into the destination, and `promoteSingleIterReductionLoop()` uses `replaceWithAdditionalYields()` to extend the destination loop with additional reduction results. Since this replaces the original destination operation, `fuseWithSiblingNodes()` recovers the new destination through the preserved body terminator and updates the dependence graph accordingly. The original source loop's results are also remapped to the corresponding results of the fused loop before the source is erased, ensuring that existing SSA users remain valid.
### Example
Two independent reductions can be fused while preserving both results.
Before:
```mlir
%0 = affine.for %i = 0 to 32
iter_args(%acc = %zero) -> f32 {
// First reduction.
}
%1 = affine.for %i = 0 to 32
iter_args(%acc = %zero) -> f32 {
// Second reduction.
}
%result = arith.addf %0, %1 : f32
```
After:
```mlir
%0:2 = affine.for %i = 0 to 32
iter_args(%acc0 = %zero, %acc1 = %zero) -> (f32, f32) {
// Both reductions.
affine.yield %next0, %next1 : f32, f32
}
%result = arith.addf %0#0, %0#1 : f32
```
Conversely, sibling fusion is rejected when the second loop depends on the first loop's result or when the resulting loop placement would violate SSA dominance.
### Testing
Regression tests exercise both legal and illegal sibling fusion scenarios using:
```bash
mlir-opt loop-fusion-sibling.mlir \
-pass-pipeline='builtin.module(func.func(affine-loop-fusion{maximal mode=sibling}))' \
--verify-each
```
which are located in `mlir/test/Dialect/Affine/loop-fusion-sibling.mlir`.
The tests cover successful fusion of independent reductions as well as rejection of transformations that would introduce invalid SSA references.
>From 37748a0029250dbb6f87ac66deee144b6c470d8a Mon Sep 17 00:00:00 2001
From: dashygo097 <dashygo097 at gmail.com>
Date: Fri, 16 Oct 2026 15:38:43 +0800
Subject: [PATCH] [mlir][affine] Sibling loop fusion does not correctly handle
SSA results from reduction loops
---
.../Dialect/Affine/Transforms/LoopFusion.cpp | 47 +++++--
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 17 ++-
.../Dialect/Affine/loop-fusion-sibling.mlir | 133 ++++++++++++++++++
3 files changed, 181 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 77eda4062ec0c1..c21e8c17131be4 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -12,6 +12,7 @@
#include "mlir/Dialect/Affine/Transforms/Passes.h"
+#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/Utils.h"
@@ -24,6 +25,7 @@
#include "mlir/IR/Builders.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
@@ -63,7 +65,6 @@ struct LoopFusion : public affine::impl::AffineLoopFusionBase<LoopFusion> {
void runOnBlock(Block *block);
void runOnOperation() override;
};
-
} // namespace
/// Returns true if node 'srcId' can be removed after fusing it with node
@@ -725,7 +726,6 @@ static bool isFusionProfitable(AffineForOp srcForOp,
}
namespace {
-
// GreedyFusion greedily fuses loop nests which have a producer/consumer or
// input-reuse relationship on a memref, with the goal of improving locality.
//
@@ -1327,6 +1327,9 @@ struct GreedyFusion {
continue;
}
+ auto *oldDstOp = dstNode->op;
+ auto *dstYield = dstAffineForOp.getBody()->getTerminator();
+
// Check if source loop is being inserted in the innermost
// destination loop. Based on this, the fused loop may be optimized
// further inside `fuseLoops`.
@@ -1335,9 +1338,12 @@ struct GreedyFusion {
affine::fuseLoops(sibAffineForOp, dstAffineForOp, bestSlice,
isInnermostInsertion);
+ dstAffineForOp = cast<AffineForOp>(dstYield->getParentOp());
+ dstNode->op = dstAffineForOp.getOperation();
+
auto dstForInst = cast<AffineForOp>(dstNode->op);
// Update operation position of fused loop nest (if needed).
- if (insertPointInst != dstForInst)
+ if (insertPointInst != oldDstOp && insertPointInst != dstForInst)
dstForInst->moveBefore(insertPointInst);
LDBG() << "Fused sibling nest " << sibId << " into destination nest "
@@ -1369,6 +1375,22 @@ struct GreedyFusion {
// TODO: Remove restrict to single load op restriction.
if (sibNode->getLoadOpCount(memref) != 1)
return false;
+
+ // Skip sibling loops with SSA dependences in either direction.
+ for (Operation *srcOp : {sibNode->op, dstNode->op}) {
+ Operation *dstOp = srcOp == sibNode->op ? dstNode->op : sibNode->op;
+
+ for (Value result : srcOp->getResults()) {
+ SetVector<Operation *> forwardSlice;
+ getForwardSlice(result, &forwardSlice);
+
+ for (Operation *op : forwardSlice) {
+ if (op == dstOp || dstOp->isAncestor(op))
+ return false;
+ }
+ }
+ }
+
// Skip if there exists a path of dependent edges between
// 'sibNode' and 'dstNode'.
if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
@@ -1431,8 +1453,8 @@ struct GreedyFusion {
}
}
- // Search for siblings by following edges through an intermediate src node.
- // Collect candidate 'dstNode' input edges in 'inEdges'.
+ // Search for siblings by following edges through an intermediate src
+ // node. Collect candidate 'dstNode' input edges in 'inEdges'.
SmallVector<MemRefDependenceGraph::Edge, 2> inEdges;
mdg->forEachMemRefInputEdge(
dstNode->id, [&](MemRefDependenceGraph::Edge inEdge) {
@@ -1444,10 +1466,11 @@ struct GreedyFusion {
inEdges.push_back(inEdge);
});
- // Search for sibling nodes to fuse by visiting output edges from each input
- // edge in 'inEdges'.
+ // Search for sibling nodes to fuse by visiting output edges from each
+ // input edge in 'inEdges'.
for (auto &inEdge : inEdges) {
- // Collect candidate output edges from each node 'inEdge.id' in 'inEdges'.
+ // Collect candidate output edges from each node 'inEdge.id' in
+ // 'inEdges'.
SmallVector<MemRefDependenceGraph::Edge, 2> outEdges;
mdg->forEachMemRefOutputEdge(
inEdge.id, [&](MemRefDependenceGraph::Edge outEdge) {
@@ -1460,7 +1483,8 @@ struct GreedyFusion {
auto *sibNode = mdg->getNode(sibNodeId);
if (!isa<AffineForOp>(sibNode->op))
return;
- // Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
+ // Check if 'sibNode/dstNode' can be input-reuse fused on
+ // 'memref'.
if (canFuseWithSibNode(sibNode, outEdge.value)) {
// Add candidate 'outEdge' to sibling node.
outEdges.push_back(outEdge);
@@ -1478,8 +1502,8 @@ struct GreedyFusion {
return false;
}
- /// Update data dependence graph state to reflect sibling fusion of 'sibNode'
- /// into 'dstNode'.
+ /// Update data dependence graph state to reflect sibling fusion of
+ /// 'sibNode' into 'dstNode'.
void updateStateAfterSiblingFusion(Node *sibNode, Node *dstNode) {
// Update 'sibNode' and 'dstNode' input/output edges to reflect fusion.
mdg->updateEdges(sibNode->id, dstNode->id);
@@ -1511,7 +1535,6 @@ struct GreedyFusion {
}
}
};
-
} // namespace
/// Run fusion on `block`.
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 68296ea3368a1f..48921ceffe481a 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -428,7 +428,7 @@ void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
// Clone 'srcForOp' into 'dstForOp' at 'srcSlice->insertPoint'.
OpBuilder b(srcSlice.insertPoint->getBlock(), srcSlice.insertPoint);
IRMapping mapper;
- b.clone(*srcForOp, mapper);
+ auto *clonedSrc = b.clone(*srcForOp, mapper);
// Update 'sliceLoopNest' upper and lower bounds from computed 'srcSlice'.
SmallVector<AffineForOp, 4> sliceLoops;
@@ -458,11 +458,20 @@ void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
// Fix up and if possible, eliminate single iteration loops.
for (AffineForOp forOp : sliceLoops) {
if (isLoopParallelAndContainsReduction(forOp) &&
- isInnermostSiblingInsertion && srcIsUnitSlice())
+ isInnermostSiblingInsertion && srcIsUnitSlice()) {
// Patch reduction loop - only ones that are sibling-fused with the
// destination loop - into the parent loop.
- (void)promoteSingleIterReductionLoop(forOp, true);
- else
+ auto *dstYield = dstForOp.getBody()->getTerminator();
+ unsigned oldParentResults = forOp->getParentOp()->getNumResults();
+ if (succeeded(promoteSingleIterReductionLoop(forOp, true)) &&
+ forOp.getOperation() == clonedSrc) {
+ AffineForOp newDst = cast<AffineForOp>(dstYield->getParentOp());
+ for (unsigned i = 0, e = srcForOp.getNumResults(); i < e; ++i) {
+ srcForOp.getResult(i).replaceAllUsesWith(
+ newDst.getResult(oldParentResults + i));
+ }
+ }
+ } else
// Promote any single iteration slice loops.
(void)promoteIfSingleIteration(forOp);
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
index 937c855b86b506..0a3f6fb38d7555 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -21,3 +21,136 @@ func.func @disjoint_stores(%0: memref<8xf32>) {
// CHECK-NOT: affine.for
return
}
+
+// Test fusion of two independent reductions with live SSA results.
+// Both original loop results must be preserved after fusion.
+
+// CHECK-LABEL: func.func @independent_reductions
+// CHECK: %[[FUSED:[0-9]+]]:2 = affine.for
+// CHECK: affine.yield {{.*}} : f32, f32
+// CHECK: arith.subf %[[FUSED]]#1, %[[FUSED]]#0
+func.func @independent_reductions(%a: memref<32xf32>) -> f32 {
+ %zero = arith.constant 0.0 : f32
+ %two = arith.constant 2.0 : f32
+
+ %first = affine.for %i = 0 to 32
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %sum = arith.addf %acc, %x : f32
+ affine.yield %sum : f32
+ }
+
+ %second = affine.for %i = 0 to 32
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %scaled = arith.mulf %x, %two : f32
+ %sum = arith.addf %acc, %scaled : f32
+ affine.yield %sum : f32
+ }
+
+ %result = arith.subf %first, %second : f32
+ return %result : f32
+}
+
+// Test that sibling fusion is skipped when the second loop directly
+// depends on the first loop's SSA result.
+
+// CHECK-LABEL: func.func @direct_ssa_dependence
+// CHECK: %[[FIRST:[0-9]+]] = affine.for
+// CHECK: %[[SECOND:[0-9]+]] = affine.for {{.*}}iter_args({{.*}} = %[[FIRST]])
+func.func @direct_ssa_dependence(%a: memref<32xf32>) -> f32 {
+ %zero = arith.constant 0.0 : f32
+
+ %first = affine.for %i = 0 to 32
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %sum = arith.addf %acc, %x : f32
+ affine.yield %sum : f32
+ }
+
+ %second = affine.for %i = 0 to 32
+ iter_args(%acc = %first) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %sum = arith.addf %acc, %x : f32
+ affine.yield %sum : f32
+ }
+
+ return %second : f32
+}
+
+// Test that transitive SSA dependencies also prevent sibling fusion.
+
+// CHECK-LABEL: func.func @transitive_ssa_dependence
+// CHECK: %[[FIRST:[0-9]+]] = affine.for
+// CHECK: %[[SEED:[0-9]+]] = arith.addf %[[FIRST]],
+// CHECK: %[[SECOND:[0-9]+]] = affine.for {{.*}}iter_args({{.*}} = %[[SEED]])
+func.func @transitive_ssa_dependence(%a: memref<32xf32>) -> f32 {
+ %zero = arith.constant 0.0 : f32
+ %one = arith.constant 1.0 : f32
+
+ %first = affine.for %i = 0 to 32
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %sum = arith.addf %acc, %x : f32
+ affine.yield %sum : f32
+ }
+
+ %seed = arith.addf %first, %one : f32
+
+ %second = affine.for %i = 0 to 32
+ iter_args(%acc = %seed) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %sum = arith.addf %acc, %x : f32
+ affine.yield %sum : f32
+ }
+
+ return %second : f32
+}
+
+// Test sibling fusion of three independent reductions that share
+// different input memrefs. All three SSA results must be preserved.
+
+// CHECK-LABEL: func.func @alternating_dot
+// CHECK: %[[FUSED:[0-9]+]]:3 = affine.for
+// CHECK: affine.yield {{.*}} : f32, f32, f32
+// CHECK: arith.addf
+// CHECK: arith.addf
+// CHECK: return
+
+func.func @alternating_dot(
+ %a: memref<32xf32>,
+ %b: memref<32xf32>,
+ %c: memref<32xf32>) -> f32 {
+ %zero = arith.constant 0.0 : f32
+
+ %tmp1 = affine.for %i = 0 to 32
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %y = affine.load %b[%i] : memref<32xf32>
+ %product = arith.mulf %x, %y : f32
+ %sum = arith.addf %acc, %product : f32
+ affine.yield %sum : f32
+ }
+
+ %tmp2 = affine.for %i = 0 to 32
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %b[%i] : memref<32xf32>
+ %y = affine.load %c[%i] : memref<32xf32>
+ %product = arith.mulf %x, %y : f32
+ %sum = arith.addf %acc, %product : f32
+ affine.yield %sum : f32
+ }
+
+ %tmp3 = affine.for %i = 0 to 32
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<32xf32>
+ %y = affine.load %c[%i] : memref<32xf32>
+ %product = arith.mulf %x, %y : f32
+ %sum = arith.addf %acc, %product : f32
+ affine.yield %sum : f32
+ }
+
+ %sum12 = arith.addf %tmp1, %tmp2 : f32
+ %result = arith.addf %sum12, %tmp3 : f32
+ return %result : f32
+}
More information about the Mlir-commits
mailing list