[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
Sun Oct 18 06:03:32 PDT 2026
https://github.com/dashygo097 updated https://github.com/llvm/llvm-project/pull/232841
>From ad49d692a9399a26c10dcb5e119ec5f4688e31d4 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 | 50 ++++--
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 17 +-
.../Dialect/Affine/loop-fusion-sibling.mlir | 167 ++++++++++++++++++
3 files changed, 218 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 77eda4062ec0c1..f4c39c4502bf98 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,25 @@ 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.
+ Operation *firstOp = sibNode->op;
+ Operation *secondOp = dstNode->op;
+
+ if (secondOp->isBeforeInBlock(firstOp))
+ std::swap(firstOp, secondOp);
+
+ for (Value result : firstOp->getResults()) {
+ for (Operation *user : result.getUsers()) {
+ // Find the ancestor operation in the siblings' block.
+ while (user && user->getBlock() != firstOp->getBlock())
+ user = user->getParentOp();
+
+ if (user && (user == secondOp || user->isBeforeInBlock(secondOp)))
+ return false;
+ }
+ }
+
// Skip if there exists a path of dependent edges between
// 'sibNode' and 'dstNode'.
if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
@@ -1431,8 +1456,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 +1469,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 +1486,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 +1505,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 +1538,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..cab22638d96c77 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -21,3 +21,170 @@ 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 that placing an SSA usage in a nested region prevents the fusion
+
+// CHECK-LABEL: func.func @intervening_nested_result_user
+// CHECK: %[[FIRST:[0-9]+]] = affine.for
+// CHECK: affine.for
+// CHECK: affine.store %[[FIRST]],
+// CHECK: %[[SECOND:[0-9]+]] = affine.for
+
+func.func @intervening_nested_result_user(
+ %a: memref<8xf32>,
+ %out: memref<4xf32>) -> f32 {
+ %zero = arith.constant 0.0 : f32
+
+ %first = affine.for %i = 0 to 8
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<8xf32>
+ %next = arith.addf %acc, %x : f32
+ affine.yield %next : f32
+ }
+
+ affine.for %j = 0 to 4 {
+ affine.store %first, %out[%j] : memref<4xf32>
+ }
+
+ %second = affine.for %i = 0 to 8
+ iter_args(%acc = %zero) -> f32 {
+ %x = affine.load %a[%i] : memref<8xf32>
+ %next = arith.addf %acc, %x : f32
+ affine.yield %next : 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