[Mlir-commits] [mlir] [MLIR][Affine] Handle cast aliases in fusion dependencies (PR #213413)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 1 06:42:13 PDT 2026
https://github.com/1sgtpepper updated https://github.com/llvm/llvm-project/pull/213413
>From 8de53cd7a3c3064ea89e5c454b869e6a660410e2 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 13:59:03 +0800
Subject: [PATCH 1/8] [MLIR][Affine] Handle cast aliases in fusion dependencies
---
.../Dialect/Affine/Analysis/CMakeLists.txt | 1 +
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 85 +++++++++++--------
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 41 ++++++++-
3 files changed, 91 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt b/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt
index 3a1996349dbed..6c52f7751f1f0 100644
--- a/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt
@@ -18,6 +18,7 @@ add_mlir_dialect_library(MLIRAffineAnalysis
MLIRControlFlowInterfaces
MLIRDialectUtils
MLIRInferTypeOpInterface
+ MLIRMemRefUtils
MLIRSideEffectInterfaces
MLIRPresburger
MLIRSCFDialect
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 321c8e34d907c..11e412f8d8704 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -19,6 +19,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IntegerSet.h"
#include "llvm/ADT/SetVector.h"
@@ -38,6 +39,49 @@ 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. Fully aliasing
+/// views are canonicalized so the MDG uses one key for the view and its source.
+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(memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(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(memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(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::skipFullyAliasingOperations(cast<MemrefValue>(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 +127,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 +142,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 +157,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 +166,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 +203,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 +227,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 +238,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 +248,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..980e68e6b9cfa 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>
@@ -884,3 +893,31 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
}
return %alloc : memref<1024x8192xf32>
}
+
+// -----
+
+// The external call receives a fully aliasing cast of the producer's memref.
+// Fusion must preserve the call between the producer and consumer loops.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_external_call
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: memref.cast
+// PRODUCER-CONSUMER-MAXIMAL: call @escape
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @cast_alias_external_call(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %comm[%i] : memref<32xf64>
+ }
+ %view = memref.cast %comm : memref<32xf64> to memref<?xf64>
+ func.call @escape(%view) : (memref<?xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape(memref<?xf64>)
>From ba43cc3744d3fee08e720ca8fa97ace576c368b7 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 14:13:28 +0800
Subject: [PATCH 2/8] Format Affine analysis changes
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 11e412f8d8704..793214e87f021 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -55,8 +55,8 @@ static void getMayEffectedValues(Operation *op,
// Memref operands have to be considered as being affected.
for (Value operand : op->getOperands()) {
if (isa<MemRefType>(operand.getType()))
- values.push_back(memref::skipFullyAliasingOperations(
- cast<MemrefValue>(operand)));
+ values.push_back(
+ memref::skipFullyAliasingOperations(cast<MemrefValue>(operand)));
}
return;
}
@@ -66,8 +66,8 @@ static void getMayEffectedValues(Operation *op,
Value effectVal = effect.getValue();
if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
isa<MemRefType>(effectVal.getType()))
- values.push_back(memref::skipFullyAliasingOperations(
- cast<MemrefValue>(effectVal)));
+ values.push_back(
+ memref::skipFullyAliasingOperations(cast<MemrefValue>(effectVal)));
};
}
>From 8a2d6150d4d468d73c91a16efdb978db9de9a82c Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 15:38:29 +0800
Subject: [PATCH 3/8] [MLIR][Affine] Make fusion alias handling symmetric
Use one canonical identity for trivial memref aliases throughout the dependence graph and loop fusion consumers. Keep raw views for precise affine accesses, and conservatively handle differing views and arbitrary memory-effecting operations.
---
.../mlir/Dialect/Affine/Analysis/Utils.h | 6 +-
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 152 ++++++++++--------
.../Dialect/Affine/Transforms/LoopFusion.cpp | 61 +++++--
mlir/lib/Dialect/Affine/Utils/CMakeLists.txt | 1 +
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 54 +++++--
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 88 ++++++++++
mlir/test/Dialect/Affine/loop-fusion.mlir | 41 ++++-
7 files changed, 310 insertions(+), 93 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index df4145db90a61..6a03fabd274fd 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
@@ -128,9 +128,9 @@ struct MemRefDependenceGraph {
// 'Node.outEdges[i].id' is the identifier of the dest node of the edge.
unsigned id;
// The SSA value on which this edge represents a dependence.
- // If the value is a memref, then the dependence is between graph nodes
- // which contain accesses to the same memref 'value'. If the value is a
- // non-memref value, then the dependence is between a graph node which
+ // If the value is a memref, then it is the canonical representative of
+ // the trivial alias class on which the dependence is based. If the value
+ // is a non-memref value, then the dependence is between a graph node which
// defines an SSA value and another graph node which uses the SSA value
// (e.g. a constant or load operation defining a value which is used inside
// a loop nest).
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 793214e87f021..bbc2fd0888d7c 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -39,13 +39,23 @@ using llvm::SmallDenseMap;
using Node = MemRefDependenceGraph::Node;
+static Value canonicalizeMemref(Value value) {
+ if (!value || !isa<MemRefType>(value.getType()))
+ return value;
+ return memref::skipFullyAliasingOperations(cast<MemrefValue>(value));
+}
+
+static bool isSameMemref(Value lhs, Value rhs) {
+ return canonicalizeMemref(lhs) == canonicalizeMemref(rhs);
+}
+
/// 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. Fully aliasing
/// views are canonicalized so the MDG uses one key for the view and its source.
template <typename... EffectTys>
-static void getMayEffectedValues(Operation *op,
+static void getMayAffectedValues(Operation *op,
SmallVectorImpl<Value> &values) {
auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
if (!memOp) {
@@ -55,8 +65,7 @@ static void getMayEffectedValues(Operation *op,
// Memref operands have to be considered as being affected.
for (Value operand : op->getOperands()) {
if (isa<MemRefType>(operand.getType()))
- values.push_back(
- memref::skipFullyAliasingOperations(cast<MemrefValue>(operand)));
+ values.push_back(canonicalizeMemref(operand));
}
return;
}
@@ -66,20 +75,18 @@ static void getMayEffectedValues(Operation *op,
Value effectVal = effect.getValue();
if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
isa<MemRefType>(effectVal.getType()))
- values.push_back(
- memref::skipFullyAliasingOperations(cast<MemrefValue>(effectVal)));
+ values.push_back(canonicalizeMemref(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`.
+/// `getMayAffectedValues` 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::skipFullyAliasingOperations(cast<MemrefValue>(memref)));
+ getMayAffectedValues<EffectTys...>(op, values);
+ return llvm::is_contained(values, canonicalizeMemref(memref));
}
// LoopNestStateCollector walks loop nests and collects load and store
@@ -99,11 +106,11 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
// This op itself is memory-effect free.
return;
- // Check operands. Eg. ops like the `call` op are handled here.
- for (Value v : op->getOperands()) {
- if (!isa<MemRefType>(v.getType()))
- continue;
- // Conservatively, we assume the memref is read and written to.
+ // Check operands. E.g., ops like the `call` op are handled here.
+ if (llvm::any_of(op->getOperands(), [](Value value) {
+ return isa<MemRefType>(value.getType());
+ })) {
+ // Conservatively, assume all memref operands are read and written.
memrefLoads.push_back(op);
memrefStores.push_back(op);
}
@@ -124,12 +131,9 @@ unsigned Node::getLoadOpCount(Value memref) const {
unsigned loadOpCount = 0;
for (Operation *loadOp : loads) {
// Common case: affine reads.
- if (auto affineLoad = dyn_cast<AffineReadOpInterface>(loadOp)) {
- if (memref == affineLoad.getMemRef())
+ if (auto affineLoad = dyn_cast<AffineReadOpInterface>(loadOp))
+ if (isSameMemref(memref, affineLoad.getMemRef()))
++loadOpCount;
- } else if (mayHaveEffect<MemoryEffects::Read>(loadOp, memref)) {
- ++loadOpCount;
- }
}
return loadOpCount;
}
@@ -140,7 +144,7 @@ unsigned Node::getStoreOpCount(Value memref) const {
for (auto *storeOp : llvm::concat<Operation *const>(stores, memrefStores)) {
// Common case: affine writes.
if (auto affineStore = dyn_cast<AffineWriteOpInterface>(storeOp)) {
- if (memref == affineStore.getMemRef())
+ if (isSameMemref(memref, affineStore.getMemRef()))
++storeOpCount;
} else if (mayHaveEffect<MemoryEffects::Write>(storeOp, memref)) {
++storeOpCount;
@@ -155,7 +159,7 @@ unsigned Node::hasStore(Value memref) const {
llvm::concat<Operation *const>(stores, memrefStores),
[&](Operation *storeOp) {
if (auto affineStore = dyn_cast<AffineWriteOpInterface>(storeOp)) {
- if (memref == affineStore.getMemRef())
+ if (isSameMemref(memref, affineStore.getMemRef()))
return true;
} else if (mayHaveEffect<MemoryEffects::Write>(storeOp, memref)) {
return true;
@@ -174,7 +178,8 @@ unsigned Node::hasFree(Value memref) const {
void Node::getStoreOpsForMemref(Value memref,
SmallVectorImpl<Operation *> *storeOps) const {
for (Operation *storeOp : stores) {
- if (memref == cast<AffineWriteOpInterface>(storeOp).getMemRef())
+ if (isSameMemref(memref,
+ cast<AffineWriteOpInterface>(storeOp).getMemRef()))
storeOps->push_back(storeOp);
}
}
@@ -183,7 +188,8 @@ void Node::getStoreOpsForMemref(Value memref,
void Node::getLoadOpsForMemref(Value memref,
SmallVectorImpl<Operation *> *loadOps) const {
for (Operation *loadOp : loads) {
- if (memref == cast<AffineReadOpInterface>(loadOp).getMemRef())
+ if (isSameMemref(memref,
+ cast<AffineReadOpInterface>(loadOp).getMemRef()))
loadOps->push_back(loadOp);
}
}
@@ -194,10 +200,12 @@ void Node::getLoadAndStoreMemrefSet(
DenseSet<Value> *loadAndStoreMemrefSet) const {
llvm::SmallDenseSet<Value, 2> loadMemrefs;
for (Operation *loadOp : loads) {
- loadMemrefs.insert(cast<AffineReadOpInterface>(loadOp).getMemRef());
+ loadMemrefs.insert(canonicalizeMemref(
+ cast<AffineReadOpInterface>(loadOp).getMemRef()));
}
for (Operation *storeOp : stores) {
- auto memref = cast<AffineWriteOpInterface>(storeOp).getMemRef();
+ auto memref = canonicalizeMemref(
+ cast<AffineWriteOpInterface>(storeOp).getMemRef());
if (loadMemrefs.count(memref) > 0)
loadAndStoreMemrefSet->insert(memref);
}
@@ -217,42 +225,44 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
Node &node = nodes.insert({newNodeId, Node(newNodeId, nodeOp)}).first->second;
for (Operation *op : collector.loadOpInsts) {
node.loads.push_back(op);
- auto memref = cast<AffineReadOpInterface>(op).getMemRef();
+ auto memref = canonicalizeMemref(
+ cast<AffineReadOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
}
for (Operation *op : collector.storeOpInsts) {
node.stores.push_back(op);
- auto memref = cast<AffineWriteOpInterface>(op).getMemRef();
+ auto memref = canonicalizeMemref(
+ cast<AffineWriteOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
}
for (Operation *op : collector.memrefLoads) {
- SmallVector<Value> effectedValues;
- getMayEffectedValues<MemoryEffects::Read>(op, effectedValues);
- if (llvm::any_of(((ValueRange)effectedValues).getTypes(),
+ SmallVector<Value> affectedValues;
+ getMayAffectedValues<MemoryEffects::Read>(op, affectedValues);
+ if (llvm::any_of(((ValueRange)affectedValues).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
// We do not know the interaction here.
return nullptr;
- for (Value memref : effectedValues)
+ for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
node.memrefLoads.push_back(op);
}
for (Operation *op : collector.memrefStores) {
- SmallVector<Value> effectedValues;
- getMayEffectedValues<MemoryEffects::Write>(op, effectedValues);
- if (llvm::any_of((ValueRange(effectedValues)).getTypes(),
+ SmallVector<Value> affectedValues;
+ getMayAffectedValues<MemoryEffects::Write>(op, affectedValues);
+ if (llvm::any_of((ValueRange(affectedValues)).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
return nullptr;
- for (Value memref : effectedValues)
+ for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
node.memrefStores.push_back(op);
}
for (Operation *op : collector.memrefFrees) {
- SmallVector<Value> effectedValues;
- getMayEffectedValues<MemoryEffects::Free>(op, effectedValues);
- if (llvm::any_of((ValueRange(effectedValues)).getTypes(),
+ SmallVector<Value> affectedValues;
+ getMayAffectedValues<MemoryEffects::Free>(op, affectedValues);
+ if (llvm::any_of((ValueRange(affectedValues)).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
return nullptr;
- for (Value memref : effectedValues)
+ for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
node.memrefFrees.push_back(op);
}
@@ -260,17 +270,16 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
return &node;
}
-/// Returns the memref being read/written by a memref/affine load/store op.
-static Value getMemRef(Operation *memOp) {
- if (auto memrefLoad = dyn_cast<memref::LoadOp>(memOp))
- return memrefLoad.getMemRef();
- if (auto affineLoad = dyn_cast<AffineReadOpInterface>(memOp))
- return affineLoad.getMemRef();
- if (auto memrefStore = dyn_cast<memref::StoreOp>(memOp))
- return memrefStore.getMemRef();
- if (auto affineStore = dyn_cast<AffineWriteOpInterface>(memOp))
- return affineStore.getMemRef();
- llvm_unreachable("unexpected op");
+/// Returns true if `op` may access `memref`, including through a fully aliasing
+/// view. Unknown operations are handled conservatively through their memory
+/// effects rather than assuming a particular operation class.
+static bool mayAccessMemRef(Operation *op, Value memref) {
+ if (auto affineRead = dyn_cast<AffineReadOpInterface>(op))
+ return isSameMemref(affineRead.getMemRef(), memref);
+ if (auto affineWrite = dyn_cast<AffineWriteOpInterface>(op))
+ return isSameMemref(affineWrite.getMemRef(), memref);
+ return mayHaveEffect<MemoryEffects::Read>(op, memref) ||
+ mayHaveEffect<MemoryEffects::Write>(op, memref);
}
/// Returns true if there may be a dependence on `memref` from srcNode's
@@ -289,16 +298,15 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
// true if there exists a conflicting read/write access involving such.
// Check whether there is a dependence from a source read/write op to a
- // destination read/write one; all expected to be memref/affine load/store.
+ // destination read/write one.
auto hasNonAffineDep = [&](ArrayRef<Operation *> srcMemOps,
ArrayRef<Operation *> dstMemOps) {
return llvm::any_of(srcMemOps, [&](Operation *srcOp) {
- Value srcMemref = getMemRef(srcOp);
- if (srcMemref != memref)
+ if (!mayAccessMemRef(srcOp, memref))
return false;
- return llvm::find_if(dstMemOps, [&](Operation *dstOp) {
- return srcMemref == getMemRef(dstOp);
- }) != dstMemOps.end();
+ return llvm::any_of(dstMemOps, [&](Operation *dstOp) {
+ return mayAccessMemRef(dstOp, memref);
+ });
});
};
@@ -335,13 +343,18 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
for (auto *srcMemOp :
llvm::concat<Operation *const>(srcNode.stores, srcNode.loads)) {
MemRefAccess srcAcc(srcMemOp);
- if (srcAcc.memref != memref)
+ if (!isSameMemref(srcAcc.memref, memref))
continue;
for (auto *destMemOp :
llvm::concat<Operation *const>(dstNode.stores, dstNode.loads)) {
MemRefAccess destAcc(destMemOp);
- if (destAcc.memref != memref)
+ if (!isSameMemref(destAcc.memref, memref))
+ continue;
+ if (srcAcc.memref != destAcc.memref) {
+ if (srcAcc.isStore() || destAcc.isStore())
+ return true;
continue;
+ }
// Check for a top-level dependence between srcNode and destNode's ops.
if (!noDependence(checkMemrefAccessDependence(
srcAcc, destAcc, getNestingDepth(srcNode.op) + 1)))
@@ -354,7 +367,7 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
bool MemRefDependenceGraph::init(bool fullAffineDependences) {
LDBG() << "--- Initializing MDG ---";
// Map from a memref to the set of ids of the nodes that have ops accessing
- // the memref.
+ // the memref. Fully aliasing views use their canonical source value here.
DenseMap<Value, SetVector<unsigned>> memrefAccesses;
// Create graph nodes.
@@ -369,14 +382,16 @@ bool MemRefDependenceGraph::init(bool fullAffineDependences) {
// Create graph node for top-level load op.
Node node(nextNodeId++, &op);
node.loads.push_back(&op);
- auto memref = cast<AffineReadOpInterface>(op).getMemRef();
+ auto memref = canonicalizeMemref(
+ cast<AffineReadOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
nodes.insert({node.id, node});
} else if (isa<AffineWriteOpInterface>(op)) {
// Create graph node for top-level store op.
Node node(nextNodeId++, &op);
node.stores.push_back(&op);
- auto memref = cast<AffineWriteOpInterface>(op).getMemRef();
+ auto memref = canonicalizeMemref(
+ cast<AffineWriteOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
nodes.insert({node.id, node});
} else if (op.getNumResults() > 0 && !op.use_empty()) {
@@ -537,6 +552,7 @@ bool MemRefDependenceGraph::hasEdge(unsigned srcId, unsigned dstId,
if (!outEdges.contains(srcId) || !inEdges.contains(dstId)) {
return false;
}
+ value = canonicalizeMemref(value);
bool hasOutEdge = llvm::any_of(outEdges.lookup(srcId), [=](const Edge &edge) {
return edge.id == dstId && (!value || edge.value == value);
});
@@ -549,6 +565,7 @@ bool MemRefDependenceGraph::hasEdge(unsigned srcId, unsigned dstId,
// Adds an edge from node 'srcId' to node 'dstId' for 'value'.
void MemRefDependenceGraph::addEdge(unsigned srcId, unsigned dstId,
Value value) {
+ value = canonicalizeMemref(value);
if (!hasEdge(srcId, dstId, value)) {
outEdges[srcId].push_back({dstId, value});
inEdges[dstId].push_back({srcId, value});
@@ -562,6 +579,7 @@ void MemRefDependenceGraph::removeEdge(unsigned srcId, unsigned dstId,
Value value) {
assert(inEdges.count(dstId) > 0);
assert(outEdges.count(srcId) > 0);
+ value = canonicalizeMemref(value);
if (isa<MemRefType>(value.getType())) {
assert(memrefEdgeCount.count(value) > 0);
memrefEdgeCount[value]--;
@@ -624,7 +642,7 @@ unsigned MemRefDependenceGraph::getIncomingMemRefAccesses(unsigned id,
Value memref) const {
unsigned inEdgeCount = 0;
for (const Edge &inEdge : inEdges.lookup(id)) {
- if (inEdge.value == memref) {
+ if (isSameMemref(inEdge.value, memref)) {
const Node *srcNode = getNode(inEdge.id);
// Only count in edges from 'srcNode' if 'srcNode' accesses 'memref'
if (srcNode->getStoreOpCount(memref) > 0)
@@ -640,7 +658,7 @@ unsigned MemRefDependenceGraph::getOutEdgeCount(unsigned id,
Value memref) const {
unsigned outEdgeCount = 0;
for (const auto &outEdge : outEdges.lookup(id))
- if (!memref || outEdge.value == memref)
+ if (!memref || isSameMemref(outEdge.value, memref))
++outEdgeCount;
return outEdgeCount;
}
@@ -743,7 +761,9 @@ void MemRefDependenceGraph::updateEdges(unsigned srcId, unsigned dstId,
SmallVector<Edge, 2> oldInEdges = inEdges[srcId];
for (auto &inEdge : oldInEdges) {
// Add edge from 'inEdge.id' to 'dstId' if it's not a private memref.
- if (!privateMemRefs.contains(inEdge.value))
+ if (!llvm::any_of(privateMemRefs, [&](Value privateMemRef) {
+ return isSameMemref(privateMemRef, inEdge.value);
+ }))
addEdge(inEdge.id, dstId, inEdge.value);
}
}
@@ -767,7 +787,9 @@ void MemRefDependenceGraph::updateEdges(unsigned srcId, unsigned dstId,
if (inEdges.count(dstId) > 0 && !privateMemRefs.empty()) {
SmallVector<Edge, 2> oldInEdges = inEdges[dstId];
for (auto &inEdge : oldInEdges)
- if (privateMemRefs.count(inEdge.value) > 0)
+ if (llvm::any_of(privateMemRefs, [&](Value privateMemRef) {
+ return isSameMemref(privateMemRef, inEdge.value);
+ }))
removeEdge(inEdge.id, dstId, inEdge.value);
}
}
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 1ec5fbfef50c3..4c6b94c26a24a 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -19,6 +19,7 @@
#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
@@ -154,7 +155,11 @@ static void getProducerCandidates(unsigned dstId,
if (any_of(srcNode->stores, [&](Operation *op) {
auto storeOp = cast<AffineWriteOpInterface>(op);
- return consumedMemrefs.count(storeOp.getMemRef()) > 0;
+ return llvm::any_of(consumedMemrefs, [&](Value consumedMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(consumedMemref),
+ cast<MemrefValue>(storeOp.getMemRef()));
+ });
}))
srcIdCandidates.push_back(srcNode->id);
}
@@ -218,7 +223,10 @@ static void gatherEscapingMemrefs(unsigned id, const MemRefDependenceGraph &mdg,
auto *node = mdg.getNode(id);
for (Operation *storeOp : node->stores) {
auto memref = cast<AffineWriteOpInterface>(storeOp).getMemRef();
- if (escapingMemRefs.count(memref))
+ if (llvm::any_of(escapingMemRefs, [&](Value escapingMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(escapingMemref), cast<MemrefValue>(memref));
+ }))
continue;
if (isEscapingMemref(memref, &mdg.block))
escapingMemRefs.insert(memref);
@@ -852,7 +860,10 @@ struct GreedyFusion {
// 1. The source is to be removed after fusion,
// OR
// 2. The destination writes to `memref`.
- if (srcEscapingMemRefs.count(memref) > 0 &&
+ if (llvm::any_of(srcEscapingMemRefs, [&](Value escapingMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(escapingMemref), cast<MemrefValue>(memref));
+ }) &&
(removeSrcNode || consumerNode->getStoreOpCount(memref) > 0))
return false;
@@ -867,7 +878,10 @@ struct GreedyFusion {
// cannot create a private memref.
if (removeSrcNode &&
any_of(mdg->outEdges[producerId], [&](const auto &edge) {
- return edge.value == memref && edge.id != consumerId;
+ return edge.id != consumerId && isa<MemRefType>(edge.value.getType()) &&
+ memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(edge.value),
+ cast<MemrefValue>(memref));
}))
return false;
@@ -972,12 +986,20 @@ struct GreedyFusion {
// producer-consumer loads/stores.
SmallVector<Operation *, 2> dstMemrefOps;
for (Operation *op : dstNode->loads)
- if (producerConsumerMemrefs.count(
- cast<AffineReadOpInterface>(op).getMemRef()) > 0)
+ if (llvm::any_of(producerConsumerMemrefs, [&](Value producerMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(producerMemref),
+ cast<MemrefValue>(
+ cast<AffineReadOpInterface>(op).getMemRef()));
+ }))
dstMemrefOps.push_back(op);
for (Operation *op : dstNode->stores)
- if (producerConsumerMemrefs.count(
- cast<AffineWriteOpInterface>(op).getMemRef()))
+ if (llvm::any_of(producerConsumerMemrefs, [&](Value producerMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(producerMemref),
+ cast<MemrefValue>(
+ cast<AffineWriteOpInterface>(op).getMemRef()));
+ }))
dstMemrefOps.push_back(op);
if (dstMemrefOps.empty())
continue;
@@ -1050,8 +1072,12 @@ struct GreedyFusion {
// Retrieve producer stores from the src loop.
SmallVector<Operation *, 2> producerStores;
for (Operation *op : srcNode->stores)
- if (producerConsumerMemrefs.count(
- cast<AffineWriteOpInterface>(op).getMemRef()))
+ if (llvm::any_of(producerConsumerMemrefs, [&](Value producerMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(producerMemref),
+ cast<MemrefValue>(
+ cast<AffineWriteOpInterface>(op).getMemRef()));
+ }))
producerStores.push_back(op);
assert(!producerStores.empty() && "Expected producer store");
@@ -1112,7 +1138,11 @@ struct GreedyFusion {
DenseMap<Value, SmallVector<Operation *, 4>> privateMemRefToStores;
dstAffineForOp.walk([&](AffineWriteOpInterface storeOp) {
Value storeMemRef = storeOp.getMemRef();
- if (privateMemrefs.count(storeMemRef) > 0)
+ if (llvm::any_of(privateMemrefs, [&](Value privateMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(privateMemref),
+ cast<MemrefValue>(storeMemRef));
+ }))
privateMemRefToStores[storeMemRef].push_back(storeOp);
});
@@ -1386,8 +1416,8 @@ struct GreedyFusion {
// Check that all stores are to the same memref if any.
DenseSet<Value> storeMemrefs;
for (auto *storeOpInst : sibNode->stores) {
- storeMemrefs.insert(
- cast<AffineWriteOpInterface>(storeOpInst).getMemRef());
+ storeMemrefs.insert(memref::skipFullyAliasingOperations(cast<MemrefValue>(
+ cast<AffineWriteOpInterface>(storeOpInst).getMemRef())));
}
return storeMemrefs.size() <= 1;
};
@@ -1457,7 +1487,10 @@ struct GreedyFusion {
if (visitedSibNodeIds->count(sibNodeId) > 0)
return;
// Skip output edge if not a sibling using the same memref.
- if (outEdge.id == dstNode->id || outEdge.value != inEdge.value)
+ if (outEdge.id == dstNode->id ||
+ !memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(outEdge.value),
+ cast<MemrefValue>(inEdge.value)))
return;
auto *sibNode = mdg->getNode(sibNodeId);
if (!isa<AffineForOp>(sibNode->op))
diff --git a/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt b/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt
index ef6e0dbf45d3a..efeac83a098e7 100644
--- a/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt
@@ -14,6 +14,7 @@ add_mlir_dialect_library(MLIRAffineUtils
MLIRArithUtils
MLIRFuncDialect
MLIRMemRefDialect
+ MLIRMemRefUtils
MLIRTransformUtils
MLIRViewLikeInterface
)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 68296ea3368a1..4e4f8fd696ffd 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -18,6 +18,7 @@
#include "mlir/Dialect/Affine/Analysis/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
@@ -37,10 +38,14 @@ static void getLoadAndStoreMemRefAccesses(Operation *opA,
DenseMap<Value, bool> &values) {
opA->walk([&](Operation *op) {
if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
- if (values.count(loadOp.getMemRef()) == 0)
- values[loadOp.getMemRef()] = false;
+ Value memref = memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(loadOp.getMemRef()));
+ if (values.count(memref) == 0)
+ values[memref] = false;
} else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
- values[storeOp.getMemRef()] = true;
+ Value memref = memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(storeOp.getMemRef()));
+ values[memref] = true;
}
});
}
@@ -50,10 +55,16 @@ static void getLoadAndStoreMemRefAccesses(Operation *opA,
/// Returns false otherwise.
static bool isDependentLoadOrStoreOp(Operation *op,
DenseMap<Value, bool> &values) {
- if (auto loadOp = dyn_cast<AffineReadOpInterface>(op))
- return values.count(loadOp.getMemRef()) > 0 && values[loadOp.getMemRef()];
- if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op))
- return values.count(storeOp.getMemRef()) > 0;
+ if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
+ Value memref = memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(loadOp.getMemRef()));
+ return values.count(memref) > 0 && values[memref];
+ }
+ if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
+ Value memref = memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(storeOp.getMemRef()));
+ return values.count(memref) > 0;
+ }
return false;
}
@@ -200,7 +211,10 @@ static unsigned getMaxLoopDepth(ArrayRef<Operation *> srcOps,
auto loadOp = dyn_cast<AffineReadOpInterface>(dstOp);
Value memref = loadOp ? loadOp.getMemRef()
: cast<AffineWriteOpInterface>(dstOp).getMemRef();
- if (producerConsumerMemrefs.count(memref) > 0)
+ if (llvm::any_of(producerConsumerMemrefs, [&](Value producerMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(producerMemref), cast<MemrefValue>(memref));
+ }))
targetDstOps.push_back(dstOp);
}
@@ -223,6 +237,19 @@ static unsigned getMaxLoopDepth(ArrayRef<Operation *> srcOps,
auto *dstOpInst = targetDstOps[j];
MemRefAccess dstAccess(dstOpInst);
+ if (!memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(srcAccess.memref),
+ cast<MemrefValue>(dstAccess.memref)))
+ continue;
+ // Affine maps are expressed in the raw view coordinates. Do not run
+ // precise dependence analysis across different views without a
+ // translation between those coordinate systems.
+ if (srcAccess.memref != dstAccess.memref) {
+ if (srcAccess.isStore() || dstAccess.isStore())
+ return 0;
+ continue;
+ }
+
unsigned numCommonLoops =
getNumCommonSurroundingLoops(*srcOpInst, *dstOpInst);
for (unsigned d = 1; d <= numCommonLoops + 1; ++d) {
@@ -328,7 +355,10 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
// 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())
+ if (load && memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(load.getMemRef()),
+ cast<MemrefValue>(
+ fusionStrategy.getSiblingFusionMemRef())))
strategyOpsA.push_back(op);
}
break;
@@ -652,6 +682,10 @@ void mlir::affine::gatherProducerConsumerMemrefs(
// memrefs from loads in 'dstOps'.
for (Operation *op : dstOps)
if (auto loadOp = dyn_cast<AffineReadOpInterface>(op))
- if (srcStoreMemRefs.count(loadOp.getMemRef()) > 0)
+ if (llvm::any_of(srcStoreMemRefs, [&](Value storeMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(storeMemref),
+ cast<MemrefValue>(loadOp.getMemRef()));
+ }))
producerConsumerMemrefs.insert(loadOp.getMemRef());
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 980e68e6b9cfa..471e2f4a9bb64 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -921,3 +921,91 @@ func.func @cast_alias_external_call(
return
}
func.func private @escape(memref<?xf64>)
+
+// -----
+
+// Affine accesses may use the fully aliasing view while the external call uses
+// the source value. The call must remain between the two loop nests.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_reverse
+// PRODUCER-CONSUMER-MAXIMAL: memref.cast
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: call @escape_static
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @cast_alias_reverse(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ %view = memref.cast %comm : memref<32xf64> to memref<?xf64>
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %view[%i] : memref<?xf64>
+ }
+ func.call @escape_static(%comm) : (memref<32xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %view[%j] : memref<?xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape_static(memref<32xf64>)
+
+// -----
+
+// Two distinct fully aliasing views must be placed in the same dependence
+// class as their source value.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_sibling
+// PRODUCER-CONSUMER-MAXIMAL: memref.cast
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: call @escape
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @cast_alias_sibling(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ %view0 = memref.cast %comm : memref<32xf64> to memref<?xf64>
+ %view1 = memref.cast %comm : memref<32xf64> to memref<?xf64>
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %view0[%i] : memref<?xf64>
+ }
+ func.call @escape(%view1) : (memref<?xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %view0[%j] : memref<?xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape(memref<?xf64>)
+
+// -----
+
+// An external call on a distinct memref must not block an otherwise legal
+// producer-consumer fusion.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_non_alias_call
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: call @escape_other
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: return
+func.func @cast_alias_non_alias_call(
+ %in: memref<32xf64>, %out: memref<32xf64>) {
+ %comm = memref.alloc() : memref<32xf64>
+ %other = memref.alloc() : memref<32xf64>
+ %view = memref.cast %other : memref<32xf64> to memref<?xf64>
+ %cst = arith.constant 1.0 : f64
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %cst : f64
+ affine.store %b, %comm[%i] : memref<32xf64>
+ }
+ func.call @escape_other(%view) : (memref<?xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ affine.store %c, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape_other(memref<?xf64>)
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 1ea42517988c3..0784e079adf5d 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1574,5 +1574,44 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
return
}
-// Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
+// -----
+
+// Unknown operations nested in an affine loop may access their memref
+// operands. Dependence checking must handle them without assuming a load/store
+// operation class.
+// CHECK-LABEL: func @nested_unknown_call
+// CHECK: func.call @escape_nested
+// CHECK: return
+func.func @nested_unknown_call(%m: memref<8xf32>, %out: memref<8xf32>) {
+ affine.for %i = 0 to 8 {
+ func.call @escape_nested(%m) : (memref<8xf32>) -> ()
+ }
+ affine.for %i = 0 to 8 {
+ %v = affine.load %m[%i] : memref<8xf32>
+ affine.store %v, %out[%i] : memref<8xf32>
+ }
+ return
+}
+func.func private @escape_nested(memref<8xf32>)
+
+// -----
+
+// Multi-memref operations must follow the same arbitrary-operation path.
+
+// CHECK-LABEL: func @nested_memref_copy
+// CHECK: memref.copy
+// CHECK: return
+func.func @nested_memref_copy(
+ %src: memref<8xf32>, %dst: memref<8xf32>, %out: memref<8xf32>) {
+ affine.for %i = 0 to 8 {
+ memref.copy %src, %dst : memref<8xf32> to memref<8xf32>
+ }
+ affine.for %i = 0 to 8 {
+ %v = affine.load %dst[%i] : memref<8xf32>
+ affine.store %v, %out[%i] : memref<8xf32>
+ }
+ return
+}
+
+// Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
>From b2345892ea76e5761d091e447d92a17a62a49e49 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 15:45:44 +0800
Subject: [PATCH 4/8] Format Affine alias handling changes
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 30 ++++++++--------
.../Dialect/Affine/Transforms/LoopFusion.cpp | 36 ++++++++++---------
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 8 ++---
3 files changed, 38 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index bbc2fd0888d7c..ed9d2d3a49c80 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -178,8 +178,7 @@ unsigned Node::hasFree(Value memref) const {
void Node::getStoreOpsForMemref(Value memref,
SmallVectorImpl<Operation *> *storeOps) const {
for (Operation *storeOp : stores) {
- if (isSameMemref(memref,
- cast<AffineWriteOpInterface>(storeOp).getMemRef()))
+ if (isSameMemref(memref, cast<AffineWriteOpInterface>(storeOp).getMemRef()))
storeOps->push_back(storeOp);
}
}
@@ -188,8 +187,7 @@ void Node::getStoreOpsForMemref(Value memref,
void Node::getLoadOpsForMemref(Value memref,
SmallVectorImpl<Operation *> *loadOps) const {
for (Operation *loadOp : loads) {
- if (isSameMemref(memref,
- cast<AffineReadOpInterface>(loadOp).getMemRef()))
+ if (isSameMemref(memref, cast<AffineReadOpInterface>(loadOp).getMemRef()))
loadOps->push_back(loadOp);
}
}
@@ -200,12 +198,12 @@ void Node::getLoadAndStoreMemrefSet(
DenseSet<Value> *loadAndStoreMemrefSet) const {
llvm::SmallDenseSet<Value, 2> loadMemrefs;
for (Operation *loadOp : loads) {
- loadMemrefs.insert(canonicalizeMemref(
- cast<AffineReadOpInterface>(loadOp).getMemRef()));
+ loadMemrefs.insert(
+ canonicalizeMemref(cast<AffineReadOpInterface>(loadOp).getMemRef()));
}
for (Operation *storeOp : stores) {
- auto memref = canonicalizeMemref(
- cast<AffineWriteOpInterface>(storeOp).getMemRef());
+ auto memref =
+ canonicalizeMemref(cast<AffineWriteOpInterface>(storeOp).getMemRef());
if (loadMemrefs.count(memref) > 0)
loadAndStoreMemrefSet->insert(memref);
}
@@ -225,14 +223,14 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
Node &node = nodes.insert({newNodeId, Node(newNodeId, nodeOp)}).first->second;
for (Operation *op : collector.loadOpInsts) {
node.loads.push_back(op);
- auto memref = canonicalizeMemref(
- cast<AffineReadOpInterface>(op).getMemRef());
+ auto memref =
+ canonicalizeMemref(cast<AffineReadOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
}
for (Operation *op : collector.storeOpInsts) {
node.stores.push_back(op);
- auto memref = canonicalizeMemref(
- cast<AffineWriteOpInterface>(op).getMemRef());
+ auto memref =
+ canonicalizeMemref(cast<AffineWriteOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
}
for (Operation *op : collector.memrefLoads) {
@@ -382,16 +380,16 @@ bool MemRefDependenceGraph::init(bool fullAffineDependences) {
// Create graph node for top-level load op.
Node node(nextNodeId++, &op);
node.loads.push_back(&op);
- auto memref = canonicalizeMemref(
- cast<AffineReadOpInterface>(op).getMemRef());
+ auto memref =
+ canonicalizeMemref(cast<AffineReadOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
nodes.insert({node.id, node});
} else if (isa<AffineWriteOpInterface>(op)) {
// Create graph node for top-level store op.
Node node(nextNodeId++, &op);
node.stores.push_back(&op);
- auto memref = canonicalizeMemref(
- cast<AffineWriteOpInterface>(op).getMemRef());
+ auto memref =
+ canonicalizeMemref(cast<AffineWriteOpInterface>(op).getMemRef());
memrefAccesses[memref].insert(node.id);
nodes.insert({node.id, node});
} else if (op.getNumResults() > 0 && !op.use_empty()) {
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 4c6b94c26a24a..c957c52504a6a 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -860,10 +860,12 @@ struct GreedyFusion {
// 1. The source is to be removed after fusion,
// OR
// 2. The destination writes to `memref`.
- if (llvm::any_of(srcEscapingMemRefs, [&](Value escapingMemref) {
- return memref::isSameViewOrTrivialAlias(
- cast<MemrefValue>(escapingMemref), cast<MemrefValue>(memref));
- }) &&
+ if (llvm::any_of(srcEscapingMemRefs,
+ [&](Value escapingMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(escapingMemref),
+ cast<MemrefValue>(memref));
+ }) &&
(removeSrcNode || consumerNode->getStoreOpCount(memref) > 0))
return false;
@@ -878,10 +880,10 @@ struct GreedyFusion {
// cannot create a private memref.
if (removeSrcNode &&
any_of(mdg->outEdges[producerId], [&](const auto &edge) {
- return edge.id != consumerId && isa<MemRefType>(edge.value.getType()) &&
- memref::isSameViewOrTrivialAlias(
- cast<MemrefValue>(edge.value),
- cast<MemrefValue>(memref));
+ return edge.id != consumerId &&
+ isa<MemRefType>(edge.value.getType()) &&
+ memref::isSameViewOrTrivialAlias(cast<MemrefValue>(edge.value),
+ cast<MemrefValue>(memref));
}))
return false;
@@ -1072,12 +1074,13 @@ struct GreedyFusion {
// Retrieve producer stores from the src loop.
SmallVector<Operation *, 2> producerStores;
for (Operation *op : srcNode->stores)
- if (llvm::any_of(producerConsumerMemrefs, [&](Value producerMemref) {
- return memref::isSameViewOrTrivialAlias(
- cast<MemrefValue>(producerMemref),
- cast<MemrefValue>(
- cast<AffineWriteOpInterface>(op).getMemRef()));
- }))
+ if (llvm::any_of(
+ producerConsumerMemrefs, [&](Value producerMemref) {
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(producerMemref),
+ cast<MemrefValue>(
+ cast<AffineWriteOpInterface>(op).getMemRef()));
+ }))
producerStores.push_back(op);
assert(!producerStores.empty() && "Expected producer store");
@@ -1416,8 +1419,9 @@ struct GreedyFusion {
// Check that all stores are to the same memref if any.
DenseSet<Value> storeMemrefs;
for (auto *storeOpInst : sibNode->stores) {
- storeMemrefs.insert(memref::skipFullyAliasingOperations(cast<MemrefValue>(
- cast<AffineWriteOpInterface>(storeOpInst).getMemRef())));
+ storeMemrefs.insert(
+ memref::skipFullyAliasingOperations(cast<MemrefValue>(
+ cast<AffineWriteOpInterface>(storeOpInst).getMemRef())));
}
return storeMemrefs.size() <= 1;
};
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 4e4f8fd696ffd..c34bfe40da89e 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -355,10 +355,10 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
// to 'memref' in 'srcForOp' to compute the slice union.
for (Operation *op : opsA) {
auto load = dyn_cast<AffineReadOpInterface>(op);
- if (load && memref::isSameViewOrTrivialAlias(
- cast<MemrefValue>(load.getMemRef()),
- cast<MemrefValue>(
- fusionStrategy.getSiblingFusionMemRef())))
+ if (load &&
+ memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(load.getMemRef()),
+ cast<MemrefValue>(fusionStrategy.getSiblingFusionMemRef())))
strategyOpsA.push_back(op);
}
break;
>From 03d0fdee46a86513d4fcb395fb18c854486670c6 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 15:57:47 +0800
Subject: [PATCH 5/8] [MLIR][Affine] Preserve memref SSA edge identity
---
.../mlir/Dialect/Affine/Analysis/Utils.h | 14 +++++++-----
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 13 ++++++-----
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 22 +++++++++++++++++++
3 files changed, 37 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index 6a03fabd274fd..8c95630edcbf1 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
@@ -128,12 +128,14 @@ struct MemRefDependenceGraph {
// 'Node.outEdges[i].id' is the identifier of the dest node of the edge.
unsigned id;
// The SSA value on which this edge represents a dependence.
- // If the value is a memref, then it is the canonical representative of
- // the trivial alias class on which the dependence is based. If the value
- // is a non-memref value, then the dependence is between a graph node which
- // defines an SSA value and another graph node which uses the SSA value
- // (e.g. a constant or load operation defining a value which is used inside
- // a loop nest).
+ // If the value is a memref and this is a memory dependence, then it is the
+ // canonical representative of the trivial alias class on which the
+ // dependence is based. Memref SSA dependences retain the defining value so
+ // that the defining operation remains observable to graph clients. If the
+ // value is a non-memref value, then the dependence is between a graph node
+ // which defines an SSA value and another graph node which uses the SSA
+ // value (e.g. a constant or load operation defining a value which is used
+ // inside a loop nest).
Value value;
};
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index ed9d2d3a49c80..d061a7d0da2f0 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -550,7 +550,6 @@ bool MemRefDependenceGraph::hasEdge(unsigned srcId, unsigned dstId,
if (!outEdges.contains(srcId) || !inEdges.contains(dstId)) {
return false;
}
- value = canonicalizeMemref(value);
bool hasOutEdge = llvm::any_of(outEdges.lookup(srcId), [=](const Edge &edge) {
return edge.id == dstId && (!value || edge.value == value);
});
@@ -563,12 +562,14 @@ bool MemRefDependenceGraph::hasEdge(unsigned srcId, unsigned dstId,
// Adds an edge from node 'srcId' to node 'dstId' for 'value'.
void MemRefDependenceGraph::addEdge(unsigned srcId, unsigned dstId,
Value value) {
- value = canonicalizeMemref(value);
+ // Keep memref SSA edges in their raw form so their defining operation stays
+ // available to graph clients. Memory-dependence callers provide the
+ // canonical representative; count both kinds by canonical identity.
if (!hasEdge(srcId, dstId, value)) {
outEdges[srcId].push_back({dstId, value});
inEdges[dstId].push_back({srcId, value});
if (isa<MemRefType>(value.getType()))
- memrefEdgeCount[value]++;
+ memrefEdgeCount[canonicalizeMemref(value)]++;
}
}
@@ -577,10 +578,10 @@ void MemRefDependenceGraph::removeEdge(unsigned srcId, unsigned dstId,
Value value) {
assert(inEdges.count(dstId) > 0);
assert(outEdges.count(srcId) > 0);
- value = canonicalizeMemref(value);
if (isa<MemRefType>(value.getType())) {
- assert(memrefEdgeCount.count(value) > 0);
- memrefEdgeCount[value]--;
+ Value canonicalValue = canonicalizeMemref(value);
+ assert(memrefEdgeCount.count(canonicalValue) > 0);
+ memrefEdgeCount[canonicalValue]--;
}
// Remove 'srcId' from 'inEdges[dstId]'.
for (auto *it = inEdges[dstId].begin(); it != inEdges[dstId].end(); ++it) {
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index c34bfe40da89e..1bfc3af59fd03 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -364,6 +364,28 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
break;
}
+ // Affine access maps are expressed in the raw view coordinates. Reject a
+ // fusion that would pair different views from one trivial alias class until
+ // a coordinate translation is available.
+ auto getMemref = [](Operation *op) -> Value {
+ if (auto load = dyn_cast<AffineReadOpInterface>(op))
+ return load.getMemRef();
+ return cast<AffineWriteOpInterface>(op).getMemRef();
+ };
+ if (llvm::any_of(strategyOpsA, [&](Operation *srcOp) {
+ return llvm::any_of(opsB, [&](Operation *dstOp) {
+ Value srcMemref = getMemref(srcOp);
+ Value dstMemref = getMemref(dstOp);
+ return srcMemref != dstMemref &&
+ memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(srcMemref),
+ cast<MemrefValue>(dstMemref));
+ });
+ })) {
+ LDBG() << "Fusion across different trivial alias views is unsupported";
+ return FusionResult::FailFusionDependence;
+ }
+
// Compute union of computation slices computed between all pairs of ops
// from 'forOpA' and 'forOpB'.
SliceComputationResult sliceComputationResult = affine::computeSliceUnion(
>From 43d392a12e9322d3f214146715cc7d5ae7bcccb1 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 16:00:42 +0800
Subject: [PATCH 6/8] Format alias view guard
---
mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp | 5 ++---
1 file changed, 2 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 1bfc3af59fd03..e51b9ff5dea27 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -377,9 +377,8 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
Value srcMemref = getMemref(srcOp);
Value dstMemref = getMemref(dstOp);
return srcMemref != dstMemref &&
- memref::isSameViewOrTrivialAlias(
- cast<MemrefValue>(srcMemref),
- cast<MemrefValue>(dstMemref));
+ memref::isSameViewOrTrivialAlias(cast<MemrefValue>(srcMemref),
+ cast<MemrefValue>(dstMemref));
});
})) {
LDBG() << "Fusion across different trivial alias views is unsupported";
>From 9d59f1c77ca4e79e69d035c0f15ae49d7fc15775 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 18:46:57 +0800
Subject: [PATCH 7/8] [MLIR][Affine] Handle unranked memref aliases
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 22 +--
.../Dialect/Affine/Transforms/LoopFusion.cpp | 2 +-
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 125 ++++++++++++++++++
mlir/test/Dialect/Affine/loop-fusion.mlir | 10 +-
4 files changed, 145 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index d061a7d0da2f0..97adf11743321 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -40,7 +40,7 @@ using llvm::SmallDenseMap;
using Node = MemRefDependenceGraph::Node;
static Value canonicalizeMemref(Value value) {
- if (!value || !isa<MemRefType>(value.getType()))
+ if (!value || !isa<BaseMemRefType>(value.getType()))
return value;
return memref::skipFullyAliasingOperations(cast<MemrefValue>(value));
}
@@ -64,7 +64,7 @@ static void getMayAffectedValues(Operation *op,
return;
// Memref operands have to be considered as being affected.
for (Value operand : op->getOperands()) {
- if (isa<MemRefType>(operand.getType()))
+ if (isa<BaseMemRefType>(operand.getType()))
values.push_back(canonicalizeMemref(operand));
}
return;
@@ -74,7 +74,7 @@ static void getMayAffectedValues(Operation *op,
for (auto &effect : effects) {
Value effectVal = effect.getValue();
if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
- isa<MemRefType>(effectVal.getType()))
+ isa<BaseMemRefType>(effectVal.getType()))
values.push_back(canonicalizeMemref(effectVal));
};
}
@@ -108,7 +108,7 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
return;
// Check operands. E.g., ops like the `call` op are handled here.
if (llvm::any_of(op->getOperands(), [](Value value) {
- return isa<MemRefType>(value.getType());
+ return isa<BaseMemRefType>(value.getType());
})) {
// Conservatively, assume all memref operands are read and written.
memrefLoads.push_back(op);
@@ -237,7 +237,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
SmallVector<Value> affectedValues;
getMayAffectedValues<MemoryEffects::Read>(op, affectedValues);
if (llvm::any_of(((ValueRange)affectedValues).getTypes(),
- [](Type type) { return !isa<MemRefType>(type); }))
+ [](Type type) { return !isa<BaseMemRefType>(type); }))
// We do not know the interaction here.
return nullptr;
for (Value memref : affectedValues)
@@ -248,7 +248,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
SmallVector<Value> affectedValues;
getMayAffectedValues<MemoryEffects::Write>(op, affectedValues);
if (llvm::any_of((ValueRange(affectedValues)).getTypes(),
- [](Type type) { return !isa<MemRefType>(type); }))
+ [](Type type) { return !isa<BaseMemRefType>(type); }))
return nullptr;
for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
@@ -258,7 +258,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
SmallVector<Value> affectedValues;
getMayAffectedValues<MemoryEffects::Free>(op, affectedValues);
if (llvm::any_of((ValueRange(affectedValues)).getTypes(),
- [](Type type) { return !isa<MemRefType>(type); }))
+ [](Type type) { return !isa<BaseMemRefType>(type); }))
return nullptr;
for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
@@ -568,7 +568,7 @@ void MemRefDependenceGraph::addEdge(unsigned srcId, unsigned dstId,
if (!hasEdge(srcId, dstId, value)) {
outEdges[srcId].push_back({dstId, value});
inEdges[dstId].push_back({srcId, value});
- if (isa<MemRefType>(value.getType()))
+ if (isa<BaseMemRefType>(value.getType()))
memrefEdgeCount[canonicalizeMemref(value)]++;
}
}
@@ -578,7 +578,7 @@ void MemRefDependenceGraph::removeEdge(unsigned srcId, unsigned dstId,
Value value) {
assert(inEdges.count(dstId) > 0);
assert(outEdges.count(srcId) > 0);
- if (isa<MemRefType>(value.getType())) {
+ if (isa<BaseMemRefType>(value.getType())) {
Value canonicalValue = canonicalizeMemref(value);
assert(memrefEdgeCount.count(canonicalValue) > 0);
memrefEdgeCount[canonicalValue]--;
@@ -669,7 +669,7 @@ void MemRefDependenceGraph::gatherDefiningNodes(
// By definition of edge, if the edge value is a non-memref value,
// then the dependence is between a graph node which defines an SSA value
// and another graph node which uses the SSA value.
- if (!isa<MemRefType>(edge.value.getType()))
+ if (!isa<BaseMemRefType>(edge.value.getType()))
definingNodes.insert(edge.id);
}
@@ -861,7 +861,7 @@ void MemRefDependenceGraph::forEachMemRefEdge(
ArrayRef<Edge> edges, const std::function<void(Edge)> &callback) {
for (const auto &edge : edges) {
// Skip if 'edge' is not a memref dependence edge.
- if (!isa<MemRefType>(edge.value.getType()))
+ if (!isa<BaseMemRefType>(edge.value.getType()))
continue;
assert(nodes.count(edge.id) > 0);
// Visit current input edge 'edge'.
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index c957c52504a6a..a22a6639fd314 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -881,7 +881,7 @@ struct GreedyFusion {
if (removeSrcNode &&
any_of(mdg->outEdges[producerId], [&](const auto &edge) {
return edge.id != consumerId &&
- isa<MemRefType>(edge.value.getType()) &&
+ isa<BaseMemRefType>(edge.value.getType()) &&
memref::isSameViewOrTrivialAlias(cast<MemrefValue>(edge.value),
cast<MemrefValue>(memref));
}))
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 471e2f4a9bb64..600c1d1d484db 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -924,6 +924,66 @@ func.func private @escape(memref<?xf64>)
// -----
+// A ranked source cast to an unranked memref must retain the dependence on the
+// source used by the affine producer and consumer.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_ranked_to_unranked
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.store
+// PRODUCER-CONSUMER-MAXIMAL: memref.cast
+// PRODUCER-CONSUMER-MAXIMAL: call @escape_unranked
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+func.func @cast_alias_ranked_to_unranked(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %comm[%i] : memref<32xf64>
+ }
+ %view = memref.cast %comm : memref<32xf64> to memref<*xf64>
+ func.call @escape_unranked(%view) : (memref<*xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape_unranked(memref<*xf64>)
+
+// -----
+
+// An unranked source cast to a ranked memref must retain the dependence on the
+// source passed to the opaque call.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_unranked_to_ranked
+// PRODUCER-CONSUMER-MAXIMAL: memref.cast
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.store
+// PRODUCER-CONSUMER-MAXIMAL: call @escape_ranked
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+func.func @cast_alias_unranked_to_ranked(
+ %in: memref<32xf64>, %comm: memref<*xf64>, %out: memref<32xf64>) {
+ %view = memref.cast %comm : memref<*xf64> to memref<32xf64>
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %view[%i] : memref<32xf64>
+ }
+ func.call @escape_ranked(%comm) : (memref<*xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %view[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape_ranked(memref<*xf64>)
+
+// -----
+
// Affine accesses may use the fully aliasing view while the external call uses
// the source value. The call must remain between the two loop nests.
@@ -1009,3 +1069,68 @@ func.func @cast_alias_non_alias_call(
return
}
func.func private @escape_other(memref<?xf64>)
+
+// -----
+
+// A zero-offset, unit-stride subview is a fully aliasing view and must retain
+// the same dependence as its source memref.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @subview_alias_external_call
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.store
+// PRODUCER-CONSUMER-MAXIMAL: memref.subview
+// PRODUCER-CONSUMER-MAXIMAL: call @escape_subview
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+func.func @subview_alias_external_call(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %comm[%i] : memref<32xf64>
+ }
+ %view = memref.subview %comm[0] [32] [1]
+ : memref<32xf64> to memref<32xf64, strided<[1], offset: 0>>
+ func.call @escape_subview(%view)
+ : (memref<32xf64, strided<[1], offset: 0>>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape_subview(memref<32xf64, strided<[1], offset: 0>>)
+
+// -----
+
+// A non-fully-aliasing subview of an unrelated memref must not block an
+// otherwise legal producer-consumer fusion.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @subview_non_alias_call
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: call @escape_subview_non_alias
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+// PRODUCER-CONSUMER-MAXIMAL: return
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+func.func @subview_non_alias_call(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ %other = memref.alloc() : memref<32xf64>
+ %view = memref.subview %other[1] [16] [1]
+ : memref<32xf64> to memref<16xf64, strided<[1], offset: 1>>
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %comm[%i] : memref<32xf64>
+ }
+ func.call @escape_subview_non_alias(%view)
+ : (memref<16xf64, strided<[1], offset: 1>>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ affine.store %c, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape_subview_non_alias(
+ memref<16xf64, strided<[1], offset: 1>>)
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 0784e079adf5d..96a8f0c6218ee 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1581,7 +1581,10 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
// operation class.
// CHECK-LABEL: func @nested_unknown_call
-// CHECK: func.call @escape_nested
+// CHECK: affine.for
+// CHECK: func.call @escape_nested
+// CHECK: affine.for
+// CHECK: affine.load
// CHECK: return
func.func @nested_unknown_call(%m: memref<8xf32>, %out: memref<8xf32>) {
affine.for %i = 0 to 8 {
@@ -1600,7 +1603,10 @@ func.func private @escape_nested(memref<8xf32>)
// Multi-memref operations must follow the same arbitrary-operation path.
// CHECK-LABEL: func @nested_memref_copy
-// CHECK: memref.copy
+// CHECK: affine.for
+// CHECK: memref.copy
+// CHECK: affine.for
+// CHECK: affine.load
// CHECK: return
func.func @nested_memref_copy(
%src: memref<8xf32>, %dst: memref<8xf32>, %out: memref<8xf32>) {
>From 8c35ea405bed0740549e20cb364c8c67c03a36f0 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 21:41:51 +0800
Subject: [PATCH 8/8] [MLIR][Affine] Preserve addressable memory effects
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 47 +++++++++-----
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 70 +++++++++++++++++++++
2 files changed, 100 insertions(+), 17 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 97adf11743321..d0a4710e359e2 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -54,38 +54,58 @@ static bool isSameMemref(Value lhs, Value rhs) {
/// (e.g. a call to an external function without a memory-effect interface) is
/// conservatively assumed to affect all its memref operands. Fully aliasing
/// views are canonicalized so the MDG uses one key for the view and its source.
+/// Returns false if an addressable effect cannot be represented by a memref
+/// value, in which case the MDG must not be used for fusion.
template <typename... EffectTys>
-static void getMayAffectedValues(Operation *op,
+static bool getMayAffectedValues(Operation *op,
SmallVectorImpl<Value> &values) {
auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
if (!memOp) {
if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
// No effects.
- return;
+ return true;
// Memref operands have to be considered as being affected.
for (Value operand : op->getOperands()) {
if (isa<BaseMemRefType>(operand.getType()))
values.push_back(canonicalizeMemref(operand));
}
- return;
+ return true;
}
SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
memOp.getEffects(effects);
for (auto &effect : effects) {
+ if (!isa<EffectTys...>(effect.getEffect()))
+ continue;
Value effectVal = effect.getValue();
- if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
- isa<BaseMemRefType>(effectVal.getType()))
+ if (!effectVal) {
+ // A value-less or symbol-associated effect on an addressable resource
+ // cannot be represented by a per-memref graph edge. Refuse to fuse the
+ // block rather than silently dropping the effect.
+ if (effect.getResource()->isAddressable())
+ return false;
+ continue;
+ }
+ if (isa<BaseMemRefType>(effectVal.getType())) {
values.push_back(canonicalizeMemref(effectVal));
+ continue;
+ }
+ // An addressable effect on a non-memref value (for example, a pointer) is
+ // equally unrepresentable by the memref dependence graph.
+ if (effect.getResource()->isAddressable())
+ return false;
};
+ return true;
}
/// Returns true if `op` may have a memory effect of type `EffectTys` on
/// `memref`, i.e., whether `memref` is among the values returned by
-/// `getMayAffectedValues` for `op`.
+/// `getMayAffectedValues` for `op`. An unrepresentable addressable effect is
+/// conservatively treated as affecting every memref.
template <typename... EffectTys>
static bool mayHaveEffect(Operation *op, Value memref) {
SmallVector<Value> values;
- getMayAffectedValues<EffectTys...>(op, values);
+ if (!getMayAffectedValues<EffectTys...>(op, values))
+ return true;
return llvm::is_contained(values, canonicalizeMemref(memref));
}
@@ -235,10 +255,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefLoads) {
SmallVector<Value> affectedValues;
- getMayAffectedValues<MemoryEffects::Read>(op, affectedValues);
- if (llvm::any_of(((ValueRange)affectedValues).getTypes(),
- [](Type type) { return !isa<BaseMemRefType>(type); }))
- // We do not know the interaction here.
+ if (!getMayAffectedValues<MemoryEffects::Read>(op, affectedValues))
return nullptr;
for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
@@ -246,9 +263,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefStores) {
SmallVector<Value> affectedValues;
- getMayAffectedValues<MemoryEffects::Write>(op, affectedValues);
- if (llvm::any_of((ValueRange(affectedValues)).getTypes(),
- [](Type type) { return !isa<BaseMemRefType>(type); }))
+ if (!getMayAffectedValues<MemoryEffects::Write>(op, affectedValues))
return nullptr;
for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
@@ -256,9 +271,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefFrees) {
SmallVector<Value> affectedValues;
- getMayAffectedValues<MemoryEffects::Free>(op, affectedValues);
- if (llvm::any_of((ValueRange(affectedValues)).getTypes(),
- [](Type type) { return !isa<BaseMemRefType>(type); }))
+ if (!getMayAffectedValues<MemoryEffects::Free>(op, affectedValues))
return nullptr;
for (Value memref : affectedValues)
memrefAccesses[memref].insert(node.id);
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 600c1d1d484db..469f9088e5017 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1134,3 +1134,73 @@ func.func @subview_non_alias_call(
}
func.func private @escape_subview_non_alias(
memref<16xf64, strided<[1], offset: 1>>)
+
+// -----
+
+// An addressable effect without an SSA memory value cannot be represented by
+// the per-memref dependence graph. Fusion must be skipped for the block.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @value_less_addressable_effect
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @value_less_addressable_effect(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ affine.store %a, %comm[%i] : memref<32xf64>
+ }
+ %effect = "test.side_effect_op"() {effects = [{effect = "write"}]} : () -> i32
+ affine.for %j = 0 to 16 {
+ %a = affine.load %comm[%j] : memref<32xf64>
+ affine.store %a, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
+// A symbol-associated addressable effect is also not representable by a
+// per-memref dependence graph. Fusion must be skipped for the block.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @symbol_addressable_effect
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @symbol_addressable_effect(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ affine.store %a, %comm[%i] : memref<32xf64>
+ }
+ "test.side_effect_op"() {
+ effects = [{effect = "write", on_reference = @effect_target}]
+ } : () -> i32
+ affine.for %j = 0 to 16 {
+ %a = affine.load %comm[%j] : memref<32xf64>
+ affine.store %a, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @effect_target()
+
+// A value-less effect on a non-addressable resource is disjoint from memref
+// accesses, so the otherwise legal producer-consumer fusion remains enabled.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @value_less_nonaddressable_effect
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+func.func @value_less_nonaddressable_effect(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ affine.store %a, %comm[%i] : memref<32xf64>
+ }
+ %effect = "test.side_effect_op"() {
+ effects = [{effect = "write", test_nonaddressable_resource}]
+ } : () -> i32
+ affine.for %j = 0 to 16 {
+ %a = affine.load %comm[%j] : memref<32xf64>
+ affine.store %a, %out[%j] : memref<32xf64>
+ }
+ return
+}
More information about the Mlir-commits
mailing list