[Mlir-commits] [mlir] [mlir][affine] Fix sibling fusion with multiple loads (PR #226759)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Sep 26 22:56:43 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
Author: speryt
<details>
<summary>Changes</summary>
Sibling fusion missed nests that load the shared memref more than once, e.g. two loops that both read `A[i]` and `A[i + 1]`. Such siblings were filtered out, and even when one of the nests had a single load, the slice was computed as the union over all load pairs. That union spans several source iterations per destination iteration and gets rejected as an incorrect slice.
Sibling nests don't depend on each other, so compute the slice from a single pair of loads instead, preferring one that gives a maximal slice, and drop the single-load restriction.
Fixes #<!-- -->219356
Used Claude to help with helper function and tests
---
Full diff: https://github.com/llvm/llvm-project/pull/226759.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp (+5-6)
- (modified) mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp (+48-7)
- (modified) mlir/test/Dialect/Affine/loop-fusion-sibling.mlir (+124-23)
``````````diff
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 77eda4062ec0c..af23a66bad747 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -1220,11 +1220,11 @@ struct GreedyFusion {
// Check if fusion would be profitable and at what depth.
- // Get unique 'sibNode' load op to 'memref'.
+ // Get a 'sibNode' load op to 'memref'. Its footprint is used to evaluate
+ // reuse for fusion profitability.
SmallVector<Operation *, 2> sibLoadOpInsts;
sibNode->getLoadOpsForMemref(memref, &sibLoadOpInsts);
- // Currently findSiblingNodeToFuse searches for siblings with one load.
- Operation *sibLoadOpInst = llvm::getSingleElement(sibLoadOpInsts);
+ Operation *sibLoadOpInst = sibLoadOpInsts.front();
// Gather 'dstNode' load ops to 'memref'.
SmallVector<Operation *, 2> dstLoadOpInsts;
@@ -1365,9 +1365,8 @@ struct GreedyFusion {
// Returns true if 'sibNode' can be fused with 'dstNode' for input reuse
// on 'memref'.
auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
- // Skip if 'outEdge' is not a read-after-write dependence.
- // TODO: Remove restrict to single load op restriction.
- if (sibNode->getLoadOpCount(memref) != 1)
+ // Skip if 'sibNode' doesn't load from 'memref'.
+ if (sibNode->getLoadOpCount(memref) == 0)
return false;
// Skip if there exists a path of dependent edges between
// 'sibNode' and 'dstNode'.
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 68296ea3368a1..2006fdd476543 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -242,6 +242,50 @@ static unsigned getMaxLoopDepth(ArrayRef<Operation *> srcOps,
return loopDepth;
}
+/// Computes in 'srcSlice' the sibling fusion slice from a single pair of loads
+/// to 'memref' in 'opsA' and 'opsB', preferring a maximal one. Read-after-read
+/// pairs don't constrain legality, so the slices of all pairs aren't unioned.
+static FusionResult computeSiblingSlice(ArrayRef<Operation *> opsA,
+ ArrayRef<Operation *> opsB,
+ Value memref, unsigned dstLoopDepth,
+ unsigned numCommonLoops,
+ bool isBackwardSlice,
+ ComputationSliceState *srcSlice) {
+ auto isLoadFromMemRef = [&](Operation *op) {
+ auto load = dyn_cast<AffineReadOpInterface>(op);
+ return load && load.getMemRef() == memref;
+ };
+
+ std::optional<ComputationSliceState> validSlice;
+ FusionResult failure = FusionResult::FailPrecondition;
+ for (Operation *a : llvm::make_filter_range(opsA, isLoadFromMemRef)) {
+ for (Operation *b : llvm::make_filter_range(opsB, isLoadFromMemRef)) {
+ ComputationSliceState slice;
+ SliceComputationResult result = affine::computeSliceUnion(
+ a, b, dstLoopDepth, numCommonLoops, isBackwardSlice, &slice);
+ if (result.value == SliceComputationResult::IncorrectSliceFailure) {
+ failure = FusionResult::FailIncorrectSlice;
+ continue;
+ }
+ if (result.value != SliceComputationResult::Success)
+ continue;
+ if (slice.isMaximal().value_or(false)) {
+ *srcSlice = slice;
+ return FusionResult::Success;
+ }
+ if (!validSlice)
+ validSlice = slice;
+ }
+ }
+
+ if (!validSlice) {
+ LDBG() << "No valid sibling fusion slice";
+ return failure;
+ }
+ *srcSlice = *validSlice;
+ return FusionResult::Success;
+}
+
// TODO: This pass performs some computation that is the same for all the depths
// (e.g., getMaxLoopDepth). Implement a version of this utility that processes
// all the depths at once or only the legal maximal depth for maximal fusion.
@@ -325,13 +369,10 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
break;
case FusionStrategy::Sibling:
// Sibling fusion (AffineLoopFusion pass) only takes into account the loads
- // to 'memref' in 'srcForOp' to compute the slice union.
- for (Operation *op : opsA) {
- auto load = dyn_cast<AffineReadOpInterface>(op);
- if (load && load.getMemRef() == fusionStrategy.getSiblingFusionMemRef())
- strategyOpsA.push_back(op);
- }
- break;
+ // to 'memref' to compute the slice.
+ return computeSiblingSlice(
+ opsA, opsB, fusionStrategy.getSiblingFusionMemRef(), dstLoopDepth,
+ numCommonLoops, isSrcForOpBeforeDstForOp, srcSlice);
}
// Compute union of computation slices computed between all pairs of ops
diff --git a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
index 937c855b86b50..5da3a85af1975 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -1,23 +1,124 @@
-// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{maximal mode=sibling}))' | FileCheck %s
-
-// Test cases specifically for sibling fusion. Note that sibling fusion test
-// cases also exist in loop-fusion*.mlir.
-
-// CHECK-LABEL: func @disjoint_stores
-func.func @disjoint_stores(%0: memref<8xf32>) {
- %alloc_1 = memref.alloc() : memref<16xf32>
- // The affine stores below are to different parts of the memrefs. Sibling
- // fusion helps improve reuse and is valid.
- affine.for %arg2 = 0 to 8 {
- %2 = affine.load %0[%arg2] : memref<8xf32>
- affine.store %2, %alloc_1[%arg2] : memref<16xf32>
- }
- affine.for %arg2 = 0 to 8 {
- %2 = affine.load %0[%arg2] : memref<8xf32>
- %3 = arith.negf %2 : f32
- affine.store %3, %alloc_1[%arg2 + 8] : memref<16xf32>
- }
- // CHECK: affine.for
- // CHECK-NOT: affine.for
- return
-}
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{maximal mode=sibling}))' -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s
+
+// Test cases specifically for sibling fusion. Note that sibling fusion test
+// cases also exist in loop-fusion*.mlir.
+
+// CHECK-LABEL: func @disjoint_stores
+func.func @disjoint_stores(%0: memref<8xf32>) {
+ %alloc_1 = memref.alloc() : memref<16xf32>
+ // The affine stores below are to different parts of the memrefs. Sibling
+ // fusion helps improve reuse and is valid.
+ affine.for %arg2 = 0 to 8 {
+ %2 = affine.load %0[%arg2] : memref<8xf32>
+ affine.store %2, %alloc_1[%arg2] : memref<16xf32>
+ }
+ affine.for %arg2 = 0 to 8 {
+ %2 = affine.load %0[%arg2] : memref<8xf32>
+ %3 = arith.negf %2 : f32
+ affine.store %3, %alloc_1[%arg2 + 8] : memref<16xf32>
+ }
+ // CHECK: affine.for
+ // CHECK-NOT: affine.for
+ return
+}
+
+// -----
+
+// Sibling nests with multiple loads to the shared memref (at different
+// offsets) are fused.
+
+// CHECK-LABEL: func @multiple_loads
+// CHECK-SAME: (%{{.*}}: memref<65xf32>, %[[B:[a-z0-9]+]]: memref<64xf32>, %[[C:[a-z0-9]+]]: memref<64xf32>)
+func.func @multiple_loads(%a: memref<65xf32>, %b: memref<64xf32>, %c: memref<64xf32>) {
+ affine.for %i = 0 to 64 {
+ %0 = affine.load %a[%i] : memref<65xf32>
+ %1 = affine.load %a[%i + 1] : memref<65xf32>
+ %2 = arith.addf %0, %1 : f32
+ affine.store %2, %b[%i] : memref<64xf32>
+ }
+ affine.for %i = 0 to 64 {
+ %0 = affine.load %a[%i] : memref<65xf32>
+ %1 = affine.load %a[%i + 1] : memref<65xf32>
+ %2 = arith.mulf %0, %1 : f32
+ affine.store %2, %c[%i] : memref<64xf32>
+ }
+ // CHECK: affine.for %[[I:.*]] = 0 to 64 {
+ // CHECK-NEXT: affine.load %{{.*}}[%[[I]]]
+ // CHECK-NEXT: affine.load %{{.*}}[%[[I]] + 1]
+ // CHECK-NEXT: arith.addf
+ // CHECK-NEXT: affine.store %{{.*}}, %[[B]][%[[I]]]
+ // CHECK-NEXT: affine.load %{{.*}}[%[[I]]]
+ // CHECK-NEXT: affine.load %{{.*}}[%[[I]] + 1]
+ // CHECK-NEXT: arith.mulf
+ // CHECK-NEXT: affine.store %{{.*}}, %[[C]][%[[I]]]
+ // CHECK-NEXT: }
+ // CHECK-NEXT: return
+ return
+}
+
+// -----
+
+// Only one of the sibling nests has multiple loads to the shared memref.
+
+// CHECK-LABEL: func @single_and_multiple_loads
+// CHECK-SAME: (%{{.*}}: memref<65xf32>, %[[B:[a-z0-9]+]]: memref<64xf32>, %[[C:[a-z0-9]+]]: memref<64xf32>)
+func.func @single_and_multiple_loads(%a: memref<65xf32>, %b: memref<64xf32>, %c: memref<64xf32>) {
+ affine.for %i = 0 to 64 {
+ %0 = affine.load %a[%i] : memref<65xf32>
+ %1 = affine.load %a[%i + 1] : memref<65xf32>
+ %2 = arith.addf %0, %1 : f32
+ affine.store %2, %b[%i] : memref<64xf32>
+ }
+ affine.for %i = 0 to 64 {
+ %0 = affine.load %a[%i] : memref<65xf32>
+ %1 = arith.mulf %0, %0 : f32
+ affine.store %1, %c[%i] : memref<64xf32>
+ }
+ // CHECK: affine.for %[[I:.*]] = 0 to 64 {
+ // CHECK-NEXT: affine.load %{{.*}}[%[[I]]]
+ // CHECK-NEXT: affine.load %{{.*}}[%[[I]] + 1]
+ // CHECK-NEXT: arith.addf
+ // CHECK-NEXT: affine.store %{{.*}}, %[[B]][%[[I]]]
+ // CHECK-NEXT: affine.load %{{.*}}[%[[I]]]
+ // CHECK-NEXT: arith.mulf
+ // CHECK-NEXT: affine.store %{{.*}}, %[[C]][%[[I]]]
+ // CHECK-NEXT: }
+ // CHECK-NEXT: return
+ return
+}
+
+// -----
+
+// 2-d sibling nests with multiple loads to the shared memref are fused at the
+// innermost depth.
+
+// CHECK-LABEL: func @multiple_loads_2d
+// CHECK-SAME: (%{{.*}}: memref<33x32xf32>, %[[B:[a-z0-9]+]]: memref<32x32xf32>, %[[C:[a-z0-9]+]]: memref<32x32xf32>)
+func.func @multiple_loads_2d(%a: memref<33x32xf32>, %b: memref<32x32xf32>, %c: memref<32x32xf32>) {
+ affine.for %i = 0 to 32 {
+ affine.for %j = 0 to 32 {
+ %0 = affine.load %a[%i, %j] : memref<33x32xf32>
+ %1 = affine.load %a[%i + 1, %j] : memref<33x32xf32>
+ %2 = arith.addf %0, %1 : f32
+ affine.store %2, %b[%i, %j] : memref<32x32xf32>
+ }
+ }
+ affine.for %i = 0 to 32 {
+ affine.for %j = 0 to 32 {
+ %0 = affine.load %a[%i, %j] : memref<33x32xf32>
+ %1 = affine.load %a[%i + 1, %j] : memref<33x32xf32>
+ %2 = arith.mulf %0, %1 : f32
+ affine.store %2, %c[%i, %j] : memref<32x32xf32>
+ }
+ }
+ // CHECK: affine.for %[[I:.*]] = 0 to 32 {
+ // CHECK-NEXT: affine.for %[[J:.*]] = 0 to 32 {
+ // CHECK: affine.store %{{.*}}, %[[B]][%[[I]], %[[J]]]
+ // CHECK-NOT: affine.for
+ // CHECK: affine.store %{{.*}}, %[[C]][%[[I]], %[[J]]]
+ // CHECK-NEXT: }
+ // CHECK-NEXT: }
+ // CHECK-NEXT: return
+ return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/226759
More information about the Mlir-commits
mailing list