[Mlir-commits] [mlir] [MLIR][Affine] Fix fusion across ops with unknown memory effects (PR #203231)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jun 11 03:02:22 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Aditya Pradhan (pradhanaditya)
<details>
<summary>Changes</summary>
The MemRefDependenceGraph answered the question "which memrefs does this op access" in two different ways. When recording accesses, it uses getMayEffectedValues, which is conservative about ops with unknown memory effects (ops without a MemoryEffectOpInterface, such as calls to external functions): they are assumed to affect all their memref operands. But when constructing dependence edges, the Node::getLoadOpCount/getStoreOpCount/hasStore/hasFree methods answered it via hasEffect, which returns false for any op without the interface.
Because of this mismatch, an external call was recorded as accessing its memref operands, yet reported no reads or writes when edges were created. The missing edges let loop fusion fuse a producer nest and a consumer nest across a call sitting between them, even when the call may write the buffer they share. In the @<!-- -->zero_tolerance case of loop-fusion-4.mlir, one nest fills a buffer, the external call may overwrite that buffer, and a later nest reads it back out. Fusion moved the read-out nest ahead of the external call, so it read the buffer's pre-call contents; if the call did write the buffer, that update was never seen.
Fix this by adding mayHaveEffect, which simply checks whether the memref is among the values getMayEffectedValues returns for the op, and use it in the four Node methods. Both sides now share the same logic and can no longer disagree.
---
Full diff: https://github.com/llvm/llvm-project/pull/203231.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Affine/Analysis/Utils.cpp (+46-34)
- (modified) mlir/test/Dialect/Affine/loop-fusion-4.mlir (+11-2)
``````````diff
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index ebe932a14694a..013478c9fa44c 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -38,6 +38,45 @@ using llvm::SmallDenseMap;
using Node = MemRefDependenceGraph::Node;
+/// Returns the values that `op` may have a memref effect of type `EffectTys`
+/// on, not considering recursive effects. An op with unknown memory effects
+/// (e.g. a call to an external function without a memory-effect interface) is
+/// conservatively assumed to affect all its memref operands.
+template <typename... EffectTys>
+static void getMayEffectedValues(Operation *op,
+ SmallVectorImpl<Value> &values) {
+ auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!memOp) {
+ if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
+ // No effects.
+ return;
+ // Memref operands have to be considered as being affected.
+ for (Value operand : op->getOperands()) {
+ if (isa<MemRefType>(operand.getType()))
+ values.push_back(operand);
+ }
+ return;
+ }
+ SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
+ memOp.getEffects(effects);
+ for (auto &effect : effects) {
+ Value effectVal = effect.getValue();
+ if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
+ isa<MemRefType>(effectVal.getType()))
+ values.push_back(effectVal);
+ };
+}
+
+/// Returns true if `op` may have a memory effect of type `EffectTys` on
+/// `memref`, i.e., whether `memref` is among the values returned by
+/// `getMayEffectedValues` for `op`.
+template <typename... EffectTys>
+static bool mayHaveEffect(Operation *op, Value memref) {
+ SmallVector<Value> values;
+ getMayEffectedValues<EffectTys...>(op, values);
+ return llvm::is_contained(values, memref);
+}
+
// LoopNestStateCollector walks loop nests and collects load and store
// operations, and whether or not a region holding op other than ForOp and IfOp
// was encountered in the loop nest.
@@ -83,7 +122,7 @@ unsigned Node::getLoadOpCount(Value memref) const {
if (auto affineLoad = dyn_cast<AffineReadOpInterface>(loadOp)) {
if (memref == affineLoad.getMemRef())
++loadOpCount;
- } else if (hasEffect<MemoryEffects::Read>(loadOp, memref)) {
+ } else if (mayHaveEffect<MemoryEffects::Read>(loadOp, memref)) {
++loadOpCount;
}
}
@@ -98,8 +137,7 @@ unsigned Node::getStoreOpCount(Value memref) const {
if (auto affineStore = dyn_cast<AffineWriteOpInterface>(storeOp)) {
if (memref == affineStore.getMemRef())
++storeOpCount;
- } else if (hasEffect<MemoryEffects::Write>(const_cast<Operation *>(storeOp),
- memref)) {
+ } else if (mayHaveEffect<MemoryEffects::Write>(storeOp, memref)) {
++storeOpCount;
}
}
@@ -114,7 +152,7 @@ unsigned Node::hasStore(Value memref) const {
if (auto affineStore = dyn_cast<AffineWriteOpInterface>(storeOp)) {
if (memref == affineStore.getMemRef())
return true;
- } else if (hasEffect<MemoryEffects::Write>(storeOp, memref)) {
+ } else if (mayHaveEffect<MemoryEffects::Write>(storeOp, memref)) {
return true;
}
return false;
@@ -123,7 +161,7 @@ unsigned Node::hasStore(Value memref) const {
unsigned Node::hasFree(Value memref) const {
return llvm::any_of(memrefFrees, [&](Operation *freeOp) {
- return hasEffect<MemoryEffects::Free>(freeOp, memref);
+ return mayHaveEffect<MemoryEffects::Free>(freeOp, memref);
});
}
@@ -160,32 +198,6 @@ void Node::getLoadAndStoreMemrefSet(
}
}
-/// Returns the values that this op has a memref effect of type `EffectTys` on,
-/// not considering recursive effects.
-template <typename... EffectTys>
-static void getEffectedValues(Operation *op, SmallVectorImpl<Value> &values) {
- auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
- if (!memOp) {
- if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
- // No effects.
- return;
- // Memref operands have to be considered as being affected.
- for (Value operand : op->getOperands()) {
- if (isa<MemRefType>(operand.getType()))
- values.push_back(operand);
- }
- return;
- }
- SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
- memOp.getEffects(effects);
- for (auto &effect : effects) {
- Value effectVal = effect.getValue();
- if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
- isa<MemRefType>(effectVal.getType()))
- values.push_back(effectVal);
- };
-}
-
/// Add `op` to MDG creating a new node and adding its memory accesses (affine
/// or non-affine to memrefAccesses (memref -> list of nodes with accesses) map.
static Node *
@@ -210,7 +222,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefLoads) {
SmallVector<Value> effectedValues;
- getEffectedValues<MemoryEffects::Read>(op, effectedValues);
+ getMayEffectedValues<MemoryEffects::Read>(op, effectedValues);
if (llvm::any_of(((ValueRange)effectedValues).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
// We do not know the interaction here.
@@ -221,7 +233,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefStores) {
SmallVector<Value> effectedValues;
- getEffectedValues<MemoryEffects::Write>(op, effectedValues);
+ getMayEffectedValues<MemoryEffects::Write>(op, effectedValues);
if (llvm::any_of((ValueRange(effectedValues)).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
return nullptr;
@@ -231,7 +243,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefFrees) {
SmallVector<Value> effectedValues;
- getEffectedValues<MemoryEffects::Free>(op, effectedValues);
+ getMayEffectedValues<MemoryEffects::Free>(op, effectedValues);
if (llvm::any_of((ValueRange(effectedValues)).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
return nullptr;
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index cf530016c201a..dc03681458e33 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -583,11 +583,15 @@ func.func @zero_tolerance(%arg0: memref<65536xcomplex<f64>>, %arg1: memref<30x13
affine.store %18, %2[%arg2] : memref<131072xi128>
affine.store %13, %1[%arg2] : memref<131072xi1>
}
- // The next two nests are fused.
+ // The next nest cannot fuse with the one following it across the opaque
+ // external call below, which may write to its memref operand in place.
// ZERO-TOLERANCE: affine.for %{{.*}} = 0 to 30
// ZERO-TOLERANCE-NEXT: affine.for %{{.*}} = 0 to 131072
// ZERO-TOLERANCE: func.call @__external_reduce_barrett
// ZERO-TOLERANCE: affine.store
+ // ZERO-TOLERANCE: call @__external_levelwise_forward_ntt
+ // ZERO-TOLERANCE-NEXT: affine.for %{{.*}} = 0 to 30
+ // ZERO-TOLERANCE-NEXT: affine.for %{{.*}} = 0 to 131072
// ZERO-TOLERANCE: affine.load
// ZERO-TOLERANCE-NEXT: affine.store
affine.for %arg2 = 0 to 30 {
@@ -611,9 +615,14 @@ func.func @zero_tolerance(%arg0: memref<65536xcomplex<f64>>, %arg1: memref<30x13
affine.store %7, %arg1[%arg2, %arg3] : memref<30x131072xi64>
}
}
- // Under maximal fusion, just one nest.
+ // Under maximal fusion, the first two nests fuse, but the last nest cannot
+ // fuse into them across the opaque external call, which may write to its
+ // memref operand in place.
// PRODUCER-CONSUMER-MAXIMAL: affine.for %{{.*}} = 0 to 30
// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 131072
+ // PRODUCER-CONSUMER-MAXIMAL: call @__external_levelwise_forward_ntt
+ // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 30
+ // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 131072
// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for %{{.*}}
memref.dealloc %2 : memref<131072xi128>
memref.dealloc %1 : memref<131072xi1>
``````````
</details>
https://github.com/llvm/llvm-project/pull/203231
More information about the Mlir-commits
mailing list