[Mlir-commits] [mlir] [mlir][affine] Fix sibling fusion after destination loop replacement (PR #213061)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 30 09:27:27 PDT 2026
https://github.com/Vaisman created https://github.com/llvm/llvm-project/pull/213061
This fixes a crash in affine sibling loop fusion when promotion replaces the
destination loop. fuseLoops previously returned void, so callers continued
using the detached destination operation.
Make fuseLoops return the replacement AffineForOp and update its callers.
Additionally, preserve external uses of sibling loop reduction results by
redirecting them to the appended destination loop results before erasing the
original sibling loop.
Testing:
- mlir/test/Dialect/Affine: 71/71 passed
>From a99f7b45f662d03836c092aea3bbbc09914118e0 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Thu, 30 Jul 2026 10:50:30 +0200
Subject: [PATCH] [mlir][affine] Fix sibling fusion after destination loop
replacement
This fixes a crash in affine sibling loop fusion when promotion replaces the
destination loop. fuseLoops previously returned void, so callers continued
using the detached destination operation.
Make fuseLoops return the replacement AffineForOp and update its callers.
Additionally, preserve external uses of sibling loop reduction results by
redirecting them to the appended destination loop results before erasing the
original sibling loop.
---
.../mlir/Dialect/Affine/LoopFusionUtils.h | 10 +--
.../Dialect/Affine/Transforms/LoopFusion.cpp | 26 +++++--
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 24 ++++---
.../Dialect/Affine/loop-fusion-sibling.mlir | 69 +++++++++++++++++++
.../lib/Dialect/Affine/TestLoopFusion.cpp | 2 +-
5 files changed, 112 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h b/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
index 0ef39fd7d1463..936bbd13fbaa5 100644
--- a/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
+++ b/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
@@ -113,11 +113,13 @@ canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp, unsigned dstLoopDepth,
/// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion
/// point and source slice loop bounds specified in 'srcSlice'.
-/// `isInnermostSiblingInsertionFusion` enables cleanup of `srcForOp that is a
+/// Returns the current destination loop, which may differ from 'dstForOp' if
+/// fusion replaces it while promoting a single-iteration reduction loop.
+/// `isInnermostSiblingInsertionFusion` enables cleanup of `srcForOp` that is a
/// single-iteration reduction loop being sibling-fused into a 'dstForOp'.
-void fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
- const ComputationSliceState &srcSlice,
- bool isInnermostSiblingInsertionFusion = false);
+AffineForOp fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
+ const ComputationSliceState &srcSlice,
+ bool isInnermostSiblingInsertionFusion = false);
/// LoopNestStats aggregates various per-loop statistics (eg. loop trip count
/// and operation count) for a loop nest up until (and including) the innermost
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 1ec5fbfef50c3..bc761b8a1721f 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -1087,7 +1087,8 @@ struct GreedyFusion {
}
// Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
- fuseLoops(srcAffineForOp, dstAffineForOp, bestSlice);
+ dstAffineForOp = fuseLoops(srcAffineForOp, dstAffineForOp, bestSlice);
+ dstNode->op = dstAffineForOp;
dstNodeChanged = true;
LDBG() << "Fused src loop " << srcId << " into dst loop " << dstId
@@ -1331,11 +1332,26 @@ struct GreedyFusion {
// destination loop. Based on this, the fused loop may be optimized
// further inside `fuseLoops`.
bool isInnermostInsertion = (bestDstLoopDepth == dstLoopDepthTest);
- // Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
- affine::fuseLoops(sibAffineForOp, dstAffineForOp, bestSlice,
- isInnermostInsertion);
+ unsigned oldDstNumResults = dstAffineForOp.getNumResults();
+ dstAffineForOp = affine::fuseLoops(sibAffineForOp, dstAffineForOp,
+ bestSlice, isInnermostInsertion);
+ dstNode->op = dstAffineForOp;
+
+ unsigned numAddedResults = sibAffineForOp.getNumResults();
+ assert(dstAffineForOp.getNumResults() ==
+ oldDstNumResults + numAddedResults &&
+ "unexpected destination loop results after sibling fusion");
+
+ // Sibling fusion appends the promoted reduction results to the existing
+ // destination loop results in the same order. Redirect external uses
+ // before erasing the original sibling loop.
+ for (unsigned i = 0; i < numAddedResults; ++i) {
+ sibAffineForOp.getResult(i).replaceAllUsesWith(
+ dstAffineForOp.getResult(oldDstNumResults + i));
+ }
+
+ auto dstForInst = dstAffineForOp;
- auto dstForInst = cast<AffineForOp>(dstNode->op);
// Update operation position of fused loop nest (if needed).
if (insertPointInst != dstForInst)
dstForInst->moveBefore(insertPointInst);
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 68296ea3368a1..9fd6096561053 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -354,8 +354,8 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
/// Patch the loop body of a forOp that is a single iteration reduction loop
/// into its containing block.
-static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
- bool siblingFusionUser) {
+static FailureOr<AffineForOp>
+promoteSingleIterReductionLoop(AffineForOp forOp, bool siblingFusionUser) {
// Check if the reduction loop is a single iteration loop.
std::optional<APInt> tripCount = forOp.getStaticTripCount();
if (!tripCount || *tripCount != 1)
@@ -417,14 +417,14 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
parentBlock->getOperations().splice(Block::iterator(forOp),
forOp.getBody()->getOperations());
forOp.erase();
- return success();
+ return newLoop;
}
/// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion point
/// and source slice loop bounds specified in 'srcSlice'.
-void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
- const ComputationSliceState &srcSlice,
- bool isInnermostSiblingInsertion) {
+AffineForOp mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
+ const ComputationSliceState &srcSlice,
+ bool isInnermostSiblingInsertion) {
// Clone 'srcForOp' into 'dstForOp' at 'srcSlice->insertPoint'.
OpBuilder b(srcSlice.insertPoint->getBlock(), srcSlice.insertPoint);
IRMapping mapper;
@@ -458,14 +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
+ AffineForOp parentForOp = forOp->getParentOfType<AffineForOp>();
+ FailureOr<AffineForOp> newParentForOp =
+ promoteSingleIterReductionLoop(forOp, true);
+ if (succeeded(newParentForOp) && parentForOp == dstForOp)
+ dstForOp = *newParentForOp;
+ } else {
// Promote any single iteration slice loops.
(void)promoteIfSingleIteration(forOp);
+ }
}
+ return dstForOp;
}
/// Collect loop nest statistics (eg. loop trip count and operation count)
diff --git a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
index 937c855b86b50..20763bbca5403 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -21,3 +21,72 @@ func.func @disjoint_stores(%0: memref<8xf32>) {
// CHECK-NOT: affine.for
return
}
+
+// CHECK-LABEL: func.func @sibling_reduction_result
+// CHECK: %[[FUSED:.*]] = affine.for {{.*}} iter_args
+// CHECK: affine.store
+// CHECK: affine.yield
+// CHECK: %[[SECOND_REDUCTION:.*]] = affine.for {{.*}} iter_args
+// CHECK: %[[SUM:.*]] = arith.addi %[[FUSED]], %[[SECOND_REDUCTION]] : i64
+// CHECK: arith.trunci %[[SUM]] : i64 to i32
+func.func @sibling_reduction_result() -> i32 {
+ %c7_i64 = arith.constant 7 : i64
+ %c3_i64 = arith.constant 3 : i64
+ %c0_i64 = arith.constant 0 : i64
+ %c97_i64 = arith.constant 97 : i64
+ %c1_i64 = arith.constant 1 : i64
+
+ %alloc = memref.alloc() : memref<8xi64>
+ %alloc_0 = memref.alloc() : memref<8xi64>
+ %alloc_1 = memref.alloc() : memref<8xi64>
+
+ affine.for %arg0 = 0 to 8 {
+ %4 = arith.index_cast %arg0 : index to i64
+ %5 = arith.addi %4, %c1_i64 : i64
+ %6 = arith.remsi %5, %c97_i64 : i64
+ affine.store %6, %alloc[%arg0] : memref<8xi64>
+ }
+
+ affine.for %arg0 = 0 to 8 {
+ affine.store %c0_i64, %alloc_0[%arg0] : memref<8xi64>
+ }
+
+ affine.for %arg0 = 0 to 8 {
+ affine.store %c0_i64, %alloc_1[%arg0] : memref<8xi64>
+ }
+
+ affine.for %arg0 = 0 to 8 {
+ %4 = affine.load %alloc[%arg0] : memref<8xi64>
+ %5 = arith.muli %4, %c3_i64 : i64
+ affine.store %5, %alloc_0[%arg0] : memref<8xi64>
+ }
+
+ affine.for %arg0 = 0 to 8 {
+ %4 = affine.load %alloc_0[%arg0] : memref<8xi64>
+ %5 = arith.addi %4, %c7_i64 : i64
+ affine.store %5, %alloc_1[%arg0] : memref<8xi64>
+ }
+
+ %0 = affine.for %arg0 = 0 to 8
+ iter_args(%arg1 = %c0_i64) -> (i64) {
+ %4 = affine.load %alloc_0[%arg0] : memref<8xi64>
+ %5 = arith.addi %arg1, %4 : i64
+ affine.yield %5 : i64
+ }
+
+ %1 = affine.for %arg0 = 0 to 8
+ iter_args(%arg1 = %c0_i64) -> (i64) {
+ %4 = affine.load %alloc_1[%arg0] : memref<8xi64>
+ %5 = arith.addi %arg1, %4 : i64
+ affine.yield %5 : i64
+ }
+
+ %2 = arith.addi %0, %1 : i64
+ %3 = arith.trunci %2 : i64 to i32
+
+ memref.dealloc %alloc : memref<8xi64>
+ memref.dealloc %alloc_0 : memref<8xi64>
+ memref.dealloc %alloc_1 : memref<8xi64>
+
+ return %3 : i32
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp b/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp
index bf11d94596fa7..988c592ed4466 100644
--- a/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp
@@ -144,7 +144,7 @@ static bool testLoopFusionUtilities(AffineForOp forOpA, AffineForOp forOpB,
FusionResult result = canFuseLoops(forOpA, forOpB, d, &sliceUnion);
if (result.value != FusionResult::Success)
continue;
- fuseLoops(forOpA, forOpB, sliceUnion);
+ forOpB = fuseLoops(forOpA, forOpB, sliceUnion);
// Note: 'forOpA' is removed to simplify test output. A proper loop
// fusion pass should perform additional checks to check safe removal.
if (forOpA.use_empty())
More information about the Mlir-commits
mailing list