[Mlir-commits] [mlir] [MLIR][Affine] Check access coverage before fusion privatization (PR #212080)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jul 25 18:33:02 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Mingfei Guo (guoriyue)
<details>
<summary>Changes</summary>
This fixes an affine loop fusion bug where privatization could create a buffer that is too small for consumer accesses.
The pass now checks that producer writes cover all consumer reads. If not, fusion continues but unsafe privatization is skipped.
---
Patch is 21.58 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/212080.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp (+161-53)
- (modified) mlir/test/Dialect/Affine/loop-fusion-4.mlir (+126)
- (added) mlir/test/mlir-runner/affine-loop-fusion.mlir (+69)
``````````diff
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 1ec5fbfef50c3..d84fac437faca 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/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/Utils.h"
@@ -272,6 +273,77 @@ getDominanceFilterForPrivateMemRefRepl(Block *sliceInsertionBlock,
return firstAncestor;
}
+static std::optional<presburger::IntegerRelation>
+getAccessRelationForPrivateBuffer(Operation *op, unsigned loopDepth) {
+ if (!isa<AffineLoadOp, AffineStoreOp>(op))
+ return std::nullopt;
+
+ // Reject domains that getAccessRelation cannot represent exactly.
+ Operation *ancestor = op->getParentOp();
+ for (; ancestor && !ancestor->hasTrait<OpTrait::AffineScope>();
+ ancestor = ancestor->getParentOp()) {
+ if (auto forOp = dyn_cast<AffineForOp>(ancestor)) {
+ if (forOp.getStepAsInt() != 1 && !forOp.hasConstantLowerBound())
+ return std::nullopt;
+ } else if (auto parallelOp = dyn_cast<AffineParallelOp>(ancestor)) {
+ if (!llvm::all_of(parallelOp.getSteps(),
+ [](int64_t step) { return step == 1; }))
+ return std::nullopt;
+ } else {
+ return std::nullopt;
+ }
+ }
+ if (!ancestor)
+ return std::nullopt;
+
+ presburger::IntegerRelation relation(
+ presburger::PresburgerSpace::getRelationSpace());
+ if (failed(MemRefAccess(op).getAccessRelation(relation)))
+ return std::nullopt;
+
+ if (getNestingDepth(op) != relation.getNumDomainVars() ||
+ loopDepth > relation.getNumDomainVars())
+ return std::nullopt;
+
+ relation.convertToLocal(presburger::VarKind::Domain, loopDepth,
+ relation.getNumDomainVars());
+ return relation;
+}
+
+/// Returns true if every load is covered by `store` in each private-buffer
+/// instance.
+static bool areLoadsCoveredByStore(Operation *store,
+ ArrayRef<Operation *> loads,
+ unsigned loopDepth) {
+ assert(store && "expected producer store");
+ assert(!loads.empty() && "expected consumer load");
+
+ auto write = getAccessRelationForPrivateBuffer(store, loopDepth);
+ if (!write)
+ return false;
+ for (Operation *load : loads) {
+ auto read = getAccessRelationForPrivateBuffer(load, loopDepth);
+ if (!read ||
+ !write->getSpace().isAligned(read->getSpace(),
+ presburger::VarKind::Domain) ||
+ write->getNumRangeVars() != read->getNumRangeVars())
+ return false;
+
+ write->mergeAndAlignSymbols(*read);
+ if (!read->isSubsetOf(*write))
+ return false;
+ }
+ return true;
+}
+
+static bool haveSameAccessFunction(ArrayRef<Operation *> stores) {
+ assert(!stores.empty() && "expected producer store");
+ MemRefAccess first(cast<AffineWriteOpInterface>(stores.front()));
+ return llvm::all_of(llvm::drop_begin(stores), [&](Operation *store) {
+ return first == MemRefAccess(cast<AffineWriteOpInterface>(store));
+ });
+}
+
/// Returns the amount of additional (redundant) computation that will be done
/// as a fraction of the total computation if `srcForOp` is fused into
/// `dstForOp` at depth `depth`. The method returns the compute cost of the
@@ -327,41 +399,60 @@ static std::optional<double> getAdditionalComputeFraction(
// Creates and returns a private (single-user) memref for fused loop rooted at
// 'forOp', with (potentially reduced) memref size based on the memref region
// written to by `storeOps` at depth 'dstLoopDepth'. 'sliceInsertionBlock'
-// specifies the block in which the slice was/will be inserted. The method
-// expects that all stores ops to the memref have the same access function.
-// Returns nullptr if the creation failed.
+// specifies the block in which the slice was/will be inserted. Returns nullptr
+// unless the fused producer covers every consumer access that would be
+// redirected to the private memref.
static Value createPrivateMemRef(AffineForOp forOp,
ArrayRef<Operation *> storeOps,
+ Operation *fusedProducerStore,
unsigned dstLoopDepth,
std::optional<unsigned> fastMemorySpace,
Block *sliceInsertionBlock,
uint64_t localBufSizeThreshold) {
assert(!storeOps.empty() && "no source stores supplied");
- // Check if all stores have the same access function; we only support this
- // case.
- // TODO: Use union of memref write regions to compute private memref footprint
- // for store ops with different access functions.
- if (storeOps.size() > 1 &&
- !std::equal(std::next(storeOps.begin()), storeOps.end(), storeOps.begin(),
- [](Operation *a, Operation *b) {
- MemRefAccess aM(cast<AffineWriteOpInterface>(a));
- MemRefAccess bM(cast<AffineWriteOpInterface>(b));
- return aM == bM;
- })) {
+ // TODO: Use a union of write regions to support stores with different access
+ // functions.
+ if (!fusedProducerStore)
+ return nullptr;
+ if (!haveSameAccessFunction(storeOps)) {
LDBG() << "Private memref creation unsupported for multiple producer "
- << "stores with different access functions.";
+ "stores with different access functions.";
return nullptr;
}
- Operation *srcStoreOp = storeOps[0];
+ Operation *srcStoreOp = storeOps.front();
+ Value oldMemRef = cast<AffineWriteOpInterface>(srcStoreOp).getMemRef();
+ if (cast<AffineWriteOpInterface>(fusedProducerStore).getMemRef() != oldMemRef)
+ return nullptr;
+
+ // Match the exact set of users that replaceAllMemRefUsesWith will rewrite.
+ // Privatization is only legal when the fused producer writes every element
+ // read from each private-buffer instance.
+ Operation *domFilter =
+ getDominanceFilterForPrivateMemRefRepl(sliceInsertionBlock, storeOps);
+ DominanceInfo domInfo(domFilter->getParentOfType<FunctionOpInterface>());
+ SmallVector<Operation *, 4> loads;
+ for (Operation *user : oldMemRef.getUsers()) {
+ if (!domInfo.dominates(domFilter, user) ||
+ hasSingleEffect<MemoryEffects::Free>(user, oldMemRef))
+ continue;
+ if (isa<AffineLoadOp>(user)) {
+ loads.push_back(user);
+ continue;
+ }
+ if (!isa<AffineStoreOp>(user) || !llvm::is_contained(storeOps, user))
+ return nullptr;
+ }
+ if (loads.empty() ||
+ !areLoadsCoveredByStore(fusedProducerStore, loads, dstLoopDepth))
+ return nullptr;
// Create builder to insert alloc op just before 'forOp'.
OpBuilder b(forOp);
// Builder to create constants at the top level.
OpBuilder top(forOp->getParentRegion());
// Create new memref type based on slice bounds.
- auto oldMemRef = cast<AffineWriteOpInterface>(srcStoreOp).getMemRef();
auto oldMemRefType = cast<MemRefType>(oldMemRef.getType());
unsigned rank = oldMemRefType.getRank();
@@ -441,12 +532,8 @@ static Value createPrivateMemRef(AffineForOp forOp,
AffineMap::get(outerIVs.size() + rank, 0, remapExprs, forOp.getContext());
// Replace all users of 'oldMemRef' with 'newMemRef'.
- Operation *domFilter =
- getDominanceFilterForPrivateMemRefRepl(sliceInsertionBlock, storeOps);
auto userFilterFn = [&](Operation *user) {
- auto domInfo = std::make_unique<DominanceInfo>(
- domFilter->getParentOfType<FunctionOpInterface>());
- return domInfo->dominates(domFilter, user);
+ return domInfo.dominates(domFilter, user);
};
LogicalResult res = replaceAllMemRefUsesWith(
oldMemRef, newMemRef, /*extraIndices=*/{}, indexRemap,
@@ -1075,17 +1162,22 @@ struct GreedyFusion {
srcId, dstId, bestSlice, fusedLoopInsPoint, srcEscapingMemRefs,
*mdg);
- DenseSet<Value> privateMemrefs;
+ DenseSet<Value> privateMemRefCandidates;
for (Value memref : producerConsumerMemrefs) {
if (canCreatePrivateMemRef(memref, srcEscapingMemRefs, srcId, dstId,
removeSrcNode)) {
- // Create a private version of this memref.
- LDBG() << "Creating private memref for " << memref;
- // Create a private version of this memref.
- privateMemrefs.insert(memref);
+ LDBG() << "Considering private memref for " << memref;
+ privateMemRefCandidates.insert(memref);
}
}
+ DenseSet<Operation *> preexistingDstStores;
+ if (!privateMemRefCandidates.empty()) {
+ dstAffineForOp.walk([&](AffineWriteOpInterface storeOp) {
+ preexistingDstStores.insert(storeOp.getOperation());
+ });
+ }
+
// Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
fuseLoops(srcAffineForOp, dstAffineForOp, bestSlice);
dstNodeChanged = true;
@@ -1098,35 +1190,51 @@ struct GreedyFusion {
if (fusedLoopInsPoint != dstAffineForOp)
dstAffineForOp->moveBefore(fusedLoopInsPoint);
- // Update edges between 'srcNode' and 'dstNode'.
- mdg->updateEdges(srcNode->id, dstNode->id, privateMemrefs,
- removeSrcNode);
-
- // Create private memrefs.
- if (!privateMemrefs.empty()) {
- // Note the block into which fusion was performed. This can be used to
- // place `alloc`s that create private memrefs.
- Block *sliceInsertionBlock = bestSlice.insertPoint->getBlock();
-
- // Gather stores for all the private-to-be memrefs.
- DenseMap<Value, SmallVector<Operation *, 4>> privateMemRefToStores;
+ // Privatization rewrites the consumer accesses to a buffer sized from
+ // the fused producer stores. Only do so when the producer's exact
+ // access set covers every rewritten consumer load. Otherwise, retain
+ // the legal fusion while keeping accesses on the original memref.
+ unsigned privateMemRefLoopDepth =
+ getNestingDepth(dstAffineForOp) + bestDstLoopDepth;
+ Block *sliceInsertionBlock = bestSlice.insertPoint->getBlock();
+ DenseMap<Value, SmallVector<Operation *, 4>> privateMemRefToStores;
+ DenseMap<Value, Operation *> privateMemRefToFusedProducerStore;
+ if (!privateMemRefCandidates.empty()) {
dstAffineForOp.walk([&](AffineWriteOpInterface storeOp) {
- Value storeMemRef = storeOp.getMemRef();
- if (privateMemrefs.count(storeMemRef) > 0)
- privateMemRefToStores[storeMemRef].push_back(storeOp);
+ Value memref = storeOp.getMemRef();
+ if (!privateMemRefCandidates.contains(memref))
+ return;
+ privateMemRefToStores[memref].push_back(storeOp);
+ if (!preexistingDstStores.contains(storeOp))
+ privateMemRefToFusedProducerStore.try_emplace(memref, storeOp);
});
+ }
+
+ // Create only private memrefs whose producer covers every rewritten
+ // load, and update the dependence graph with exactly those memrefs.
+ DenseSet<Value> privateMemrefs;
+ SmallVector<Value, 4> newPrivateMemrefs;
+ for (auto &[oldMemRef, stores] : privateMemRefToStores) {
+ Value newMemRef = createPrivateMemRef(
+ dstAffineForOp, stores,
+ privateMemRefToFusedProducerStore.lookup(oldMemRef),
+ privateMemRefLoopDepth, fastMemorySpace, sliceInsertionBlock,
+ localBufSizeThreshold);
+ if (!newMemRef) {
+ LDBG() << "Skipping private memref with unsupported or uncovered "
+ "consumer accesses: "
+ << oldMemRef;
+ continue;
+ }
+ privateMemrefs.insert(oldMemRef);
+ newPrivateMemrefs.push_back(newMemRef);
+ }
+
+ // Update edges between 'srcNode' and 'dstNode'.
+ mdg->updateEdges(srcId, dstId, privateMemrefs, removeSrcNode);
- // Replace original memrefs with private memrefs. Note that all the
- // loads and stores on these memrefs will be replaced with a new
- // loads and stores. Any reference to the original ones becomes
- // invalid after this point.
- for (auto &memrefToStoresPair : privateMemRefToStores) {
- ArrayRef<Operation *> storesForMemref = memrefToStoresPair.second;
- Value newMemRef = createPrivateMemRef(
- dstAffineForOp, storesForMemref, bestDstLoopDepth,
- fastMemorySpace, sliceInsertionBlock, localBufSizeThreshold);
- if (!newMemRef)
- continue;
+ if (!newPrivateMemrefs.empty()) {
+ for (Value newMemRef : newPrivateMemrefs) {
// Create new node in dependence graph for 'newMemRef' alloc op.
unsigned newMemRefNodeId = mdg->addNode(newMemRef.getDefiningOp());
// Add edge from 'newMemRef' node to dstNode.
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index cf530016c201a..c216fc14b9590 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -884,3 +884,129 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
}
return %alloc : memref<1024x8192xf32>
}
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/212039.
+// Do not privatize a producer buffer when its exact, strided write set does
+// not cover the consumer read set. Fusion itself remains legal.
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func.func @private_buffer_strided_holes(
+// PRODUCER-CONSUMER-MAXIMAL-NOT: memref.alloc
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}} : memref<32xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.load {{.*}} : memref<32xf64>
+// PRODUCER-CONSUMER-MAXIMAL-NOT: memref.alloc
+// PRODUCER-CONSUMER-MAXIMAL: return
+func.func @private_buffer_strided_holes(
+ %in: memref<32xf64>, %comm: memref<32xf64>,
+ %out: memref<32xf64>) {
+ affine.for %i = 0 to 8 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ affine.store %a, %comm[2 * %i] : memref<32xf64>
+ }
+ affine.for %j = 0 to 16 {
+ %b = affine.load %comm[%j] : memref<32xf64>
+ affine.store %b, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
+// -----
+
+// Preserve privatization when the strided producer covers every read.
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func.func @private_buffer_strided_coverage(
+// PRODUCER-CONSUMER-MAXIMAL: memref.alloc() : memref<1xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}}[0] : memref<1xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.load {{.*}}[0] : memref<1xf64>
+func.func @private_buffer_strided_coverage(
+ %in: memref<16xf64>, %out: memref<16xf64>) {
+ %comm = memref.alloc() : memref<32xf64>
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<16xf64>
+ affine.store %a, %comm[2 * %i] : memref<32xf64>
+ }
+ affine.for %j = 0 to 16 {
+ %b = affine.load %comm[2 * %j] : memref<32xf64>
+ affine.store %b, %out[%j] : memref<16xf64>
+ }
+ return
+}
+
+// -----
+
+// Although the second read's union across all outer iterations equals the
+// producer write union, it is not covered in the same private-buffer instance.
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func.func @private_buffer_outer_iv_mismatch(
+// PRODUCER-CONSUMER-MAXIMAL-NOT: memref<1xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: memref.alloc() : memref<64xf64>
+// PRODUCER-CONSUMER-MAXIMAL-NOT: memref<1xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.for %[[J:.*]] = 0 to 16 {
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %[[I:.*]] = 0 to 16 {
+// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}} : memref<64xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.load {{.*}} : memref<64xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.load {{.*}} : memref<64xf64>
+// PRODUCER-CONSUMER-MAXIMAL-NOT: memref<1xf64>
+// PRODUCER-CONSUMER-MAXIMAL: return
+func.func @private_buffer_outer_iv_mismatch(
+ %in: memref<2x4x16xf64>, %out: memref<2x4x16xf64>) {
+ affine.for %instance = 0 to 2 {
+ affine.for %batch = 0 to 4 {
+ %comm = memref.alloc() : memref<64xf64>
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%instance, %batch, %i] : memref<2x4x16xf64>
+ affine.store %a, %comm[16 * %batch + %i] : memref<64xf64>
+ }
+ affine.for %j = 0 to 16 {
+ %a = affine.load %comm[16 * %batch + %j] : memref<64xf64>
+ %b = affine.load %comm[-16 * %batch + %j + 48] : memref<64xf64>
+ %sum = arith.addf %a, %b : f64
+ affine.store %sum, %out[%instance, %batch, %j]
+ : memref<2x4x16xf64>
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// Access relations conservatively approximate a non-unit affine.for with a
+// dynamic lower bound. Do not use that approximation to prove coverage.
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func.func @private_buffer_dynamic_lb_stride(
+// PRODUCER-CONSUMER-MAXIMAL-NOT: memref.alloc
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.for {{.*}} step 2
+// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}} : memref<4x?xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.for {{.*}} step 2
+// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}} : memref<4x?xf64>
+// PRODUCER-CONSUMER-MAXIMAL: affine.load {{.*}} : memref<4x?xf64>
+// PRODUCER-CONSUMER-MAXIMAL-NOT: memref.alloc
+// PRODUCER-CONSUMER-MAXIMAL: return
+func.func @private_buffer_dynamic_lb_stride(
+ %lb: index, %in: memref<4xf64>, %comm: memref<4x?xf64>,
+ %out: memref<4xf64>) {
+ affine.for %batch = 0 to 4 {
+ affine.for %k = %lb to affine_map<(d0) -> (d0 + 8)>(%lb) step 2 {
+ %a = affine.load %in[%batch] : memref<4xf64>
+ affine.store %a, %comm[%batch, %k] : memref<4x?xf64>
+ }
+ }
+ affine.for %batch = 0 to 4 {
+ %a = affine.load %comm[%batch, %lb] : memref<4x?xf64>
+ %b = affine.load %comm[%batch, %lb + 1] : memref<4x?xf64>
+ %c = affine.load %comm[%batch, %lb + 2] : memref<4x?xf64>
+ %d = affine.load %comm[%batch, %lb + 3] : memref<4x?xf64>
+ %ab = arith.addf %a, %b : f64
+ %cd = arith.addf %c, %d : f64
+ %sum = arith.addf %ab, %cd : f64
+ affine.store %sum, %out[%batch] : memref<4xf64>
+ }
+ return
+}
diff --git a/mlir/test/mlir-runner/affine-loop-fusion.mlir b/mlir/test/mlir-runner/affine-loop-fusion.mlir
new file mode 100644
index 0000000000000..ecf1858d5c456
--- /dev/null
+++ b/mlir/test/mlir-runner/affine-loop-fusion.mlir
@@ -0,0 +1,69 @@
+// RUN: mlir-opt %s \
+// RUN: -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer maximal}),func.func(lower-affine),generate-runtime-verification,convert-scf-to-cf,convert-to-llvm)' | \
+// RUN: mlir-runner -e main -entry-point-result=i64 \
+// RUN: %if target={{s390x-.*}} %{ -argext-abi-check=false %} | FileCheck %s
+// XFAIL: system-aix
+
+// CHECK: {{^0$}}
+
+// The producer overwrites the even elements; the odd elements keep their
+// initial values. Runtime verification makes an undersized private buffer fail
+// deterministically instead of relying on undefined out-of-bounds behavior.
+memref.global "private" constant @expected : memref<16xf64> =
+ dense<[1000.0, 101.0, 1001.0, 103.0, 1002.0, 105.0, 1003.0, 107.0,
+ 1004.0, 109.0, 1005.0, 111.0, 1006.0, 113.0, 1007.0, 115.0]>
+
+func.func @kernel(%in: memref<32xf64>, %comm: memref<32xf64>,
+ %out: memref<32xf64>) {
+ affine.for %i = 0 to 8 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ affine.store %a, %comm[2 * %i] : memref<32xf64>
+ }
+ affine.for %j = 0 to 16 {
+ %b = affine.load %comm[%...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/212080
More information about the Mlir-commits
mailing list