[Mlir-commits] [mlir] [MLIR][Affine] Fix fusion of producers that update a memref in place (PR #210533)
Hamza Qureshi
llvmlistbot at llvm.org
Sat Jul 18 22:29:53 PDT 2026
https://github.com/hamzaqureshi5 updated https://github.com/llvm/llvm-project/pull/210533
>From 40a7e8c8b2158d4b695cf33c2581b5de67fab6f5 Mon Sep 17 00:00:00 2001
From: Hamza Qureshi <hamza7771.861 at gmail.com>
Date: Sat, 18 Jul 2026 23:17:43 +0500
Subject: [PATCH] [MLIR][Affine] Fix fusion of producers that update a memref
in place
Affine producer-consumer fusion places the fused nest at an insertion
point between the source and the destination nest. When the source nest
is retained after fusion (e.g., non-maximal fusion with an escaping
memref) and updates a memref in place (loads from and stores to it),
this placement is unsound: the fused nest recomputes the source's
values reading from that memref, but the retained source nest has
already overwritten it, so the source's computation is applied twice.
Fix this by placing the fused nest *before* the retained source nest,
where the in-place updated memrefs still hold their original values.
The recomputed stores are redirected to a private memref, so the
retained source nest also still reads the original values. An
anti-dependence edge from the destination to the source node is added
to the dependence graph to preserve this placement in subsequent fusion
decisions. When no such placement can be shown to preserve all
dependences, or the in-place updated memref can't be privatized, bail
out on the fusion candidate.
Loads that only read values stored earlier in the same iteration (e.g.,
initialize-then-reuse patterns) don't observe the memref's incoming
state and don't require the new placement.
Update @same_memref_load_store and @same_memref_load_multiple_stores,
whose CHECK lines encoded the miscompiled output; their fused nests are
unchanged but now precede the retained source nests.
Fixes: https://github.com/llvm/llvm-project/issues/210490
---
.../Dialect/Affine/Transforms/LoopFusion.cpp | 140 ++++++++++++++++++
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 87 ++++++++++-
2 files changed, 221 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 1ec5fbfef50c3..5ec6ef7998431 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -124,6 +124,107 @@ static bool canRemoveSrcNodeAfterFusion(
return true;
}
+/// Returns true if `node`'s loads of `memref` may observe the memory state
+/// of `memref` from before the node's execution. This doesn't hold when
+/// every load of `memref` is dominated by a store to the same location
+/// within the same iteration (e.g., an initialize-then-reuse pattern): all
+/// such loads only read values produced by the node itself.
+static bool mayReadIncomingMemState(const MemRefDependenceGraph::Node &node,
+ Value memref,
+ const DominanceInfo &domInfo) {
+ for (Operation *loadOp : node.loads) {
+ if (cast<AffineReadOpInterface>(loadOp).getMemRef() != memref)
+ continue;
+ MemRefAccess loadAccess(loadOp);
+ bool coveredByStore = llvm::any_of(node.stores, [&](Operation *storeOp) {
+ return cast<AffineWriteOpInterface>(storeOp).getMemRef() == memref &&
+ domInfo.properlyDominates(storeOp, loadOp) &&
+ MemRefAccess(storeOp) == loadAccess;
+ });
+ if (!coveredByStore)
+ return true;
+ }
+ return false;
+}
+
+/// Returns true if the fused loop nest can be placed just before the source
+/// nest `srcId` while preserving all dependences. This placement is needed
+/// when the source nest is retained after fusion and updates the memrefs in
+/// `inPlaceUpdatedMemrefs` in place (loads from and stores to them): the
+/// fused nest recomputes the source's values reading from these memrefs,
+/// which is only sound while they still hold their original values, i.e.,
+/// before the retained source nest overwrites them. See
+/// https://github.com/llvm/llvm-project/issues/210490.
+static bool
+canPlaceFusedNestBeforeSrcNode(const MemRefDependenceGraph &mdg, unsigned srcId,
+ unsigned dstId,
+ ArrayRef<Value> inPlaceUpdatedMemrefs,
+ const DenseSet<Value> &privateMemrefs) {
+ const MemRefDependenceGraph::Node *srcNode = mdg.getNode(srcId);
+ const MemRefDependenceGraph::Node *dstNode = mdg.getNode(dstId);
+ Operation *srcOp = srcNode->op;
+ Operation *dstOp = dstNode->op;
+
+ // The recomputed stores to an in-place updated memref must be redirected to
+ // a private memref; otherwise they would clobber values the retained source
+ // nest still has to read. Private memref creation happens only after
+ // fusion, so check its requirements here.
+ for (Value memref : inPlaceUpdatedMemrefs) {
+ // The memref must have been selected for privatization.
+ if (!privateMemrefs.contains(memref))
+ return false;
+ // Stores to the memref in the destination nest would in general prevent
+ // private memref creation as they don't share the source stores' access
+ // function.
+ if (dstNode->getStoreOpCount(memref) > 0)
+ return false;
+ // Private memref creation requires all stores to have the same access
+ // function.
+ SmallVector<Operation *, 4> storeOps;
+ srcNode->getStoreOpsForMemref(memref, &storeOps);
+ if (storeOps.size() > 1 &&
+ !std::equal(std::next(storeOps.begin()), storeOps.end(),
+ storeOps.begin(), [](Operation *a, Operation *b) {
+ return MemRefAccess(cast<AffineWriteOpInterface>(a)) ==
+ MemRefAccess(cast<AffineWriteOpInterface>(b));
+ }))
+ return false;
+ }
+
+ // All dependences from the source nest to the destination nest must be
+ // satisfied by the privatized recomputation; any other dependence would be
+ // violated by the new placement.
+ for (const auto &edge : mdg.outEdges.lookup(srcId))
+ if (edge.id == dstId && !privateMemrefs.contains(edge.value))
+ return false;
+
+ // The new placement makes the fused nest cross all graph nodes between the
+ // source and the destination nest; conservatively require that there are
+ // none.
+ for (const auto &idAndNode : mdg.nodes) {
+ Operation *op = idAndNode.second.op;
+ if (op->getBlock() == srcOp->getBlock() && srcOp->isBeforeInBlock(op) &&
+ op->isBeforeInBlock(dstOp))
+ return false;
+ }
+
+ // All SSA values used in the destination nest must dominate the source nest
+ // for the destination nest to be placed before it.
+ DominanceInfo domInfo;
+ WalkResult walkResult = dstOp->walk([&](Operation *op) {
+ for (Value operand : op->getOperands()) {
+ Operation *defOp = operand.getDefiningOp();
+ // Block arguments and values defined within the destination nest remain
+ // available; any other value must dominate the source nest.
+ if (defOp && !dstOp->isAncestor(defOp) &&
+ !domInfo.properlyDominates(defOp, srcOp))
+ return WalkResult::interrupt();
+ }
+ return WalkResult::advance();
+ });
+ return !walkResult.wasInterrupted();
+}
+
/// Returns in 'srcIdCandidates' the producer fusion candidates for consumer
/// 'dstId'. Candidates are sorted by node id order. This order corresponds to
/// the program order when the 'mdg' is created. However, program order is not
@@ -1086,6 +1187,38 @@ struct GreedyFusion {
}
}
+ // Collect memrefs the source nest updates in place: memrefs it
+ // stores to while also reading their incoming values. If the source
+ // nest is retained after fusion, the fused nest has to be placed
+ // before it: the fused nest recomputes the source's values reading
+ // from these memrefs, which is only sound while they still hold
+ // their original values, i.e., before the retained source nest
+ // overwrites them (https://github.com/llvm/llvm-project/issues/210490).
+ // Bail out on this candidate if no legal such placement exists.
+ SmallVector<Value, 2> inPlaceUpdatedMemrefs;
+ if (!removeSrcNode) {
+ DominanceInfo domInfo;
+ for (Operation *store : srcNode->stores) {
+ Value memref = cast<AffineWriteOpInterface>(store).getMemRef();
+ if (srcNode->getLoadOpCount(memref) > 0 &&
+ !llvm::is_contained(inPlaceUpdatedMemrefs, memref) &&
+ mayReadIncomingMemState(*srcNode, memref, domInfo))
+ inPlaceUpdatedMemrefs.push_back(memref);
+ }
+ }
+ if (!inPlaceUpdatedMemrefs.empty()) {
+ if (!canPlaceFusedNestBeforeSrcNode(
+ *mdg, srcId, dstId, inPlaceUpdatedMemrefs, privateMemrefs)) {
+ LDBG() << "Can't fuse: the source nest updates memref(s) in "
+ "place, is retained after fusion, and the fused nest "
+ "can't be placed before it";
+ continue;
+ }
+ // Place the fused nest right before the retained source nest so
+ // that it reads the in-place updated memrefs' original values.
+ fusedLoopInsPoint = srcNode->op;
+ }
+
// Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
fuseLoops(srcAffineForOp, dstAffineForOp, bestSlice);
dstNodeChanged = true;
@@ -1102,6 +1235,13 @@ struct GreedyFusion {
mdg->updateEdges(srcNode->id, dstNode->id, privateMemrefs,
removeSrcNode);
+ // If the fused nest was placed before the retained source nest, it
+ // reads the original values of the memrefs the source nest updates
+ // in place: record these anti-dependences to preserve the placement
+ // in subsequent fusion decisions.
+ for (Value memref : inPlaceUpdatedMemrefs)
+ mdg->addEdge(dstId, srcId, memref);
+
// Create private memrefs.
if (!privateMemrefs.empty()) {
// Note the block into which fusion was performed. This can be used to
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index cf530016c201a..0b63b83d21d73 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -296,8 +296,6 @@ module {
// PRODUCER-CONSUMER-LABEL: func @same_memref_load_store
func.func @same_memref_load_store(%producer : memref<32xf32>, %consumer: memref<16xf32>){
%cst = arith.constant 2.000000e+00 : f32
- // Source isn't removed.
- // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 32
affine.for %arg3 = 0 to 32 {
%0 = affine.load %producer[%arg3] : memref<32xf32>
%2 = arith.mulf %0, %cst : f32
@@ -308,6 +306,9 @@ func.func @same_memref_load_store(%producer : memref<32xf32>, %consumer: memref<
%2 = arith.addf %0, %cst : f32
affine.store %2, %consumer[%arg3] : memref<16xf32>
}
+ // The source nest updates %producer in place and isn't removed; the fused
+ // nest is placed before it so that its recomputation of the source's
+ // values reads %producer's original values (issue #210490).
// Fused nest.
// PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 16
// PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>
@@ -317,6 +318,12 @@ func.func @same_memref_load_store(%producer : memref<32xf32>, %consumer: memref<
// PRODUCER-CONSUMER-NEXT: arith.addf
// PRODUCER-CONSUMER-NEXT: affine.store
// PRODUCER-CONSUMER-NEXT: }
+ // Source isn't removed, and runs after the fused nest.
+ // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 32
+ // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>
+ // PRODUCER-CONSUMER-NEXT: arith.mulf
+ // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<32xf32>
+ // PRODUCER-CONSUMER-NEXT: }
return
}
@@ -326,10 +333,6 @@ func.func @same_memref_load_store(%producer : memref<32xf32>, %consumer: memref<
// ALL-LABEL: func @same_memref_load_multiple_stores
func.func @same_memref_load_multiple_stores(%producer : memref<32xf32>, %producer_2 : memref<32xf32>, %consumer: memref<16xf32>){
%cst = arith.constant 2.000000e+00 : f32
- // Ensure that source isn't removed during both producer-consumer fusion and
- // sibling fusion.
- // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 32
- // ALL: affine.for %{{.*}} = 0 to 32
affine.for %arg3 = 0 to 32 {
%0 = affine.load %producer[%arg3] : memref<32xf32>
%2 = arith.mulf %0, %cst : f32
@@ -342,6 +345,10 @@ func.func @same_memref_load_multiple_stores(%producer : memref<32xf32>, %produce
%2 = arith.addf %0, %1 : f32
affine.store %2, %consumer[%arg3] : memref<16xf32>
}
+ // The source nest updates %producer in place and isn't removed during both
+ // producer-consumer fusion and sibling fusion; the fused nest is placed
+ // before it so that its recomputation of the source's values reads
+ // %producer's original values (issue #210490).
// Fused nest.
// PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 16
// PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>
@@ -356,6 +363,74 @@ func.func @same_memref_load_multiple_stores(%producer : memref<32xf32>, %produce
// ALL: affine.for %{{.*}} = 0 to 16
// ALL: mulf
// ALL: addf
+ // Source isn't removed, and runs after the fused nest.
+ // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 32
+ // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>
+ // PRODUCER-CONSUMER-NEXT: arith.mulf
+ // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<32xf32>
+ // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<32xf32>
+ // PRODUCER-CONSUMER-NEXT: }
+ // ALL: affine.for %{{.*}} = 0 to 32
+ return
+}
+
+// -----
+
+// Reproducer for https://github.com/llvm/llvm-project/issues/210490: the
+// producer nest updates %p in place and is retained after fusion since the
+// fusion isn't maximal and %p escapes. The fused nest must be placed before
+// the retained producer nest so that its recomputation of the producer's
+// values reads %p's original values; placing it after would apply the
+// multiplication twice on the consumed prefix of %p.
+// PRODUCER-CONSUMER-LABEL: func @in_place_producer_prefix_consumer
+func.func @in_place_producer_prefix_consumer(%p: memref<4xf32>, %d: memref<2xf32>, %c: f32) {
+ affine.for %i = 0 to 4 {
+ %0 = affine.load %p[%i] : memref<4xf32>
+ %1 = arith.mulf %0, %c : f32
+ affine.store %1, %p[%i] : memref<4xf32>
+ }
+ affine.for %i = 0 to 2 {
+ %0 = affine.load %p[%i] : memref<4xf32>
+ affine.store %0, %d[%i] : memref<2xf32>
+ }
+ // Fused nest, placed before the retained producer nest.
+ // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 2
+ // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<4xf32>
+ // PRODUCER-CONSUMER-NEXT: arith.mulf
+ // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
+ // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[0] : memref<1xf32>
+ // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<2xf32>
+ // PRODUCER-CONSUMER-NEXT: }
+ // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 4
+ // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<4xf32>
+ // PRODUCER-CONSUMER-NEXT: arith.mulf
+ // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<4xf32>
+ // PRODUCER-CONSUMER-NEXT: }
+ return
+}
+
+// -----
+
+// The consumer nest also stores to the in-place updated memref %p: %p can't
+// be privatized in the fused nest, so there is no legal placement for the
+// fused nest before the retained producer nest. Fusion must not happen.
+// PRODUCER-CONSUMER-LABEL: func @in_place_producer_storing_consumer
+func.func @in_place_producer_storing_consumer(%p: memref<4xf32>, %d: memref<2xf32>, %c: f32) {
+ affine.for %i = 0 to 4 {
+ %0 = affine.load %p[%i] : memref<4xf32>
+ %1 = arith.mulf %0, %c : f32
+ affine.store %1, %p[%i] : memref<4xf32>
+ }
+ affine.for %i = 0 to 2 {
+ %0 = affine.load %p[%i] : memref<4xf32>
+ affine.store %0, %d[%i] : memref<2xf32>
+ affine.store %c, %p[%i] : memref<4xf32>
+ }
+ // PRODUCER-CONSUMER-NOT: memref.alloc
+ // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 4
+ // PRODUCER-CONSUMER: arith.mulf
+ // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 2
+ // PRODUCER-CONSUMER-NOT: arith.mulf
return
}
More information about the Mlir-commits
mailing list