[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:55:54 PDT 2026
https://github.com/speryt created https://github.com/llvm/llvm-project/pull/226759
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
>From 3c2b79ecda69f859c1099a351594b4173f59ee6a Mon Sep 17 00:00:00 2001
From: s-peryt <1652033+s-peryt at users.noreply.github.com>
Date: Sat, 26 Sep 2026 22:37:05 -0700
Subject: [PATCH] [mlir][affine] Fix sibling fusion with multiple loads to the
shared memref
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 wiht helper function and tests
---
.../Dialect/Affine/Transforms/LoopFusion.cpp | 11 +-
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 55 ++++++-
.../Dialect/Affine/loop-fusion-sibling.mlir | 147 +++++++++++++++---
3 files changed, 177 insertions(+), 36 deletions(-)
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
+}
More information about the Mlir-commits
mailing list