[Mlir-commits] [mlir] [MLIR][Affine] Handle aliases and effects in fusion dependencies (PR #213413)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 5 00:06:22 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 01/23] [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 02/23] 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 03/23] [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 04/23] 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 05/23] [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 06/23] 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 07/23] [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 08/23] [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
+}
>From 19f57849890548c75474c79b174bb07bac850ea4 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 10:03:41 +0800
Subject: [PATCH 09/23] [MLIR][Affine] Preserve unrepresentable memory effects
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 29 +++----
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 92 +++++++++++++++++++--
2 files changed, 96 insertions(+), 25 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index d0a4710e359e2..c4dcf258ba928 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -54,8 +54,9 @@ 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.
+/// Returns false if a selected effect cannot be represented by a memref value.
+/// The MDG has no resource-level edges, so dropping such an effect would make
+/// fusion unsound even when its resource is non-addressable.
template <typename... EffectTys>
static bool getMayAffectedValues(Operation *op,
SmallVectorImpl<Value> &values) {
@@ -77,29 +78,16 @@ static bool getMayAffectedValues(Operation *op,
if (!isa<EffectTys...>(effect.getEffect()))
continue;
Value effectVal = effect.getValue();
- 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())
+ if (!effectVal || !isa<BaseMemRefType>(effectVal.getType()))
return false;
+ values.push_back(canonicalizeMemref(effectVal));
};
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`. An unrepresentable addressable effect is
+/// `getMayAffectedValues` for `op`. An unrepresentable effect is
/// conservatively treated as affecting every memref.
template <typename... EffectTys>
static bool mayHaveEffect(Operation *op, Value memref) {
@@ -135,7 +123,10 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
memrefStores.push_back(op);
}
} else {
- // Non-affine loads and stores.
+ // Non-affine loads, stores, and frees. Allocation effects are
+ // intentionally omitted: they do not access existing memory, and
+ // allocation results are handled by existing SSA and local-allocation
+ // analysis instead of the memref access graph.
if (hasEffect<MemoryEffects::Read>(op))
memrefLoads.push_back(op);
if (hasEffect<MemoryEffects::Write>(op))
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 469f9088e5017..01967ad414335 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1182,21 +1182,101 @@ func.func @symbol_addressable_effect(
}
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.
+// An addressable effect on a non-memref SSA value is not representable by the
+// per-memref dependence graph. Fusion must be skipped for the block.
-// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @value_less_nonaddressable_effect
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @value_addressable_nonmemref_effect
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
// PRODUCER-CONSUMER-MAXIMAL: affine.for
-// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
-func.func @value_less_nonaddressable_effect(
+func.func @value_addressable_nonmemref_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}]
+ effects = [{effect = "write", on_result}]
+ } : () -> i32
+ affine.for %j = 0 to 16 {
+ %a = affine.load %comm[%j] : memref<32xf64>
+ affine.store %a, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
+// A value-less effect on a non-addressable resource cannot be represented by
+// the per-memref dependence graph. Effects on the same resource must still be
+// ordered, so fusion must be skipped for the block.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @value_less_nonaddressable_write_effect
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @value_less_nonaddressable_write_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", test_nonaddressable_resource}]
+ } : () -> i32
+ }
+ "test.side_effect_op"() {
+ effects = [{effect = "read", 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
+}
+
+// Value-less reads on a non-addressable resource are also unrepresentable.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @value_less_nonaddressable_read_effect
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+func.func @value_less_nonaddressable_read_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 = "read", test_nonaddressable_resource}]
+ } : () -> i32
+ affine.for %j = 0 to 16 {
+ %a = affine.load %comm[%j] : memref<32xf64>
+ affine.store %a, %out[%j] : memref<32xf64>
+ "test.side_effect_op"() {
+ effects = [{effect = "write", test_nonaddressable_resource}]
+ } : () -> i32
+ }
+ return
+}
+
+// Free effects follow the same conservative rule.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @value_less_nonaddressable_free_effect
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @value_less_nonaddressable_free_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 = "free", test_nonaddressable_resource}]
+ } : () -> i32
+ }
+ "test.side_effect_op"() {
+ effects = [{effect = "read", test_nonaddressable_resource}]
} : () -> i32
affine.for %j = 0 to 16 {
%a = affine.load %comm[%j] : memref<32xf64>
>From 6f2aaefb86d23be48a3ddeacb0f6b3f2f2d45b98 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 17:31:30 +0800
Subject: [PATCH 10/23] [MLIR][Affine] Fail closed on incomplete fusion effects
---
.../mlir/Dialect/Affine/Analysis/Utils.h | 7 +-
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 57 ++++----
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 136 ++++++++++++------
mlir/test/Dialect/Affine/loop-fusion.mlir | 6 +-
4 files changed, 128 insertions(+), 78 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index 8c95630edcbf1..f78fe772d77e9 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
@@ -129,7 +129,7 @@ struct MemRefDependenceGraph {
unsigned id;
// The SSA value on which this edge represents a dependence.
// 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
+ // canonical representative of the view-like storage 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
@@ -160,8 +160,9 @@ struct MemRefDependenceGraph {
// side-effect-free operations with zero results and no regions. Assigns each
// node in the graph a node id based on the order in block. Fails if certain
// kinds of operations, for which `Node` creation isn't supported, are
- // encountered (unknown region holding ops). If `fullAffineDependences` is
- // set, affine memory dependence analysis is performed before concluding that
+ // encountered (unknown effects or region holding ops). If
+ // `fullAffineDependences` is set, affine memory dependence analysis is
+ // performed before concluding that
// conflicting affine memory accesses lead to a dependence check; otherwise, a
// pair of conflicting affine memory accesses (where one of them is a store
// and they are to the same memref) always leads to an edge (conservatively).
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index c4dcf258ba928..3a731193bd887 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -42,7 +42,10 @@ using Node = MemRefDependenceGraph::Node;
static Value canonicalizeMemref(Value value) {
if (!value || !isa<BaseMemRefType>(value.getType()))
return value;
- return memref::skipFullyAliasingOperations(cast<MemrefValue>(value));
+ // Use the storage source for graph identity. Keep the original view in
+ // MemRefAccess so affine maps remain expressed in their original
+ // coordinate systems.
+ return memref::skipViewLikeOps(cast<MemrefValue>(value));
}
static bool isSameMemref(Value lhs, Value rhs) {
@@ -50,28 +53,19 @@ static bool isSameMemref(Value lhs, Value 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.
+/// on, not considering recursive effects. View-like values are canonicalized
+/// to their storage source so the MDG uses one key for a view chain while raw
+/// views remain available for affine-coordinate analysis. Unknown operations
+/// cannot be represented by the memref-keyed graph, so return false for them.
/// Returns false if a selected effect cannot be represented by a memref value.
-/// The MDG has no resource-level edges, so dropping such an effect would make
-/// fusion unsound even when its resource is non-addressable.
+/// The MDG has no resource-level or all-memory edges, so dropping such an
+/// effect would make fusion unsound.
template <typename... EffectTys>
static bool getMayAffectedValues(Operation *op,
SmallVectorImpl<Value> &values) {
auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
- if (!memOp) {
- if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
- // No effects.
- 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 true;
- }
+ if (!memOp)
+ return !hasUnknownEffects(op);
SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
memOp.getEffects(effects);
for (auto &effect : effects) {
@@ -111,17 +105,12 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
} else {
auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
if (!memInterface) {
- if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
- // This op itself is memory-effect free.
+ if (!hasUnknownEffects(op))
return;
- // Check operands. E.g., ops like the `call` op are handled here.
- if (llvm::any_of(op->getOperands(), [](Value value) {
- return isa<BaseMemRefType>(value.getType());
- })) {
- // Conservatively, assume all memref operands are read and written.
- memrefLoads.push_back(op);
- memrefStores.push_back(op);
- }
+ // Unknown effects may reach memory through globals or other state not
+ // represented by SSA memref operands. Keep the graph fail-closed.
+ memrefLoads.push_back(op);
+ memrefStores.push_back(op);
} else {
// Non-affine loads, stores, and frees. Allocation effects are
// intentionally omitted: they do not access existing memory, and
@@ -272,9 +261,9 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
return &node;
}
-/// 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.
+/// Returns true if `op` may access `memref`, including through a view-like
+/// operation. 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);
@@ -295,6 +284,10 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
assert(srcNode.op->getBlock() == dstNode.op->getBlock());
if (!isa<AffineForOp>(srcNode.op) || !isa<AffineForOp>(dstNode.op))
return true;
+ // Deallocation invalidates the whole storage object. Affine access
+ // relations cannot prove a free harmless by comparing indexed accesses.
+ if (srcNode.hasFree(memref) || dstNode.hasFree(memref))
+ return true;
// Conservatively handle dependences involving non-affine load/stores. Return
// true if there exists a conflicting read/write access involving such.
@@ -369,7 +362,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. Fully aliasing views use their canonical source value here.
+ // the memref. View-like values use their canonical storage source here.
DenseMap<Value, SetVector<unsigned>> memrefAccesses;
// Create graph nodes.
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 01967ad414335..fc82b34f870bf 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1041,37 +1041,6 @@ 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>)
-
-// -----
-
// A zero-offset, unit-stride subview is a fully aliasing view and must retain
// the same dependence as its source memref.
@@ -1104,36 +1073,123 @@ 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.
+// A non-zero-offset subview of the producer's storage must retain a
+// dependence even though its affine coordinates use a different view.
-// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @subview_non_alias_call
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @subview_overlap_memref_store
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: memref.subview
+// PRODUCER-CONSUMER-MAXIMAL: memref.store
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @subview_overlap_memref_store(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 1.0 : f64
+ 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[1] [16] [1]
+ : memref<32xf64> to memref<16xf64, strided<[1], offset: 1>>
+ memref.store %cst, %view[%c0]
+ : 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
+}
+
+// -----
+
+// A non-zero-offset subview of unrelated storage must not block an otherwise
+// legal producer-consumer fusion.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @subview_non_alias_memref_store
// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
-// PRODUCER-CONSUMER-MAXIMAL: call @escape_subview_non_alias
+// PRODUCER-CONSUMER-MAXIMAL: memref.store
// 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(
+func.func @subview_non_alias_memref_store(
%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>>
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 1.0 : f64
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>>) -> ()
+ memref.store %cst, %view[%c0]
+ : 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>>)
+
+// -----
+
+// A represented free must remain a dependence through full affine filtering.
+// Otherwise fusing the first and third loops would move a store to %a past
+// its deallocation.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @representable_free_between_loops
+// PRODUCER-CONSUMER-MAXIMAL: %[[A:.*]] = memref.alloc
+// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}}, %[[A]][
+// PRODUCER-CONSUMER-MAXIMAL: memref.dealloc %[[A]]
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+func.func @representable_free_between_loops(
+ %in: memref<32xf64>, %out: memref<32xf64>) {
+ %a = memref.alloc() : memref<32xf64>
+ %b = memref.alloc() : memref<32xf64>
+ affine.for %i = 0 to 16 {
+ %v = affine.load %in[%i] : memref<32xf64>
+ affine.store %v, %a[%i] : memref<32xf64>
+ affine.store %v, %b[%i] : memref<32xf64>
+ }
+ affine.for %k = 0 to 1 {
+ memref.dealloc %a : memref<32xf64>
+ }
+ affine.for %j = 0 to 16 {
+ %v = affine.load %b[%j] : memref<32xf64>
+ affine.store %v, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
+// -----
+
+// An unknown call without memref operands can access a global memref and must
+// not become an isolated node that fusion can cross.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @unknown_call_global_effect
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: call @touch_global
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @unknown_call_global_effect(
+ %in: memref<32xf64>, %out: memref<32xf64>) {
+ %global = memref.get_global @fusion_global : memref<32xf64>
+ affine.for %i = 0 to 16 {
+ %v = affine.load %in[%i] : memref<32xf64>
+ affine.store %v, %global[%i] : memref<32xf64>
+ }
+ func.call @touch_global() : () -> ()
+ affine.for %j = 0 to 16 {
+ %v = affine.load %global[%j] : memref<32xf64>
+ affine.store %v, %out[%j] : memref<32xf64>
+ }
+ return
+}
+memref.global "private" @fusion_global : memref<32xf64>
+func.func private @touch_global()
// -----
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 96a8f0c6218ee..79d9b55cf6e1b 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1576,9 +1576,9 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
// -----
-// Unknown operations nested in an affine loop may access their memref
-// operands. Dependence checking must handle them without assuming a load/store
-// operation class.
+// Unknown operations nested in an affine loop may access memory outside their
+// explicit operands. Fusion must leave the block unchanged when their effects
+// cannot be represented by the memref dependence graph.
// CHECK-LABEL: func @nested_unknown_call
// CHECK: affine.for
>From c3c459a8eeedda0f9e2407473d33bbc01e5a1fd6 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 17:47:49 +0800
Subject: [PATCH 11/23] [MLIR][Affine] Restrict unknown-effect bailout to calls
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 34 ++++++++++++++++-----
mlir/test/Dialect/Affine/loop-fusion-3.mlir | 15 ++++-----
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 9 ++----
3 files changed, 37 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 3a731193bd887..009b0a9126e85 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -22,6 +22,7 @@
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IntegerSet.h"
+#include "mlir/Interfaces/CallInterfaces.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/Support/Debug.h"
@@ -55,8 +56,10 @@ static bool isSameMemref(Value lhs, Value rhs) {
/// Returns the values that `op` may have a memref effect of type `EffectTys`
/// on, not considering recursive effects. View-like values are canonicalized
/// to their storage source so the MDG uses one key for a view chain while raw
-/// views remain available for affine-coordinate analysis. Unknown operations
-/// cannot be represented by the memref-keyed graph, so return false for them.
+/// views remain available for affine-coordinate analysis. Unknown calls cannot
+/// be represented by the memref-keyed graph because their effects are not
+/// limited to explicit memref operands. Other unknown operations retain the
+/// existing operand-based fallback.
/// Returns false if a selected effect cannot be represented by a memref value.
/// The MDG has no resource-level or all-memory edges, so dropping such an
/// effect would make fusion unsound.
@@ -64,8 +67,16 @@ template <typename... EffectTys>
static bool getMayAffectedValues(Operation *op,
SmallVectorImpl<Value> &values) {
auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
- if (!memOp)
- return !hasUnknownEffects(op);
+ if (!memOp) {
+ if (!hasUnknownEffects(op))
+ return true;
+ if (isa<CallOpInterface>(op))
+ return false;
+ for (Value operand : op->getOperands())
+ if (isa<BaseMemRefType>(operand.getType()))
+ values.push_back(canonicalizeMemref(operand));
+ return true;
+ }
SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
memOp.getEffects(effects);
for (auto &effect : effects) {
@@ -107,10 +118,17 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
if (!memInterface) {
if (!hasUnknownEffects(op))
return;
- // Unknown effects may reach memory through globals or other state not
- // represented by SSA memref operands. Keep the graph fail-closed.
- memrefLoads.push_back(op);
- memrefStores.push_back(op);
+ if (isa<CallOpInterface>(op)) {
+ // Calls may access memory not represented by explicit operands.
+ memrefLoads.push_back(op);
+ memrefStores.push_back(op);
+ } else if (llvm::any_of(op->getOperands(), [](Value value) {
+ return isa<BaseMemRefType>(value.getType());
+ })) {
+ // Conservatively, assume all memref operands are read and written.
+ memrefLoads.push_back(op);
+ memrefStores.push_back(op);
+ }
} else {
// Non-affine loads, stores, and frees. Allocation effects are
// intentionally omitted: they do not access existing memory, and
diff --git a/mlir/test/Dialect/Affine/loop-fusion-3.mlir b/mlir/test/Dialect/Affine/loop-fusion-3.mlir
index 70d6c82105543..d8204d0afc376 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-3.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-3.mlir
@@ -869,7 +869,8 @@ func.func @call_op_prevents_fusion(%arg0: memref<16xf32>){
// -----
func.func private @some_function()
-func.func @call_op_does_not_prevent_fusion(%arg0: memref<16xf32>){
+func.func @call_op_without_memref_operands_prevents_fusion(
+ %arg0: memref<16xf32>) {
%A = memref.alloc() : memref<16xf32>
%cst_1 = arith.constant 1.000000e+00 : f32
affine.for %arg1 = 0 to 16 {
@@ -885,9 +886,10 @@ func.func @call_op_does_not_prevent_fusion(%arg0: memref<16xf32>){
}
return
}
-// CHECK-LABEL: func @call_op_does_not_prevent_fusion
+// CHECK-LABEL: func @call_op_without_memref_operands_prevents_fusion
+// CHECK: affine.for
+// CHECK: call @some_function() : () -> ()
// CHECK: affine.for
-// CHECK-NOT: affine.for
// -----
@@ -1281,14 +1283,13 @@ func.func @unknown_memref_def_op() {
affine.for %i1 = 0 to 10 {
%0 = affine.load %may_alias[%i1] : memref<10xf32>
}
- // Fusion happens, but memref isn't privatized since %may_alias's origin is
- // unknown.
+ // The unknown call prevents fusion because its memory effects are not
+ // limited to the returned memref.
// CHECK: call
// CHECK-NEXT: affine.for
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+ // CHECK: affine.for
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
- // CHECK-NEXT: }
- // CHECK-NOT: affine.for
return
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index fc82b34f870bf..7fa411ae6c134 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1137,19 +1137,16 @@ func.func @subview_non_alias_memref_store(
// -----
// A represented free must remain a dependence through full affine filtering.
-// Otherwise fusing the first and third loops would move a store to %a past
-// its deallocation.
+// Otherwise fusing the first and third loops could move a store to %a past its
+// deallocation.
// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @representable_free_between_loops
// PRODUCER-CONSUMER-MAXIMAL: %[[A:.*]] = memref.alloc
// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}}, %[[A]][
// PRODUCER-CONSUMER-MAXIMAL: memref.dealloc %[[A]]
-// PRODUCER-CONSUMER-MAXIMAL: affine.for
-// PRODUCER-CONSUMER-MAXIMAL: affine.load
func.func @representable_free_between_loops(
- %in: memref<32xf64>, %out: memref<32xf64>) {
+ %in: memref<32xf64>, %b: memref<32xf64>, %out: memref<32xf64>) {
%a = memref.alloc() : memref<32xf64>
- %b = memref.alloc() : memref<32xf64>
affine.for %i = 0 to 16 {
%v = affine.load %in[%i] : memref<32xf64>
affine.store %v, %a[%i] : memref<32xf64>
>From df1b83d3b665cc51006a395ce2e336f92e536005 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 17:55:52 +0800
Subject: [PATCH 12/23] [MLIR][Affine] Reuse effect classification in fusion
collection
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 19 +++++++------------
1 file changed, 7 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 009b0a9126e85..5e02828b4714f 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -116,19 +116,14 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
} else {
auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
if (!memInterface) {
- if (!hasUnknownEffects(op))
+ SmallVector<Value> affectedValues;
+ if (getMayAffectedValues<MemoryEffects::Read>(op, affectedValues) &&
+ affectedValues.empty())
return;
- if (isa<CallOpInterface>(op)) {
- // Calls may access memory not represented by explicit operands.
- memrefLoads.push_back(op);
- memrefStores.push_back(op);
- } else if (llvm::any_of(op->getOperands(), [](Value value) {
- return isa<BaseMemRefType>(value.getType());
- })) {
- // Conservatively, assume all memref operands are read and written.
- memrefLoads.push_back(op);
- memrefStores.push_back(op);
- }
+ // Unknown calls may access memory not represented by explicit
+ // operands; other unknown operations reach here with memref operands.
+ memrefLoads.push_back(op);
+ memrefStores.push_back(op);
} else {
// Non-affine loads, stores, and frees. Allocation effects are
// intentionally omitted: they do not access existing memory, and
>From 970ada299ed3921bd69da4b9ce7d778891f4869d Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 20:11:42 +0800
Subject: [PATCH 13/23] [MLIR][Affine] Fail closed on unknown fusion effects
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 28 ++++---------
mlir/test/Dialect/Affine/loop-fusion-2.mlir | 24 +++++------
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 45 +++++++++++++++++++++
mlir/test/Dialect/Affine/loop-fusion.mlir | 34 ++++++++--------
4 files changed, 82 insertions(+), 49 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 5e02828b4714f..70a11acb8b8e0 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -22,7 +22,6 @@
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IntegerSet.h"
-#include "mlir/Interfaces/CallInterfaces.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/Support/Debug.h"
@@ -56,27 +55,17 @@ static bool isSameMemref(Value lhs, Value rhs) {
/// Returns the values that `op` may have a memref effect of type `EffectTys`
/// on, not considering recursive effects. View-like values are canonicalized
/// to their storage source so the MDG uses one key for a view chain while raw
-/// views remain available for affine-coordinate analysis. Unknown calls cannot
-/// be represented by the memref-keyed graph because their effects are not
-/// limited to explicit memref operands. Other unknown operations retain the
-/// existing operand-based fallback.
-/// Returns false if a selected effect cannot be represented by a memref value.
-/// The MDG has no resource-level or all-memory edges, so dropping such an
-/// effect would make fusion unsound.
+/// views remain available for affine-coordinate analysis. Unknown effects
+/// cannot be represented by the memref-keyed graph.
+/// Returns false if a selected effect cannot be represented by a memref value
+/// or if the operation's effects are unknown. The MDG has no resource-level or
+/// all-memory edges, so dropping such an effect would make fusion unsound.
template <typename... EffectTys>
static bool getMayAffectedValues(Operation *op,
SmallVectorImpl<Value> &values) {
auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
- if (!memOp) {
- if (!hasUnknownEffects(op))
- return true;
- if (isa<CallOpInterface>(op))
- return false;
- for (Value operand : op->getOperands())
- if (isa<BaseMemRefType>(operand.getType()))
- values.push_back(canonicalizeMemref(operand));
- return true;
- }
+ if (!memOp)
+ return !hasUnknownEffects(op);
SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
memOp.getEffects(effects);
for (auto &effect : effects) {
@@ -120,8 +109,7 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
if (getMayAffectedValues<MemoryEffects::Read>(op, affectedValues) &&
affectedValues.empty())
return;
- // Unknown calls may access memory not represented by explicit
- // operands; other unknown operations reach here with memref operands.
+ // Unknown effects cannot be represented by the memref-keyed graph.
memrefLoads.push_back(op);
memrefStores.push_back(op);
} else {
diff --git a/mlir/test/Dialect/Affine/loop-fusion-2.mlir b/mlir/test/Dialect/Affine/loop-fusion-2.mlir
index b26a539b2f7d5..2ef55b01164e2 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-2.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-2.mlir
@@ -20,15 +20,15 @@ func.func @should_fuse_at_depth_above_loop_carried_dependence(%arg0: memref<64x4
affine.for %i3 = 0 to 4 {
affine.for %i4 = 0 to 16 {
%v = affine.load %arg1[16 * %i3 - %i4 + 15, %i2] : memref<64x4xf32>
- "op0"(%v) : (f32) -> ()
+ %unused0 = arith.addf %v, %v : f32
}
affine.for %i5 = 0 to 4 {
affine.for %i6 = 0 to 16 {
%v = affine.load %arg0[16 * %i5 - %i6 + 15, %i3] : memref<64x4xf32>
- "op1"(%v) : (f32) -> ()
+ %unused1 = arith.addf %v, %v : f32
}
affine.for %i7 = 0 to 16 {
- %r = "op2"() : () -> (f32)
+ %r = arith.constant 0.0 : f32
%v = affine.load %out[16 * %i5 + %i7, %i2] : memref<64x4xf32>
%s = arith.addf %v, %r : f32
affine.store %s, %out[16 * %i5 + %i7, %i2] : memref<64x4xf32>
@@ -57,15 +57,15 @@ func.func @should_fuse_at_depth_above_loop_carried_dependence(%arg0: memref<64x4
// CHECK-NEXT: affine.for %{{.*}} = 0 to 4 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, %{{.*}}] : memref<64x4xf32>
- // CHECK-NEXT: "op0"(%{{.*}}) : (f32) -> ()
+ // CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 4 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, %{{.*}}] : memref<64x4xf32>
- // CHECK-NEXT: "op1"(%{{.*}}) : (f32) -> ()
+ // CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
- // CHECK-NEXT: %{{.*}} = "op2"() : () -> f32
+ // CHECK-NEXT: %{{.*}} = arith.constant {{.*}} : f32
// CHECK: affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<64x4xf32>
// CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<64x4xf32>
@@ -244,7 +244,7 @@ func.func @slice_tile(%arg0: memref<128x8xf32>, %arg1: memref<32x8xf32>, %0 : f3
affine.for %k = 0 to 8 {
affine.for %kk = 0 to 16 {
%v = affine.load %arg0[16 * %k + %kk, %j] : memref<128x8xf32>
- %r = "foo"(%v) : (f32) -> f32
+ %r = arith.addf %v, %v : f32
}
affine.for %ii = 0 to 16 {
%v = affine.load %arg1[16 * %i + %ii, %j] : memref<32x8xf32>
@@ -264,7 +264,7 @@ func.func @slice_tile(%arg0: memref<128x8xf32>, %arg1: memref<32x8xf32>, %0 : f3
// CHECK-NEXT: affine.for %{{.*}} = 0 to 8 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<128x8xf32>
-// CHECK-NEXT: "foo"(%{{.*}}) : (f32) -> f32
+// CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<32x8xf32>
@@ -463,26 +463,26 @@ func.func @should_not_slice_past_slice_barrier() {
%0 = memref.alloc() : memref<100x16xf32>
affine.for %i0 = 0 to 100 {
affine.for %i1 = 0 to 16 {
- %1 = "op1"() : () -> f32
+ %1 = arith.constant 0.0 : f32
affine.store %1, %0[%i0, %i1] : memref<100x16xf32>
} {slice_fusion_barrier = true}
}
affine.for %i2 = 0 to 100 {
affine.for %i3 = 0 to 16 {
%2 = affine.load %0[%i2, %i3] : memref<100x16xf32>
- "op2"(%2) : (f32) -> ()
+ %unused = arith.addf %2, %2 : f32
}
}
// The 'slice_fusion_barrier' attribute on '%i1' prevents slicing the
// iteration space of '%i1' and any enclosing loop nests.
// CHECK: affine.for %{{.*}} = 0 to 100 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
-// CHECK-NEXT: %{{.*}} = "op1"() : () -> f32
+// CHECK-NEXT: %{{.*}} = arith.constant {{.*}} : f32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
// CHECK-NEXT: } {slice_fusion_barrier = true}
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
// CHECK-NEXT: affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
-// CHECK-NEXT: "op2"(%{{.*}}) : (f32) -> ()
+// CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: }
return
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 7fa411ae6c134..9b13cbc969823 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1188,6 +1188,51 @@ func.func @unknown_call_global_effect(
memref.global "private" @fusion_global : memref<32xf64>
func.func private @touch_global()
+// An unknown non-call operation can access memory not represented by its
+// operands and must not become an isolated node that fusion can cross.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @unknown_non_call_without_operands_prevents_fusion
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: }
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: "unknown.touch_global"() : () -> ()
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+func.func @unknown_non_call_without_operands_prevents_fusion(
+ %in: memref<32xf64>, %out: memref<32xf64>) {
+ %global = memref.get_global @fusion_global : memref<32xf64>
+ affine.for %i = 0 to 16 {
+ %v = affine.load %in[%i] : memref<32xf64>
+ affine.store %v, %global[%i] : memref<32xf64>
+ }
+ "unknown.touch_global"() : () -> ()
+ affine.for %j = 0 to 16 {
+ %v = affine.load %global[%j] : memref<32xf64>
+ affine.store %v, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
+// Explicit operands do not make the effects of an unknown operation complete.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @unknown_non_call_with_memref_operand_may_have_implicit_effects
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: }
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: "unknown.touch_global"(%{{.*}}) : (memref<32xf64>) -> ()
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+func.func @unknown_non_call_with_memref_operand_may_have_implicit_effects(
+ %in: memref<32xf64>, %unrelated: memref<32xf64>, %out: memref<32xf64>) {
+ %global = memref.get_global @fusion_global : memref<32xf64>
+ affine.for %i = 0 to 16 {
+ %v = affine.load %in[%i] : memref<32xf64>
+ affine.store %v, %global[%i] : memref<32xf64>
+ }
+ "unknown.touch_global"(%unrelated) : (memref<32xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %v = affine.load %global[%j] : memref<32xf64>
+ affine.store %v, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
// -----
// An addressable effect without an SSA memory value cannot be represented by
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 79d9b55cf6e1b..5e4064d9c54ae 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -596,7 +596,7 @@ func.func @permute_and_fuse() {
affine.for %i4 = 0 to 10 {
affine.for %i5 = 0 to 20 {
%v0 = affine.load %m[%i4, %i5, %i3] : memref<10x20x30xf32>
- "foo"(%v0) : (f32) -> ()
+ %unused = arith.addf %v0, %v0 : f32
}
}
}
@@ -605,7 +605,7 @@ func.func @permute_and_fuse() {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 20 {
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
// CHECK-NEXT: affine.load %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
-// CHECK-NEXT: "foo"(%{{.*}}) : (f32) -> ()
+// CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
@@ -633,7 +633,7 @@ func.func @fuse_reshape_64_16_4(%in : memref<64xf32>) {
affine.for %i1 = 0 to 16 {
affine.for %i2 = 0 to 4 {
%w = affine.load %out[%i1, %i2] : memref<16x4xf32>
- "foo"(%w) : (f32) -> ()
+ %unused = arith.addf %w, %w : f32
}
}
return
@@ -665,7 +665,7 @@ func.func @fuse_reshape_16_4_64() {
affine.for %i2 = 0 to 64 {
%w = affine.load %out[%i2] : memref<64xf32>
- "foo"(%w) : (f32) -> ()
+ %unused = arith.addf %w, %w : f32
}
// CHECK: affine.for %{{.*}} = 0 to 64 {
// CHECK-NEXT: affine.apply [[$MAP0]](%{{.*}})
@@ -674,7 +674,7 @@ func.func @fuse_reshape_16_4_64() {
// CHECK-NEXT: affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
// CHECK-NEXT: affine.load %{{.*}}[0] : memref<1xf32>
-// CHECK-NEXT: "foo"(%{{.*}}) : (f32) -> ()
+// CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: return
return
@@ -697,7 +697,7 @@ func.func @R6_to_R2_reshape_square() -> memref<64x9xi32> {
affine.for %i3 = 0 to 3 {
affine.for %i4 = 0 to 16 {
affine.for %i5 = 0 to 1 {
- %val = "foo"(%i0, %i1, %i2, %i3, %i4, %i5) : (index, index, index, index, index, index) -> i32
+ %val = arith.constant 0 : i32
affine.store %val, %in[%i0, %i1, %i2, %i3, %i4, %i5] : memref<2x2x3x3x16x1xi32>
}
}
@@ -758,7 +758,7 @@ func.func @R6_to_R2_reshape_square() -> memref<64x9xi32> {
// CHECK-NEXT: affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
// CHECK-NEXT: affine.apply [[$MAP3]](%{{.*}}, %{{.*}})
// CHECK-NEXT: affine.apply [[$MAP4]](%{{.*}}, %{{.*}})
-// CHECK-NEXT: "foo"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index, index, index, index) -> i32
+// CHECK-NEXT: %{{.*}} = arith.constant {{.*}} : i32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, 0, 0, 0, 0, 0] : memref<1x1x1x1x1x1xi32>
// CHECK-NEXT: affine.apply [[$MAP11]](%{{.*}}, %{{.*}})
// CHECK-NEXT: affine.apply [[$MAP12]](%{{.*}})
@@ -812,7 +812,7 @@ func.func @should_fuse_reduction_at_depth_of_one() {
affine.for %i1 = 0 to 100 {
%v0 = affine.load %b[%i0] : memref<10xf32>
%v1 = affine.load %a[%i0, %i1] : memref<10x100xf32>
- %v2 = "maxf"(%v0, %v1) : (f32, f32) -> f32
+ %v2 = arith.addf %v0, %v1 : f32
affine.store %v2, %b[%i0] : memref<10xf32>
}
}
@@ -832,7 +832,7 @@ func.func @should_fuse_reduction_at_depth_of_one() {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 100 {
// CHECK-NEXT: affine.load %{{.*}}[0] : memref<1xf32>
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x100xf32>
- // CHECK-NEXT: "maxf"(%{{.*}}, %{{.*}}) : (f32, f32) -> f32
+ // CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 100 {
@@ -856,10 +856,10 @@ func.func @should_fuse_at_src_depth1_and_dst_depth1() {
affine.for %i0 = 0 to 100 {
affine.for %i1 = 0 to 16 {
%v0 = affine.load %a[%i0, %i1] : memref<100x16xf32>
- "op0"(%v0) : (f32) -> ()
+ %unused0 = arith.addf %v0, %v0 : f32
}
affine.for %i2 = 0 to 16 {
- %v1 = "op1"() : () -> (f32)
+ %v1 = arith.constant 0.0 : f32
affine.store %v1, %b[%i0, %i2] : memref<100x16xf32>
}
}
@@ -867,7 +867,7 @@ func.func @should_fuse_at_src_depth1_and_dst_depth1() {
affine.for %i3 = 0 to 100 {
affine.for %i4 = 0 to 16 {
%v2 = affine.load %b[%i3, %i4] : memref<100x16xf32>
- "op2"(%v2) : (f32) -> ()
+ %unused2 = arith.addf %v2, %v2 : f32
}
}
// We can slice iterations of the '%i0' and '%i1' loops in the source
@@ -878,15 +878,15 @@ func.func @should_fuse_at_src_depth1_and_dst_depth1() {
// CHECK: affine.for %{{.*}} = 0 to 100 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<100x16xf32>
- // CHECK-NEXT: "op0"(%{{.*}}) : (f32) -> ()
+ // CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
- // CHECK-NEXT: %{{.*}} = "op1"() : () -> f32
+ // CHECK-NEXT: %{{.*}} = arith.constant {{.*}} : f32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {
// CHECK-NEXT: affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
- // CHECK-NEXT: "op2"(%{{.*}}) : (f32) -> ()
+ // CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: return
@@ -1304,7 +1304,7 @@ func.func @R3_to_R2_reshape() {
affine.for %i0 = 0 to 2 {
affine.for %i1 = 0 to 3 {
affine.for %i2 = 0 to 16 {
- %val = "foo"(%i0, %i1, %i2) : (index, index, index) -> i32
+ %val = arith.constant 0 : i32
affine.store %val, %in[%i0, %i1, %i2] : memref<2x3x16xi32>
}
}
@@ -1329,7 +1329,7 @@ func.func @R3_to_R2_reshape() {
// CHECK: affine.for %{{.*}} = 0 to 32 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 3 {
// CHECK-NEXT: affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
-// CHECK-NEXT: "foo"(%{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index) -> i32
+// CHECK-NEXT: %{{.*}} = arith.constant {{.*}} : i32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xi32>
// CHECK-NEXT: affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
// CHECK-NEXT: affine.apply [[$MAP2]](%{{.*}})
>From 6a8f8f19e5d92ce381a839da064702612d038dd5 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 22:27:34 +0800
Subject: [PATCH 14/23] [MLIR][Affine] Guard fusion against unmodeled loop
effects
---
.../mlir/Dialect/Affine/Analysis/Utils.h | 10 +-
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 7 +-
.../Dialect/Affine/Transforms/LoopFusion.cpp | 8 +-
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 49 ++++++++++
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 93 +++++++++++++++++++
mlir/test/Dialect/Affine/loop-fusion.mlir | 18 +++-
6 files changed, 169 insertions(+), 16 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index f78fe772d77e9..32ce462f3b552 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
@@ -32,9 +32,8 @@ class AffineForOp;
class AffineValueMap;
struct MemRefAccess;
-// 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.
+// LoopNestStateCollector walks loop nests and collects affine and non-affine
+// memory operations.
struct LoopNestStateCollector {
SmallVector<AffineForOp, 4> forOps;
// Affine loads.
@@ -48,8 +47,7 @@ struct LoopNestStateCollector {
// Free operations.
SmallVector<Operation *, 4> memrefFrees;
- // Collects load and store operations, and whether or not a region holding op
- // other than ForOp and IfOp was encountered in the loop nest.
+ // Collects affine and non-affine memory operations in the loop nest.
void collect(Operation *opToWalk);
};
@@ -248,7 +246,7 @@ struct MemRefDependenceGraph {
ArrayRef<Operation *> memrefStores,
ArrayRef<Operation *> memrefFrees);
- void clearNodeLoadAndStores(unsigned id);
+ void clearNodeMemoryOps(unsigned id);
// Calls 'callback' for each input edge incident to node 'id' which carries a
// memref dependence.
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 70a11acb8b8e0..e273c3f40e799 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -817,7 +817,7 @@ void MemRefDependenceGraph::updateEdges(unsigned sibId, unsigned dstId) {
}
}
-// Adds ops in 'loads' and 'stores' to node at 'id'.
+// Adds all collected memory operations to node at 'id'.
void MemRefDependenceGraph::addToNode(unsigned id, ArrayRef<Operation *> loads,
ArrayRef<Operation *> stores,
ArrayRef<Operation *> memrefLoads,
@@ -831,10 +831,13 @@ void MemRefDependenceGraph::addToNode(unsigned id, ArrayRef<Operation *> loads,
llvm::append_range(node->memrefFrees, memrefFrees);
}
-void MemRefDependenceGraph::clearNodeLoadAndStores(unsigned id) {
+void MemRefDependenceGraph::clearNodeMemoryOps(unsigned id) {
Node *node = getNode(id);
node->loads.clear();
node->stores.clear();
+ node->memrefLoads.clear();
+ node->memrefStores.clear();
+ node->memrefFrees.clear();
}
// Calls 'callback' for each input edge incident to node 'id' which carries a
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index a22a6639fd314..8a6dd4804a06b 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -1175,8 +1175,8 @@ struct GreedyFusion {
LoopNestStateCollector dstLoopCollector;
dstLoopCollector.collect(dstAffineForOp);
- // Clear and add back loads and stores.
- mdg->clearNodeLoadAndStores(dstNode->id);
+ // Clear and add back all memory operations.
+ mdg->clearNodeMemoryOps(dstNode->id);
mdg->addToNode(
dstId, dstLoopCollector.loadOpInsts, dstLoopCollector.storeOpInsts,
dstLoopCollector.memrefLoads, dstLoopCollector.memrefStores,
@@ -1527,8 +1527,8 @@ struct GreedyFusion {
auto dstForInst = cast<AffineForOp>(dstNode->op);
LoopNestStateCollector dstLoopCollector;
dstLoopCollector.collect(dstForInst);
- // Clear and add back loads and stores
- mdg->clearNodeLoadAndStores(dstNode->id);
+ // Clear and add back all memory operations.
+ mdg->clearNodeMemoryOps(dstNode->id);
mdg->addToNode(dstNode->id, dstLoopCollector.loadOpInsts,
dstLoopCollector.storeOpInsts, dstLoopCollector.memrefLoads,
dstLoopCollector.memrefStores, dstLoopCollector.memrefFrees);
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index e51b9ff5dea27..5126c5d317f8e 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -191,6 +191,43 @@ gatherLoadsAndStores(AffineForOp forOp,
return !hasIfOp;
}
+// The fusion transformation clones the complete source loop body into the
+// destination schedule. The affine slice and dependence analysis below does
+// not model the order of non-affine memory effects within that schedule, so
+// keep such effects out of candidate loops until that analysis is extended.
+static bool hasUnmodeledMemoryEffects(AffineForOp forOp) {
+ LoopNestStateCollector collector;
+ collector.collect(forOp);
+ return !collector.memrefLoads.empty() || !collector.memrefStores.empty() ||
+ !collector.memrefFrees.empty();
+}
+
+// Returns true when accesses in two loop bodies use different views of the
+// same storage and at least one access writes. The raw view remains important
+// for affine coordinate analysis, but a storage-level conflict is enough to
+// reject fusion because this check protects the frame condition of the
+// complete loop bodies before strategy-specific filtering.
+static bool hasConflictingStorageAccesses(ArrayRef<Operation *> firstOps,
+ ArrayRef<Operation *> secondOps) {
+ for (Operation *firstOp : firstOps) {
+ MemRefAccess firstAccess(firstOp);
+ Value firstStorage = memref::skipViewLikeOps(
+ cast<MemrefValue>(firstAccess.memref));
+ for (Operation *secondOp : secondOps) {
+ MemRefAccess secondAccess(secondOp);
+ if (firstAccess.memref == secondAccess.memref)
+ continue;
+ Value secondStorage = memref::skipViewLikeOps(
+ cast<MemrefValue>(secondAccess.memref));
+ if (firstStorage != secondStorage)
+ continue;
+ if (firstAccess.isStore() || secondAccess.isStore())
+ return true;
+ }
+ }
+ return false;
+}
+
/// Returns the maximum loop depth at which we could fuse producer loop
/// 'srcForOp' into consumer loop 'dstForOp' without violating data dependences.
// TODO: Generalize this check for sibling and more generic fusion scenarios.
@@ -289,6 +326,12 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
return FusionResult::FailPrecondition;
}
+ if (hasUnmodeledMemoryEffects(srcForOp) ||
+ hasUnmodeledMemoryEffects(dstForOp)) {
+ LDBG() << "Cannot fuse loop nests with unmodeled non-affine memory effects";
+ return FusionResult::FailFusionDependence;
+ }
+
// Return 'failure' if no valid insertion point for fused loop nest in 'block'
// exists which would preserve dependences.
if (!getFusedLoopNestInsertionPoint(srcForOp, dstForOp)) {
@@ -316,6 +359,12 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
return FusionResult::FailPrecondition;
}
+ if (hasConflictingStorageAccesses(opsA, opsB)) {
+ LDBG() << "Fusion would change ordering of accesses through different "
+ "views of the same storage";
+ return FusionResult::FailFusionDependence;
+ }
+
// Return 'failure' if fusing loops at depth 'dstLoopDepth' wouldn't preserve
// loop dependences.
// TODO: Enable this check for sibling and more generic loop fusion
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 9b13cbc969823..c78a1b97fb330 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1382,3 +1382,96 @@ func.func @value_less_nonaddressable_free_effect(
}
return
}
+
+// Non-affine memory effects inside candidate loops are not included in the
+// affine slice analysis, so fusion must preserve the two complete loops.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @non_affine_effects_inside_candidate_loops
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: memref.load
+// PRODUCER-CONSUMER-MAXIMAL: memref.store
+// PRODUCER-CONSUMER-MAXIMAL: }
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: memref.load
+func.func @non_affine_effects_inside_candidate_loops(
+ %in: memref<8xi32>, %out: memref<8xi32>) {
+ %tmp = memref.alloc() : memref<8xi32>
+ %counter = memref.alloc() : memref<1xi32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : i32
+ affine.for %i = 0 to 8 {
+ %v = affine.load %in[%i] : memref<8xi32>
+ %old = memref.load %counter[%c0] : memref<1xi32>
+ %next = arith.addi %old, %c1 : i32
+ memref.store %next, %counter[%c0] : memref<1xi32>
+ affine.store %v, %tmp[%i] : memref<8xi32>
+ }
+ affine.for %j = 0 to 8 {
+ %v = affine.load %tmp[%j] : memref<8xi32>
+ %current = memref.load %counter[%c0] : memref<1xi32>
+ %result = arith.addi %v, %current : i32
+ affine.store %result, %out[%j] : memref<8xi32>
+ }
+ return
+}
+
+// A source read and destination write through overlapping non-trivial views
+// must not be interleaved by producer-consumer fusion.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @overlapping_subviews_inside_loops
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+// PRODUCER-CONSUMER-MAXIMAL: affine.store
+// PRODUCER-CONSUMER-MAXIMAL: }
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.store
+func.func @overlapping_subviews_inside_loops(%root: memref<32xf32>) {
+ %a = memref.subview %root[0] [16] [1]
+ : memref<32xf32> to memref<16xf32, strided<[1], offset: 0>>
+ %b = memref.subview %root[1] [16] [1]
+ : memref<32xf32> to memref<16xf32, strided<[1], offset: 1>>
+ %tmp = memref.alloc() : memref<16xf32>
+ affine.for %i = 0 to 16 {
+ %v = affine.load %a[%i]
+ : memref<16xf32, strided<[1], offset: 0>>
+ affine.store %v, %tmp[%i] : memref<16xf32>
+ }
+ affine.for %j = 0 to 16 {
+ %v = affine.load %tmp[%j] : memref<16xf32>
+ affine.store %v, %b[%j]
+ : memref<16xf32, strided<[1], offset: 1>>
+ }
+ return
+}
+
+// The reverse read/write direction is covered independently so the guard is
+// not accidentally limited to source-read/destination-write pairs.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @overlapping_subviews_write_read
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.store
+// PRODUCER-CONSUMER-MAXIMAL: }
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load
+func.func @overlapping_subviews_write_read(
+ %root: memref<32xf32>, %out: memref<16xf32>) {
+ %a = memref.subview %root[0] [16] [1]
+ : memref<32xf32> to memref<16xf32, strided<[1], offset: 0>>
+ %b = memref.subview %root[1] [16] [1]
+ : memref<32xf32> to memref<16xf32, strided<[1], offset: 1>>
+ %tmp = memref.alloc() : memref<16xf32>
+ %cst = arith.constant 1.0 : f32
+ affine.for %i = 0 to 16 {
+ affine.store %cst, %a[%i]
+ : memref<16xf32, strided<[1], offset: 0>>
+ affine.store %cst, %tmp[%i] : memref<16xf32>
+ }
+ affine.for %j = 0 to 16 {
+ %v = affine.load %tmp[%j] : memref<16xf32>
+ %other = affine.load %b[%j]
+ : memref<16xf32, strided<[1], offset: 1>>
+ %sum = arith.addf %v, %other : f32
+ affine.store %sum, %out[%j] : memref<16xf32>
+ }
+ return
+}
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 5e4064d9c54ae..b55eee5c27b6e 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1583,15 +1583,20 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
// CHECK-LABEL: func @nested_unknown_call
// CHECK: affine.for
// CHECK: func.call @escape_nested
-// CHECK: affine.for
+// CHECK: affine.store
+// CHECK: }
+// CHECK-NEXT: affine.for
// CHECK: affine.load
// CHECK: return
func.func @nested_unknown_call(%m: memref<8xf32>, %out: memref<8xf32>) {
+ %tmp = memref.alloc() : memref<8xf32>
affine.for %i = 0 to 8 {
+ %v = affine.load %m[%i] : memref<8xf32>
func.call @escape_nested(%m) : (memref<8xf32>) -> ()
+ affine.store %v, %tmp[%i] : memref<8xf32>
}
affine.for %i = 0 to 8 {
- %v = affine.load %m[%i] : memref<8xf32>
+ %v = affine.load %tmp[%i] : memref<8xf32>
affine.store %v, %out[%i] : memref<8xf32>
}
return
@@ -1605,16 +1610,21 @@ func.func private @escape_nested(memref<8xf32>)
// CHECK-LABEL: func @nested_memref_copy
// CHECK: affine.for
// CHECK: memref.copy
-// CHECK: affine.for
+// CHECK: affine.store
+// CHECK: }
+// CHECK-NEXT: affine.for
// CHECK: affine.load
// CHECK: return
func.func @nested_memref_copy(
%src: memref<8xf32>, %dst: memref<8xf32>, %out: memref<8xf32>) {
+ %tmp = memref.alloc() : memref<8xf32>
affine.for %i = 0 to 8 {
memref.copy %src, %dst : memref<8xf32> to memref<8xf32>
+ %v = affine.load %src[%i] : memref<8xf32>
+ affine.store %v, %tmp[%i] : memref<8xf32>
}
affine.for %i = 0 to 8 {
- %v = affine.load %dst[%i] : memref<8xf32>
+ %v = affine.load %tmp[%i] : memref<8xf32>
affine.store %v, %out[%i] : memref<8xf32>
}
return
>From dd0c490fd9d3d6e8067dc50c019613c20f443ecf Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 22:30:36 +0800
Subject: [PATCH 15/23] [MLIR][Affine] Format fusion legality guard
---
mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 5126c5d317f8e..ef2a70e8a850b 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -211,14 +211,14 @@ static bool hasConflictingStorageAccesses(ArrayRef<Operation *> firstOps,
ArrayRef<Operation *> secondOps) {
for (Operation *firstOp : firstOps) {
MemRefAccess firstAccess(firstOp);
- Value firstStorage = memref::skipViewLikeOps(
- cast<MemrefValue>(firstAccess.memref));
+ Value firstStorage =
+ memref::skipViewLikeOps(cast<MemrefValue>(firstAccess.memref));
for (Operation *secondOp : secondOps) {
MemRefAccess secondAccess(secondOp);
if (firstAccess.memref == secondAccess.memref)
continue;
- Value secondStorage = memref::skipViewLikeOps(
- cast<MemrefValue>(secondAccess.memref));
+ Value secondStorage =
+ memref::skipViewLikeOps(cast<MemrefValue>(secondAccess.memref));
if (firstStorage != secondStorage)
continue;
if (firstAccess.isStore() || secondAccess.isStore())
>From 9ecdf23e730c4739be30457a0de91a31a0a65bfa Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 22:35:13 +0800
Subject: [PATCH 16/23] [MLIR][Affine] Clarify fusion effect collector
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 5 ++---
1 file changed, 2 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index e273c3f40e799..3d454a2cbc150 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -91,9 +91,8 @@ static bool mayHaveEffect(Operation *op, Value memref) {
return llvm::is_contained(values, canonicalizeMemref(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.
+// LoopNestStateCollector walks loop nests and collects affine and non-affine
+// memory operations.
void LoopNestStateCollector::collect(Operation *opToWalk) {
opToWalk->walk([&](Operation *op) {
if (auto forOp = dyn_cast<AffineForOp>(op)) {
>From 2fb02dbbe67af15804fdd96ea23a6ee772ed1f0f Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 2 Aug 2026 22:47:21 +0800
Subject: [PATCH 17/23] [MLIR][Affine] Keep fusion tests within effect model
---
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 18 +++++++++++++-----
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 10 +++++++---
.../Affine/loop-fusion-dependence-check.mlir | 7 +------
3 files changed, 21 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index ef2a70e8a850b..b6927b61ff52d 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -194,12 +194,20 @@ gatherLoadsAndStores(AffineForOp forOp,
// The fusion transformation clones the complete source loop body into the
// destination schedule. The affine slice and dependence analysis below does
// not model the order of non-affine memory effects within that schedule, so
-// keep such effects out of candidate loops until that analysis is extended.
+// keep representable effects out of candidate loops until that analysis is
+// extended. Unknown effects are rejected during MDG construction.
static bool hasUnmodeledMemoryEffects(AffineForOp forOp) {
- LoopNestStateCollector collector;
- collector.collect(forOp);
- return !collector.memrefLoads.empty() || !collector.memrefStores.empty() ||
- !collector.memrefFrees.empty();
+ bool hasEffects = false;
+ forOp.walk([&](Operation *op) {
+ if (hasEffects || isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
+ return;
+ if (!isa<MemoryEffectOpInterface>(op))
+ return;
+ hasEffects = hasEffect<MemoryEffects::Read>(op) ||
+ hasEffect<MemoryEffects::Write>(op) ||
+ hasEffect<MemoryEffects::Free>(op);
+ });
+ return hasEffects;
}
// Returns true when accesses in two loop bodies use different views of the
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index c78a1b97fb330..81ac12a2aa6db 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -364,7 +364,8 @@ func.func @same_memref_load_multiple_stores(%producer : memref<32xf32>, %produce
#map = affine_map<()[s0] -> (s0 + 5)>
#map1 = affine_map<()[s0] -> (s0 + 17)>
-// Test with non-int/float memref types.
+// Test with non-int/float memref types. The non-affine memref.load in the
+// consumer is not modeled by the affine slice analysis, so fusion is refused.
// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @memref_index_type
func.func @memref_index_type() {
@@ -388,9 +389,12 @@ func.func @memref_index_type() {
%7 = memref.load %alloc[%5, %6] : memref<8x18xf32>
affine.store %7, %alloc_1[%arg3] : memref<3xf32>
}
- // Expect fusion.
+ // Do not fuse a candidate loop containing a non-affine memory effect.
// PRODUCER-CONSUMER-MAXIMAL: affine.for
- // PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+ // PRODUCER-CONSUMER-MAXIMAL: affine.store
+ // PRODUCER-CONSUMER-MAXIMAL: }
+ // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+ // PRODUCER-CONSUMER-MAXIMAL: memref.load
// PRODUCER-CONSUMER-MAXIMAL: return
return
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir b/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir
index 2c53852a8cec9..1314080e4864f 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir
@@ -109,7 +109,6 @@ func.func @should_not_fuse_across_intermediate_store() {
affine.for %i0 = 0 to 10 {
// expected-remark at -1 {{block-level dependence preventing fusion of loop nest 0 into loop nest 1 at depth 0}}
%v0 = affine.load %0[%i0] : memref<10xf32>
- "op0"(%v0) : (f32) -> ()
}
// Should not fuse loop nests '%i0' and '%i1' across top-level store.
@@ -118,7 +117,6 @@ func.func @should_not_fuse_across_intermediate_store() {
affine.for %i1 = 0 to 10 {
// expected-remark at -1 {{block-level dependence preventing fusion of loop nest 1 into loop nest 0 at depth 0}}
%v1 = affine.load %0[%i1] : memref<10xf32>
- "op1"(%v1) : (f32) -> ()
}
return
}
@@ -138,7 +136,6 @@ func.func @should_not_fuse_across_intermediate_load() {
// Should not fuse loop nests '%i0' and '%i1' across top-level load.
%v0 = affine.load %0[%c0] : memref<10xf32>
- "op0"(%v0) : (f32) -> ()
affine.for %i1 = 0 to 10 {
// expected-remark at -1 {{block-level dependence preventing fusion of loop nest 1 into loop nest 0 at depth 0}}
@@ -165,7 +162,6 @@ func.func @should_not_fuse_across_ssa_value_def() {
// Loop nest '%i0" cannot be fused past load from '%1' due to RAW dependence.
%v1 = affine.load %1[%c0] : memref<10xf32>
- "op0"(%v1) : (f32) -> ()
// Loop nest '%i1' cannot be fused past SSA value def '%c2' which it uses.
%c2 = arith.constant 2 : index
@@ -323,7 +319,6 @@ func.func @should_not_fuse_across_ssa_value_def_at_depth1() {
// RAW dependence from store in loop nest '%i1' to 'load %1' prevents
// fusion loop nest '%i1' into loops after load.
%v1 = affine.load %1[%i0, %c0] : memref<10x10xf32>
- "op0"(%v1) : (f32) -> ()
// Loop nest '%i2' cannot be fused past SSA value def '%c2' which it uses.
%c2 = arith.constant 2 : index
@@ -334,4 +329,4 @@ func.func @should_not_fuse_across_ssa_value_def_at_depth1() {
}
}
return
-}
\ No newline at end of file
+}
>From aebb4646519c029103fb412c7db9c7b28fc1d3ff Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Mon, 3 Aug 2026 10:16:13 +0800
Subject: [PATCH 18/23] [MLIR][Affine] Preserve fusion effect and SSA
dependences
---
.../mlir/Dialect/Affine/Analysis/Utils.h | 45 +++--
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 183 ++++++++++++-----
.../Dialect/Affine/Transforms/LoopFusion.cpp | 44 +++--
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 24 ++-
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 186 ++++++++++++++++++
.../Affine/loop-fusion-dependence-check.mlir | 53 +++++
6 files changed, 444 insertions(+), 91 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index 32ce462f3b552..b431527561e85 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
@@ -46,6 +46,9 @@ struct LoopNestStateCollector {
SmallVector<Operation *, 4> memrefStores;
// Free operations.
SmallVector<Operation *, 4> memrefFrees;
+ // True when an operation has effects that cannot be represented by the
+ // memref-keyed dependence graph.
+ bool hasUnmodeledMemoryEffects = false;
// Collects affine and non-affine memory operations in the loop nest.
void collect(Operation *opToWalk);
@@ -125,15 +128,13 @@ struct MemRefDependenceGraph {
// If this edge is stored in Edge = Node.outEdges[i], then
// '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 and this is a memory dependence, then it is the
- // canonical representative of the view-like storage 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).
+ // Whether this is a storage-level memory dependence or an exact SSA
+ // dependence.
+ enum class Kind { Memory, SSA };
+ Kind kind;
+ // The value on which this edge represents a dependence. Memory edges carry
+ // the canonical representative of a view-like storage class. SSA edges
+ // retain the exact defining value, including for memrefs.
Value value;
};
@@ -193,16 +194,18 @@ struct MemRefDependenceGraph {
// argument to) the block. Returns false otherwise.
bool writesToLiveInOrEscapingMemrefs(unsigned id) const;
- // Returns true iff there is an edge from node 'srcId' to node 'dstId' which
- // is for 'value' if non-null, or for any value otherwise. Returns false
- // otherwise.
- bool hasEdge(unsigned srcId, unsigned dstId, Value value = nullptr) const;
+ // Returns true iff there is an edge of 'kind' from node 'srcId' to node
+ // 'dstId' which is for 'value' if non-null, or for any value otherwise.
+ bool hasEdge(unsigned srcId, unsigned dstId, Value value = nullptr,
+ Edge::Kind kind = Edge::Kind::Memory) const;
// Adds an edge from node 'srcId' to node 'dstId' for 'value'.
- void addEdge(unsigned srcId, unsigned dstId, Value value);
+ void addEdge(unsigned srcId, unsigned dstId, Value value,
+ Edge::Kind kind = Edge::Kind::Memory);
// Removes an edge from node 'srcId' to node 'dstId' for 'value'.
- void removeEdge(unsigned srcId, unsigned dstId, Value value);
+ void removeEdge(unsigned srcId, unsigned dstId, Value value,
+ Edge::Kind kind = Edge::Kind::Memory);
// Returns true if there is a path in the dependence graph from node 'srcId'
// to node 'dstId'. Returns false otherwise. `srcId`, `dstId`, and the
@@ -248,18 +251,18 @@ struct MemRefDependenceGraph {
void clearNodeMemoryOps(unsigned id);
- // Calls 'callback' for each input edge incident to node 'id' which carries a
- // memref dependence.
+ // Calls 'callback' for each input edge incident to node 'id' whose value is a
+ // memref, whether it carries a memory or an SSA dependence.
void forEachMemRefInputEdge(unsigned id,
const std::function<void(Edge)> &callback);
- // Calls 'callback' for each output edge from node 'id' which carries a
- // memref dependence.
+ // Calls 'callback' for each output edge from node 'id' whose value is a
+ // memref, whether it carries a memory or an SSA dependence.
void forEachMemRefOutputEdge(unsigned id,
const std::function<void(Edge)> &callback);
- // Calls 'callback' for each edge in 'edges' which carries a memref
- // dependence.
+ // Calls 'callback' for each edge in 'edges' whose value is a memref, whether
+ // it carries a memory or an SSA dependence.
void forEachMemRefEdge(ArrayRef<Edge> edges,
const std::function<void(Edge)> &callback);
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 3d454a2cbc150..b4d0bfe970270 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -52,6 +52,80 @@ static bool isSameMemref(Value lhs, Value rhs) {
return canonicalizeMemref(lhs) == canonicalizeMemref(rhs);
}
+static bool edgeMatchesMemref(const MemRefDependenceGraph::Edge &edge,
+ Value memref) {
+ if (edge.kind == MemRefDependenceGraph::Edge::Kind::Memory)
+ return isSameMemref(edge.value, memref);
+ return edge.value == memref;
+}
+
+/// Returns true if an affine memory interface operation has an effect outside
+/// the access represented by its affine interface. The affine access is the
+/// graph's precise model; any additional effect must be represented
+/// separately or the graph must fail closed.
+static bool hasUnmodeledAffineMemoryEffects(Operation *op) {
+ auto affineRead = dyn_cast<AffineReadOpInterface>(op);
+ auto affineWrite = dyn_cast<AffineWriteOpInterface>(op);
+ if (!affineRead && !affineWrite)
+ return false;
+
+ // An operation with both roles has no single affine access model here.
+ if (affineRead && affineWrite)
+ return true;
+
+ auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!memInterface) {
+ // The affine access interface does not establish that the operation has
+ // no additional memory effects.
+ return true;
+ }
+
+ SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
+ memInterface.getEffects(effects);
+ if (effects.size() != 1)
+ return true;
+
+ const auto &effect = effects.front();
+ if (effect.getResource() != SideEffects::DefaultResource::get() ||
+ effect.getEffectOnFullRegion())
+ return true;
+ Value memref = affineRead ? affineRead.getMemRef() : affineWrite.getMemRef();
+ if (effect.getValue() != memref)
+ return true;
+ if (affineRead && !isa<MemoryEffects::Read>(effect.getEffect()))
+ return true;
+ if (affineWrite && !isa<MemoryEffects::Write>(effect.getEffect()))
+ return true;
+ return false;
+}
+
+/// Returns true if an effect can be represented by the MDG's per-memref model.
+/// Allocation is represented only for a memref result defined by the same
+/// operation; candidate-loop legality rejects all allocation effects because
+/// slicing can otherwise change their dynamic execution count.
+static bool isRepresentableMemoryEffect(
+ Operation *op,
+ const SideEffects::EffectInstance<MemoryEffects::Effect> &effect) {
+ if (isa<MemoryEffects::Read, MemoryEffects::Write, MemoryEffects::Free>(
+ effect.getEffect()))
+ return effect.getValue() &&
+ effect.getResource() == SideEffects::DefaultResource::get() &&
+ isa<BaseMemRefType>(effect.getValue().getType());
+
+ if (isa<MemoryEffects::Allocate>(effect.getEffect())) {
+ Value value = effect.getValue();
+ // Allocation effects on a result are represented by the exact SSA edge
+ // from the defining operation. Standard memref allocations use a
+ // full-region allocation effect, so this check must be based on the
+ // result identity rather than effectOnFullRegion.
+ return effect.getResource() == SideEffects::DefaultResource::get() &&
+ value && isa<BaseMemRefType>(value.getType()) &&
+ value.getDefiningOp() == op;
+ }
+
+ return false;
+}
+
/// Returns the values that `op` may have a memref effect of type `EffectTys`
/// on, not considering recursive effects. View-like values are canonicalized
/// to their storage source so the MDG uses one key for a view chain while raw
@@ -98,12 +172,16 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
if (auto forOp = dyn_cast<AffineForOp>(op)) {
forOps.push_back(forOp);
} else if (isa<AffineReadOpInterface>(op)) {
+ hasUnmodeledMemoryEffects |= hasUnmodeledAffineMemoryEffects(op);
loadOpInsts.push_back(op);
} else if (isa<AffineWriteOpInterface>(op)) {
+ hasUnmodeledMemoryEffects |= hasUnmodeledAffineMemoryEffects(op);
storeOpInsts.push_back(op);
} else {
auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
if (!memInterface) {
+ if (hasUnknownEffects(op))
+ hasUnmodeledMemoryEffects = true;
SmallVector<Value> affectedValues;
if (getMayAffectedValues<MemoryEffects::Read>(op, affectedValues) &&
affectedValues.empty())
@@ -112,10 +190,16 @@ void LoopNestStateCollector::collect(Operation *opToWalk) {
memrefLoads.push_back(op);
memrefStores.push_back(op);
} else {
- // Non-affine loads, stores, and frees. Allocation effects are
- // intentionally omitted: they do not access existing memory, and
- // allocation results are handled by existing SSA and local-allocation
- // analysis instead of the memref access graph.
+ SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4>
+ effects;
+ memInterface.getEffects(effects);
+ for (const auto &effect : effects)
+ if (!isRepresentableMemoryEffect(op, effect))
+ hasUnmodeledMemoryEffects = true;
+
+ // Non-affine loads, stores, and frees. Allocation effects on local
+ // memref results remain represented by SSA edges; other allocation or
+ // resource effects make the graph fail closed.
if (hasEffect<MemoryEffects::Read>(op))
memrefLoads.push_back(op);
if (hasEffect<MemoryEffects::Write>(op))
@@ -219,6 +303,8 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
// all loads and store accesses it contains.
LoopNestStateCollector collector;
collector.collect(nodeOp);
+ if (collector.hasUnmodeledMemoryEffects)
+ return nullptr;
unsigned newNodeId = mdg.nextNodeId++;
Node &node = nodes.insert({newNodeId, Node(newNodeId, nodeOp)}).first->second;
for (Operation *op : collector.loadOpInsts) {
@@ -374,6 +460,8 @@ bool MemRefDependenceGraph::init(bool fullAffineDependences) {
return false;
forToNodeMap[&op] = node->id;
} else if (isa<AffineReadOpInterface>(op)) {
+ if (hasUnmodeledAffineMemoryEffects(&op))
+ return false;
// Create graph node for top-level load op.
Node node(nextNodeId++, &op);
node.loads.push_back(&op);
@@ -382,6 +470,8 @@ bool MemRefDependenceGraph::init(bool fullAffineDependences) {
memrefAccesses[memref].insert(node.id);
nodes.insert({node.id, node});
} else if (isa<AffineWriteOpInterface>(op)) {
+ if (hasUnmodeledAffineMemoryEffects(&op))
+ return false;
// Create graph node for top-level store op.
Node node(nextNodeId++, &op);
node.stores.push_back(&op);
@@ -444,7 +534,7 @@ bool MemRefDependenceGraph::init(bool fullAffineDependences) {
continue;
assert(forToNodeMap.count(*it) > 0 && "missing mapping");
unsigned userLoopNestId = forToNodeMap[*it];
- addEdge(node.id, userLoopNestId, value);
+ addEdge(node.id, userLoopNestId, value, Edge::Kind::SSA);
}
}
}
@@ -468,7 +558,7 @@ bool MemRefDependenceGraph::init(bool fullAffineDependences) {
// Check precise affine deps if asked for; otherwise, conservative.
if (!fullAffineDependences ||
mayDependence(*srcNode, *dstNode, srcMemRef))
- addEdge(srcId, dstId, srcMemRef);
+ addEdge(srcId, dstId, srcMemRef, Edge::Kind::Memory);
}
}
}
@@ -504,14 +594,14 @@ void MemRefDependenceGraph::removeNode(unsigned id) {
if (inEdges.count(id) > 0) {
SmallVector<Edge, 2> oldInEdges = inEdges[id];
for (auto &inEdge : oldInEdges) {
- removeEdge(inEdge.id, id, inEdge.value);
+ removeEdge(inEdge.id, id, inEdge.value, inEdge.kind);
}
}
// Remove each edge in 'outEdges[id]'.
if (outEdges.contains(id)) {
SmallVector<Edge, 2> oldOutEdges = outEdges[id];
for (auto &outEdge : oldOutEdges) {
- removeEdge(id, outEdge.id, outEdge.value);
+ removeEdge(id, outEdge.id, outEdge.value, outEdge.kind);
}
}
// Erase remaining node state.
@@ -543,28 +633,29 @@ bool MemRefDependenceGraph::writesToLiveInOrEscapingMemrefs(unsigned id) const {
// is for 'value' if non-null, or for any value otherwise. Returns false
// otherwise.
bool MemRefDependenceGraph::hasEdge(unsigned srcId, unsigned dstId,
- Value value) const {
+ Value value, Edge::Kind kind) const {
if (!outEdges.contains(srcId) || !inEdges.contains(dstId)) {
return false;
}
bool hasOutEdge = llvm::any_of(outEdges.lookup(srcId), [=](const Edge &edge) {
- return edge.id == dstId && (!value || edge.value == value);
+ return edge.id == dstId && edge.kind == kind &&
+ (!value || edge.value == value);
});
bool hasInEdge = llvm::any_of(inEdges.lookup(dstId), [=](const Edge &edge) {
- return edge.id == srcId && (!value || edge.value == value);
+ return edge.id == srcId && edge.kind == kind &&
+ (!value || edge.value == value);
});
return hasOutEdge && hasInEdge;
}
// Adds an edge from node 'srcId' to node 'dstId' for 'value'.
void MemRefDependenceGraph::addEdge(unsigned srcId, unsigned dstId,
- Value 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});
+ Value value, Edge::Kind kind) {
+ if (!hasEdge(srcId, dstId, value, kind)) {
+ outEdges[srcId].push_back({dstId, kind, value});
+ inEdges[dstId].push_back({srcId, kind, value});
+ // Allocation cleanup is storage-based, so both memory and SSA memref
+ // edges keep the canonical storage object live.
if (isa<BaseMemRefType>(value.getType()))
memrefEdgeCount[canonicalizeMemref(value)]++;
}
@@ -572,7 +663,7 @@ void MemRefDependenceGraph::addEdge(unsigned srcId, unsigned dstId,
// Removes an edge from node 'srcId' to node 'dstId' for 'value'.
void MemRefDependenceGraph::removeEdge(unsigned srcId, unsigned dstId,
- Value value) {
+ Value value, Edge::Kind kind) {
assert(inEdges.count(dstId) > 0);
assert(outEdges.count(srcId) > 0);
if (isa<BaseMemRefType>(value.getType())) {
@@ -582,14 +673,14 @@ void MemRefDependenceGraph::removeEdge(unsigned srcId, unsigned dstId,
}
// Remove 'srcId' from 'inEdges[dstId]'.
for (auto *it = inEdges[dstId].begin(); it != inEdges[dstId].end(); ++it) {
- if ((*it).id == srcId && (*it).value == value) {
+ if ((*it).id == srcId && (*it).kind == kind && (*it).value == value) {
inEdges[dstId].erase(it);
break;
}
}
// Remove 'dstId' from 'outEdges[srcId]'.
for (auto *it = outEdges[srcId].begin(); it != outEdges[srcId].end(); ++it) {
- if ((*it).id == dstId && (*it).value == value) {
+ if ((*it).id == dstId && (*it).kind == kind && (*it).value == value) {
outEdges[srcId].erase(it);
break;
}
@@ -638,7 +729,7 @@ unsigned MemRefDependenceGraph::getIncomingMemRefAccesses(unsigned id,
Value memref) const {
unsigned inEdgeCount = 0;
for (const Edge &inEdge : inEdges.lookup(id)) {
- if (isSameMemref(inEdge.value, memref)) {
+ if (edgeMatchesMemref(inEdge, memref)) {
const Node *srcNode = getNode(inEdge.id);
// Only count in edges from 'srcNode' if 'srcNode' accesses 'memref'
if (srcNode->getStoreOpCount(memref) > 0)
@@ -654,7 +745,7 @@ unsigned MemRefDependenceGraph::getOutEdgeCount(unsigned id,
Value memref) const {
unsigned outEdgeCount = 0;
for (const auto &outEdge : outEdges.lookup(id))
- if (!memref || isSameMemref(outEdge.value, memref))
+ if (!memref || edgeMatchesMemref(outEdge, memref))
++outEdgeCount;
return outEdgeCount;
}
@@ -663,10 +754,9 @@ unsigned MemRefDependenceGraph::getOutEdgeCount(unsigned id,
void MemRefDependenceGraph::gatherDefiningNodes(
unsigned id, DenseSet<unsigned> &definingNodes) const {
for (const Edge &edge : inEdges.lookup(id))
- // 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<BaseMemRefType>(edge.value.getType()))
+ // SSA edges include memref definitions; memory edges are storage
+ // dependences even when their canonical value is a memref.
+ if (edge.kind == Edge::Kind::SSA)
definingNodes.insert(edge.id);
}
@@ -758,9 +848,9 @@ void MemRefDependenceGraph::updateEdges(unsigned srcId, unsigned dstId,
for (auto &inEdge : oldInEdges) {
// Add edge from 'inEdge.id' to 'dstId' if it's not a private memref.
if (!llvm::any_of(privateMemRefs, [&](Value privateMemRef) {
- return isSameMemref(privateMemRef, inEdge.value);
+ return edgeMatchesMemref(inEdge, privateMemRef);
}))
- addEdge(inEdge.id, dstId, inEdge.value);
+ addEdge(inEdge.id, dstId, inEdge.value, inEdge.kind);
}
}
// For each edge in 'outEdges[srcId]': remove edge from 'srcId' to 'dstId'.
@@ -770,10 +860,10 @@ void MemRefDependenceGraph::updateEdges(unsigned srcId, unsigned dstId,
for (auto &outEdge : oldOutEdges) {
// Remove any out edges from 'srcId' to 'dstId' across memrefs.
if (outEdge.id == dstId)
- removeEdge(srcId, outEdge.id, outEdge.value);
+ removeEdge(srcId, outEdge.id, outEdge.value, outEdge.kind);
else if (removeSrcId) {
- addEdge(dstId, outEdge.id, outEdge.value);
- removeEdge(srcId, outEdge.id, outEdge.value);
+ addEdge(dstId, outEdge.id, outEdge.value, outEdge.kind);
+ removeEdge(srcId, outEdge.id, outEdge.value, outEdge.kind);
}
}
}
@@ -784,9 +874,9 @@ void MemRefDependenceGraph::updateEdges(unsigned srcId, unsigned dstId,
SmallVector<Edge, 2> oldInEdges = inEdges[dstId];
for (auto &inEdge : oldInEdges)
if (llvm::any_of(privateMemRefs, [&](Value privateMemRef) {
- return isSameMemref(privateMemRef, inEdge.value);
+ return edgeMatchesMemref(inEdge, privateMemRef);
}))
- removeEdge(inEdge.id, dstId, inEdge.value);
+ removeEdge(inEdge.id, dstId, inEdge.value, inEdge.kind);
}
}
@@ -799,8 +889,8 @@ void MemRefDependenceGraph::updateEdges(unsigned sibId, unsigned dstId) {
if (inEdges.count(sibId) > 0) {
SmallVector<Edge, 2> oldInEdges = inEdges[sibId];
for (auto &inEdge : oldInEdges) {
- addEdge(inEdge.id, dstId, inEdge.value);
- removeEdge(inEdge.id, sibId, inEdge.value);
+ addEdge(inEdge.id, dstId, inEdge.value, inEdge.kind);
+ removeEdge(inEdge.id, sibId, inEdge.value, inEdge.kind);
}
}
@@ -810,8 +900,8 @@ void MemRefDependenceGraph::updateEdges(unsigned sibId, unsigned dstId) {
if (outEdges.count(sibId) > 0) {
SmallVector<Edge, 2> oldOutEdges = outEdges[sibId];
for (auto &outEdge : oldOutEdges) {
- addEdge(dstId, outEdge.id, outEdge.value);
- removeEdge(sibId, outEdge.id, outEdge.value);
+ addEdge(dstId, outEdge.id, outEdge.value, outEdge.kind);
+ removeEdge(sibId, outEdge.id, outEdge.value, outEdge.kind);
}
}
}
@@ -839,29 +929,30 @@ void MemRefDependenceGraph::clearNodeMemoryOps(unsigned id) {
node->memrefFrees.clear();
}
-// Calls 'callback' for each input edge incident to node 'id' which carries a
-// memref dependence.
+// Calls 'callback' for each input edge incident to node 'id' whose value is a
+// memref, whether it carries a memory or an SSA dependence.
void MemRefDependenceGraph::forEachMemRefInputEdge(
unsigned id, const std::function<void(Edge)> &callback) {
if (inEdges.count(id) > 0)
forEachMemRefEdge(inEdges.at(id), callback);
}
-// Calls 'callback' for each output edge from node 'id' which carries a
-// memref dependence.
+// Calls 'callback' for each output edge from node 'id' whose value is a
+// memref, whether it carries a memory or an SSA dependence.
void MemRefDependenceGraph::forEachMemRefOutputEdge(
unsigned id, const std::function<void(Edge)> &callback) {
if (outEdges.count(id) > 0)
forEachMemRefEdge(outEdges.at(id), callback);
}
-// Calls 'callback' for each edge in 'edges' which carries a memref
-// dependence.
+// Calls 'callback' for each edge in 'edges' whose value is a memref, whether it
+// carries a memory or an SSA dependence.
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<BaseMemRefType>(edge.value.getType()))
+ // Skip non-memref SSA edges; memory edges always carry a memref value.
+ if (edge.kind == Edge::Kind::SSA &&
+ !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 8a6dd4804a06b..c8e3163e0ce24 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -772,8 +772,9 @@ namespace {
// *) Update graph state to reflect the fusion of 'sibNode' into 'dstNode'.
//
// Given a graph where top-level operations are vertices in the set 'V' and
-// edges in the set 'E' are dependences between vertices, this algorithm
-// takes O(V) time for initialization, and has runtime O(V + E).
+// edges in the set 'E' are dependences between vertices, graph traversal is
+// linear in the graph size. Graph construction and fusion legality also
+// perform pairwise access and dependence checks for candidate loop depths.
//
// This greedy algorithm is not 'maximal' due to the current restriction of
// fusing along single producer consumer edges, but there is a TODO: to fix
@@ -880,10 +881,12 @@ struct GreedyFusion {
// cannot create a private memref.
if (removeSrcNode &&
any_of(mdg->outEdges[producerId], [&](const auto &edge) {
- return edge.id != consumerId &&
- isa<BaseMemRefType>(edge.value.getType()) &&
- memref::isSameViewOrTrivialAlias(cast<MemrefValue>(edge.value),
- cast<MemrefValue>(memref));
+ if (edge.id == consumerId)
+ return false;
+ if (edge.kind == MemRefDependenceGraph::Edge::Kind::SSA)
+ return edge.value == memref;
+ return memref::isSameViewOrTrivialAlias(
+ cast<MemrefValue>(edge.value), cast<MemrefValue>(memref));
}))
return false;
@@ -1104,14 +1107,13 @@ struct GreedyFusion {
srcId, dstId, bestSlice, fusedLoopInsPoint, srcEscapingMemRefs,
*mdg);
- DenseSet<Value> privateMemrefs;
+ DenseSet<Value> privateMemrefsToCreate;
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);
+ privateMemrefsToCreate.insert(memref);
}
}
@@ -1127,12 +1129,11 @@ 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()) {
+ // Create private memrefs. Only record a memref for edge updates after
+ // its replacement has succeeded; a failed private-buffer creation
+ // leaves the original memref uses and their dependences intact.
+ DenseSet<Value> privateMemrefs;
+ if (!privateMemrefsToCreate.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();
@@ -1141,7 +1142,7 @@ struct GreedyFusion {
DenseMap<Value, SmallVector<Operation *, 4>> privateMemRefToStores;
dstAffineForOp.walk([&](AffineWriteOpInterface storeOp) {
Value storeMemRef = storeOp.getMemRef();
- if (llvm::any_of(privateMemrefs, [&](Value privateMemref) {
+ if (llvm::any_of(privateMemrefsToCreate, [&](Value privateMemref) {
return memref::isSameViewOrTrivialAlias(
cast<MemrefValue>(privateMemref),
cast<MemrefValue>(storeMemRef));
@@ -1154,16 +1155,21 @@ struct GreedyFusion {
// loads and stores. Any reference to the original ones becomes
// invalid after this point.
for (auto &memrefToStoresPair : privateMemRefToStores) {
+ // Capture the pre-replacement value; the store operands are
+ // rewritten before createPrivateMemRef returns.
+ Value oldMemRef = memrefToStoresPair.first;
ArrayRef<Operation *> storesForMemref = memrefToStoresPair.second;
Value newMemRef = createPrivateMemRef(
dstAffineForOp, storesForMemref, bestDstLoopDepth,
fastMemorySpace, sliceInsertionBlock, localBufSizeThreshold);
if (!newMemRef)
continue;
+ privateMemrefs.insert(oldMemRef);
// Create new node in dependence graph for 'newMemRef' alloc op.
unsigned newMemRefNodeId = mdg->addNode(newMemRef.getDefiningOp());
// Add edge from 'newMemRef' node to dstNode.
- mdg->addEdge(newMemRefNodeId, dstId, newMemRef);
+ mdg->addEdge(newMemRefNodeId, dstId, newMemRef,
+ MemRefDependenceGraph::Edge::Kind::SSA);
}
// One or more entries for 'newMemRef' alloc op are inserted into
// the DenseMap mdg->nodes. Since an insertion may cause DenseMap to
@@ -1171,6 +1177,10 @@ struct GreedyFusion {
dstNode = mdg->getNode(dstId);
}
+ // Update edges between 'srcNode' and 'dstNode' after IR mutation and
+ // only for private memrefs whose replacements succeeded.
+ mdg->updateEdges(srcId, dstId, privateMemrefs, removeSrcNode);
+
// Collect dst loop stats after memref privatization transformation.
LoopNestStateCollector dstLoopCollector;
dstLoopCollector.collect(dstAffineForOp);
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index b6927b61ff52d..c48b902f8489a 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -193,19 +193,29 @@ gatherLoadsAndStores(AffineForOp forOp,
// The fusion transformation clones the complete source loop body into the
// destination schedule. The affine slice and dependence analysis below does
-// not model the order of non-affine memory effects within that schedule, so
-// keep representable effects out of candidate loops until that analysis is
-// extended. Unknown effects are rejected during MDG construction.
+// not model the order or execution count of non-affine effects within that
+// schedule, so reject every such effect here. This check is intentionally
+// independent of MDG construction because this utility is also called
+// directly by tests and other clients.
static bool hasUnmodeledMemoryEffects(AffineForOp forOp) {
+ LoopNestStateCollector collector;
+ collector.collect(forOp);
+ if (collector.hasUnmodeledMemoryEffects)
+ return true;
+
bool hasEffects = false;
forOp.walk([&](Operation *op) {
if (hasEffects || isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
return;
- if (!isa<MemoryEffectOpInterface>(op))
+ auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!memInterface) {
+ hasEffects = hasUnknownEffects(op);
return;
- hasEffects = hasEffect<MemoryEffects::Read>(op) ||
- hasEffect<MemoryEffects::Write>(op) ||
- hasEffect<MemoryEffects::Free>(op);
+ }
+
+ SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
+ memInterface.getEffects(effects);
+ hasEffects = !effects.empty();
});
return hasEffects;
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 81ac12a2aa6db..0887e288fd384 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -401,6 +401,29 @@ func.func @memref_index_type() {
// -----
+// Keep a positive affine-only case for non-float element types. The element
+// type itself is not a reason to reject fusion.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @memref_index_element_type_fuses
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: return
+func.func @memref_index_element_type_fuses(
+ %in: memref<3xindex>, %out: memref<3xindex>) {
+ %tmp = memref.alloc() : memref<3xindex>
+ affine.for %i = 0 to 3 {
+ %v = affine.load %in[%i] : memref<3xindex>
+ affine.store %v, %tmp[%i] : memref<3xindex>
+ }
+ affine.for %j = 0 to 3 {
+ %v = affine.load %tmp[%j] : memref<3xindex>
+ affine.store %v, %out[%j] : memref<3xindex>
+ }
+ return
+}
+
+// -----
+
#map = affine_map<(d0) -> (d0)>
#map1 =affine_map<(d0) -> (d0 + 1)>
@@ -1045,6 +1068,146 @@ func.func private @escape(memref<?xf64>)
// -----
+// A representable store through a cast alias must retain the producer-consumer
+// dependence without relying on an opaque external call.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_memref_store
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: memref.cast
+// PRODUCER-CONSUMER-MAXIMAL: memref.store
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @cast_alias_memref_store(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 1.0 : f64
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ affine.store %a, %comm[%i] : memref<32xf64>
+ }
+ %view = memref.cast %comm : memref<32xf64> to memref<?xf64>
+ memref.store %cst, %view[%c0] : memref<?xf64>
+ affine.for %j = 0 to 8 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ affine.store %c, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
+// The same representable effect on unrelated storage must not block fusion.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_unrelated_memref_store
+// PRODUCER-CONSUMER-MAXIMAL: memref.store
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: return
+func.func @cast_alias_unrelated_memref_store(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ %other = memref.alloc() : memref<32xf64>
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 1.0 : f64
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ affine.store %a, %comm[%i] : memref<32xf64>
+ }
+ %view = memref.cast %other : memref<32xf64> to memref<?xf64>
+ memref.store %cst, %view[%c0] : memref<?xf64>
+ affine.for %j = 0 to 8 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ affine.store %c, %out[%j] : memref<32xf64>
+ }
+ return
+}
+
+// -----
+
+// A private replacement of view A must not remove the exact SSA dependence on
+// sibling view B. The second fusion opportunity is deliberately separated by
+// a use of the value produced by B; dropping that edge would move the fused
+// destination before B's definition.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @incremental_ssa_view_edge_preserved
+// PRODUCER-CONSUMER-MAXIMAL: %[[B:.*]] = memref.subview
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>
+// PRODUCER-CONSUMER-MAXIMAL: memref.dim %[[B]]
+func.func @incremental_ssa_view_edge_preserved(
+ %in: memref<32xf32>, %out: memref<8xf32>) {
+ %base = memref.alloc() : memref<32xf32>
+ %a = memref.cast %base : memref<32xf32> to memref<?xf32>
+ %y = memref.alloc() : memref<16xf32>
+ %dout = memref.alloc() : memref<16xf32>
+ affine.for %py = 0 to 16 {
+ %v = affine.load %in[%py] : memref<32xf32>
+ affine.store %v, %y[%py] : memref<16xf32>
+ }
+ affine.for %d = 0 to 16 {
+ %v = affine.load %y[%d] : memref<16xf32>
+ affine.store %v, %dout[%d] : memref<16xf32>
+ }
+ affine.for %pa = 0 to 16 {
+ %v = affine.load %in[%pa] : memref<32xf32>
+ affine.store %v, %a[%pa] : memref<?xf32>
+ }
+ %b = memref.subview %base[0] [32] [1]
+ : memref<32xf32> to memref<32xf32, strided<[1], offset: 0>>
+ %c0 = arith.constant 0 : index
+ affine.for %c = 0 to 8 {
+ %dim = memref.dim %b, %c0
+ : memref<32xf32, strided<[1], offset: 0>>
+ %av = affine.load %a[%c] : memref<?xf32>
+ %yv = affine.load %y[%c] : memref<16xf32>
+ %sum = arith.addf %av, %yv : f32
+ affine.store %sum, %out[%c] : memref<8xf32>
+ }
+ return
+}
+
+// -----
+
+// If private-buffer creation fails for multiple producer stores with different
+// access functions, the original graph edges must remain available to the next
+// fusion attempt.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @failed_private_memref_keeps_ssa_edge
+// PRODUCER-CONSUMER-MAXIMAL: %[[FAILED_B:.*]] = memref.subview
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>
+// PRODUCER-CONSUMER-MAXIMAL: memref.dim %[[FAILED_B]]
+func.func @failed_private_memref_keeps_ssa_edge(
+ %in: memref<32xf32>, %out: memref<8xf32>) {
+ %base = memref.alloc() : memref<32xf32>
+ %a = memref.cast %base : memref<32xf32> to memref<?xf32>
+ %y = memref.alloc() : memref<16xf32>
+ %dout = memref.alloc() : memref<16xf32>
+ affine.for %py = 0 to 16 {
+ %v = affine.load %in[%py] : memref<32xf32>
+ affine.store %v, %y[%py] : memref<16xf32>
+ }
+ affine.for %d = 0 to 16 {
+ %v = affine.load %y[%d] : memref<16xf32>
+ affine.store %v, %dout[%d] : memref<16xf32>
+ }
+ affine.for %pa = 0 to 16 {
+ %v = affine.load %in[%pa] : memref<32xf32>
+ affine.store %v, %a[%pa] : memref<?xf32>
+ affine.store %v, %a[%pa + 1] : memref<?xf32>
+ }
+ %b = memref.subview %base[0] [32] [1]
+ : memref<32xf32> to memref<32xf32, strided<[1], offset: 0>>
+ %c0 = arith.constant 0 : index
+ affine.for %c = 0 to 8 {
+ %dim = memref.dim %b, %c0
+ : memref<32xf32, strided<[1], offset: 0>>
+ %av = affine.load %a[%c] : memref<?xf32>
+ %yv = affine.load %y[%c] : memref<16xf32>
+ %sum = arith.addf %av, %yv : f32
+ affine.store %sum, %out[%c] : memref<8xf32>
+ }
+ return
+}
+
+// -----
+
// A zero-offset, unit-stride subview is a fully aliasing view and must retain
// the same dependence as its source memref.
@@ -1419,6 +1582,29 @@ func.func @non_affine_effects_inside_candidate_loops(
return
}
+// Allocation is an effect even when it has no memref value. Slicing the source
+// loop must not change the number or order of such effects.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @allocation_effect_inside_candidate_loop
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
+// PRODUCER-CONSUMER-MAXIMAL: }
+// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+func.func @allocation_effect_inside_candidate_loop(
+ %in: memref<16xi32>, %out: memref<16xi32>) {
+ %tmp = memref.alloc() : memref<16xi32>
+ affine.for %i = 0 to 16 {
+ %v = affine.load %in[%i] : memref<16xi32>
+ "test.side_effect_op"() {effects = [{effect = "allocate"}]} : () -> i32
+ affine.store %v, %tmp[%i] : memref<16xi32>
+ }
+ affine.for %j = 0 to 8 {
+ %v = affine.load %tmp[%j] : memref<16xi32>
+ affine.store %v, %out[%j] : memref<16xi32>
+ }
+ return
+}
+
// A source read and destination write through overlapping non-trivial views
// must not be interleaved by producer-consumer fusion.
diff --git a/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir b/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir
index 1314080e4864f..cd4a798ebf088 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-dependence-check.mlir
@@ -1,4 +1,5 @@
// RUN: mlir-opt -allow-unregistered-dialect %s -test-loop-fusion=test-loop-fusion-dependence-check -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -test-loop-fusion=test-loop-fusion-utilities -split-input-file | FileCheck %s --check-prefix=UTILITY
// -----
@@ -302,6 +303,58 @@ func.func @should_not_fuse_across_store_in_loop_at_depth1() {
// -----
+// The public utility must reject an unknown effect even when no MDG was built
+// by the caller first.
+
+// UTILITY-LABEL: func.func @utility_rejects_unknown_effect_in_candidate_loop
+// UTILITY: affine.for
+// UTILITY: "unknown.effect"() : () -> ()
+// UTILITY-NEXT: affine.store
+// UTILITY-NEXT: }
+// UTILITY-NEXT: affine.for
+func.func @utility_rejects_unknown_effect_in_candidate_loop(
+ %in: memref<8xi32>, %out: memref<8xi32>) {
+ %tmp = memref.alloc() : memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %v = affine.load %in[%i] : memref<8xi32>
+ "unknown.effect"() : () -> ()
+ affine.store %v, %tmp[%i] : memref<8xi32>
+ }
+ affine.for %j = 0 to 8 {
+ %v = affine.load %tmp[%j] : memref<8xi32>
+ affine.store %v, %out[%j] : memref<8xi32>
+ }
+ return
+}
+
+// -----
+
+// Allocation effects are not safe to clone into a sliced loop body, even when
+// the effect has no associated memref value.
+
+// UTILITY-LABEL: func.func @utility_rejects_allocate_effect_in_candidate_loop
+// UTILITY: affine.for
+// UTILITY: "test.side_effect_op"
+// UTILITY-NEXT: affine.store
+// UTILITY-NEXT: }
+// UTILITY-NEXT: affine.for
+func.func @utility_rejects_allocate_effect_in_candidate_loop(
+ %in: memref<8xi32>, %out: memref<8xi32>) {
+ %tmp = memref.alloc() : memref<8xi32>
+ affine.for %i = 0 to 8 {
+ %v = affine.load %in[%i] : memref<8xi32>
+ "test.side_effect_op"() {effects = [{effect = "allocate"}]} : () -> i32
+ affine.store %v, %tmp[%i] : memref<8xi32>
+ }
+ affine.for %j = 0 to 8 {
+ %v = affine.load %tmp[%j] : memref<8xi32>
+ affine.store %v, %out[%j] : memref<8xi32>
+ }
+ return
+}
+
+// -----
+
// CHECK-LABEL: func @should_not_fuse_across_ssa_value_def_at_depth1() {
func.func @should_not_fuse_across_ssa_value_def_at_depth1() {
%0 = memref.alloc() : memref<10x10xf32>
>From 0e43b4a1bff1280d32b0715740861fc58d2542d9 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Mon, 3 Aug 2026 10:28:21 +0800
Subject: [PATCH 19/23] [MLIR][Affine] Fix fusion effect regressions
---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 18 +++++++++---------
.../Dialect/Affine/Transforms/LoopFusion.cpp | 4 ++--
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 4 ++--
.../Dialect/Affine/loop-fusion-utilities.mlir | 12 +++++-------
4 files changed, 18 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index b4d0bfe970270..83ee9e468f158 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -115,11 +115,11 @@ static bool isRepresentableMemoryEffect(
if (isa<MemoryEffects::Allocate>(effect.getEffect())) {
Value value = effect.getValue();
// Allocation effects on a result are represented by the exact SSA edge
- // from the defining operation. Standard memref allocations use a
- // full-region allocation effect, so this check must be based on the
- // result identity rather than effectOnFullRegion.
- return effect.getResource() == SideEffects::DefaultResource::get() &&
- value && isa<BaseMemRefType>(value.getType()) &&
+ // from the defining operation. Standard memref allocations use both the
+ // default and automatic-allocation resources, so the result identity is
+ // the relevant representation rather than the resource or
+ // effectOnFullRegion.
+ return value && isa<BaseMemRefType>(value.getType()) &&
value.getDefiningOp() == op;
}
@@ -632,8 +632,8 @@ bool MemRefDependenceGraph::writesToLiveInOrEscapingMemrefs(unsigned id) const {
// Returns true iff there is an edge from node 'srcId' to node 'dstId' which
// is for 'value' if non-null, or for any value otherwise. Returns false
// otherwise.
-bool MemRefDependenceGraph::hasEdge(unsigned srcId, unsigned dstId,
- Value value, Edge::Kind kind) const {
+bool MemRefDependenceGraph::hasEdge(unsigned srcId, unsigned dstId, Value value,
+ Edge::Kind kind) const {
if (!outEdges.contains(srcId) || !inEdges.contains(dstId)) {
return false;
}
@@ -649,8 +649,8 @@ 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, Edge::Kind kind) {
+void MemRefDependenceGraph::addEdge(unsigned srcId, unsigned dstId, Value value,
+ Edge::Kind kind) {
if (!hasEdge(srcId, dstId, value, kind)) {
outEdges[srcId].push_back({dstId, kind, value});
inEdges[dstId].push_back({srcId, kind, value});
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index c8e3163e0ce24..431a335ffdeff 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -885,8 +885,8 @@ struct GreedyFusion {
return false;
if (edge.kind == MemRefDependenceGraph::Edge::Kind::SSA)
return edge.value == memref;
- return memref::isSameViewOrTrivialAlias(
- cast<MemrefValue>(edge.value), cast<MemrefValue>(memref));
+ return memref::isSameViewOrTrivialAlias(cast<MemrefValue>(edge.value),
+ cast<MemrefValue>(memref));
}))
return false;
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 0887e288fd384..aa749562513e6 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1588,8 +1588,8 @@ func.func @non_affine_effects_inside_candidate_loops(
// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @allocation_effect_inside_candidate_loop
// PRODUCER-CONSUMER-MAXIMAL: affine.for
// PRODUCER-CONSUMER-MAXIMAL: test.side_effect_op
-// PRODUCER-CONSUMER-MAXIMAL: }
-// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
func.func @allocation_effect_inside_candidate_loop(
%in: memref<16xi32>, %out: memref<16xi32>) {
%tmp = memref.alloc() : memref<16xi32>
diff --git a/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir b/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir
index 11435ad2d0203..6b9715280bab0 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir
@@ -9,14 +9,13 @@ func.func @slice_depth1_loop_nest() {
}
affine.for %i1 = 0 to 5 {
%1 = affine.load %0[%i1] : memref<100xf32>
- "prevent.dce"(%1) : (f32) -> ()
+ "test.side_effect_op"() : () -> i32
}
// CHECK: affine.for %[[IV0:.*]] = 0 to 5 {
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%[[IV0]]] : memref<100xf32>
// CHECK-NEXT: affine.load %{{.*}}[%[[IV0]]] : memref<100xf32>
- // CHECK-NEXT: "prevent.dce"(%{{.*}}) : (f32) -> ()
- // CHECK-NEXT: }
- // CHECK-NEXT: return
+ // CHECK: }
+ // CHECK: return
return
}
@@ -81,7 +80,7 @@ func.func @should_fuse_avoiding_dependence_cycle() {
affine.for %i1 = 0 to 10 {
affine.store %cf7, %a[%i1] : memref<10xf32>
%v1 = affine.load %c[%i1] : memref<10xf32>
- "prevent.dce"(%v1) : (f32) -> ()
+ "test.side_effect_op"() : () -> i32
}
affine.for %i2 = 0 to 10 {
%v2 = affine.load %b[%i2] : memref<10xf32>
@@ -100,8 +99,7 @@ func.func @should_fuse_avoiding_dependence_cycle() {
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
// CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
- // CHECK-NEXT: "prevent.dce"
- // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+ // CHECK: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
// CHECK-NEXT: }
// CHECK-NEXT: return
>From 2f618e1150f2435878582629411bf8696cfb8bac Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Mon, 3 Aug 2026 10:35:01 +0800
Subject: [PATCH 20/23] [MLIR][Affine] Keep utility fusion accesses visible
---
mlir/test/Dialect/Affine/loop-fusion-utilities.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir b/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir
index 6b9715280bab0..a169b3875dfaf 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-utilities.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -allow-unregistered-dialect -test-loop-fusion=test-loop-fusion-utilities -split-input-file -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -test-loop-fusion=test-loop-fusion-utilities -split-input-file | FileCheck %s
// CHECK-LABEL: func @slice_depth1_loop_nest() {
func.func @slice_depth1_loop_nest() {
>From a5ad395a8aa721192e4365bee46a25358ae97f19 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Wed, 5 Aug 2026 12:41:01 +0800
Subject: [PATCH 21/23] [MLIR][Affine] Guard loop sinking against unmodeled
effects
---
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 20 +++++++++
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 45 +++++++++++++++++++++
2 files changed, 65 insertions(+)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 90bc57e950cf1..1c9933a631313 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1462,6 +1462,26 @@ AffineForOp mlir::affine::sinkSequentialLoops(AffineForOp forOp) {
if (loops.size() < 2)
return forOp;
+ // Dependence components only model affine memory accesses. Do not reorder a
+ // loop nest containing any other memory effect, even if that effect can be
+ // represented by a different analysis, because its schedule is not covered
+ // by the dependence components used below.
+ LoopNestStateCollector collector;
+ collector.collect(loops[0]);
+ if (collector.hasUnmodeledMemoryEffects)
+ return forOp;
+ if (loops[0]
+ ->walk([&](Operation *op) {
+ if (isa<AffineForOp, AffineReadOpInterface,
+ AffineWriteOpInterface>(op) ||
+ op->mightHaveTrait<OpTrait::IsTerminator>())
+ return WalkResult::advance();
+ return isMemoryEffectFree(op) ? WalkResult::advance()
+ : WalkResult::interrupt();
+ })
+ .wasInterrupted())
+ return forOp;
+
// Gather dependence components for dependences between all ops in loop nest
// rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth].
unsigned maxLoopDepth = loops.size();
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index aa749562513e6..2ab31dde01df0 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1,6 +1,7 @@
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer}))' -split-input-file | FileCheck %s --check-prefix=PRODUCER-CONSUMER
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{compute-tolerance=0.0}))' -split-input-file | FileCheck %s --check-prefix=ZERO-TOLERANCE
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer maximal}))' -split-input-file | FileCheck %s --check-prefix=PRODUCER-CONSUMER-MAXIMAL
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer maximal}))' -split-input-file | FileCheck %s --check-prefix=NO-REORDER
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{maximal mode=sibling}))' -split-input-file | FileCheck %s --check-prefix=SIBLING-MAXIMAL
// All fusion: producer-consumer and sibling.
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s --check-prefix=ALL
@@ -1311,6 +1312,8 @@ func.func @subview_non_alias_memref_store(
// PRODUCER-CONSUMER-MAXIMAL: %[[A:.*]] = memref.alloc
// PRODUCER-CONSUMER-MAXIMAL: affine.store {{.*}}, %[[A]][
// PRODUCER-CONSUMER-MAXIMAL: memref.dealloc %[[A]]
+// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.store {{.*}}, %[[A]][
+// PRODUCER-CONSUMER-MAXIMAL: return
func.func @representable_free_between_loops(
%in: memref<32xf64>, %b: memref<32xf64>, %out: memref<32xf64>) {
%a = memref.alloc() : memref<32xf64>
@@ -1665,3 +1668,45 @@ func.func @overlapping_subviews_write_read(
}
return
}
+
+// -----
+
+// The destination nest has an affine loop-carried dependence that makes its
+// outer loop sequential. The non-affine recurrence is not part of the affine
+// dependence components, so sinking the sequential loop would change its
+// schedule even though fusion is later rejected.
+
+#linear = affine_map<(d0, d1) -> (d0 * 2 + d1)>
+
+// NO-REORDER-LABEL: func.func @non_affine_effect_prevents_loop_sinking(
+// NO-REORDER: affine.for %{{.*}} = 0 to 5 {
+// NO-REORDER-NEXT: affine.for %{{.*}} = 0 to 3 {
+// NO-REORDER: affine.for %{{.*}} = 0 to 4 {
+// NO-REORDER-NEXT: affine.for %{{.*}} = 0 to 3 {
+func.func @non_affine_effect_prevents_loop_sinking() {
+ %dep = memref.alloc() : memref<5x3xi64>
+ %state = memref.alloc() : memref<1xi64>
+ %out = memref.alloc() : memref<4x3xi64>
+ %c0 = arith.constant 0 : index
+ %zero = arith.constant 0 : i64
+ memref.store %zero, %state[%c0] : memref<1xi64>
+
+ affine.for %p = 0 to 5 {
+ affine.for %q = 0 to 3 {
+ affine.store %zero, %dep[%p, %q] : memref<5x3xi64>
+ }
+ }
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 3 {
+ %v = affine.load %dep[%i, %j] : memref<5x3xi64>
+ affine.store %v, %dep[%i + 1, %j] : memref<5x3xi64>
+ %linear = affine.apply #linear(%i, %j)
+ %linear_i64 = arith.index_cast %linear : index to i64
+ %old = memref.load %state[%c0] : memref<1xi64>
+ %next = arith.addi %old, %linear_i64 : i64
+ memref.store %next, %state[%c0] : memref<1xi64>
+ affine.store %v, %out[%i, %j] : memref<4x3xi64>
+ }
+ }
+ return
+}
>From 47644d77b1c7a517060d282d028709d432aea5ee Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Wed, 5 Aug 2026 12:45:11 +0800
Subject: [PATCH 22/23] [MLIR][Affine] Format loop sinking guard
---
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 1c9933a631313..3f8ee61c894eb 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1472,8 +1472,8 @@ AffineForOp mlir::affine::sinkSequentialLoops(AffineForOp forOp) {
return forOp;
if (loops[0]
->walk([&](Operation *op) {
- if (isa<AffineForOp, AffineReadOpInterface,
- AffineWriteOpInterface>(op) ||
+ if (isa<AffineForOp, AffineReadOpInterface, AffineWriteOpInterface>(
+ op) ||
op->mightHaveTrait<OpTrait::IsTerminator>())
return WalkResult::advance();
return isMemoryEffectFree(op) ? WalkResult::advance()
>From 44849b1dd6f8ee291ab422aa987782e667e50395 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Wed, 5 Aug 2026 15:05:55 +0800
Subject: [PATCH 23/23] [MLIR][Affine] Guard sinking against view aliases
---
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 33 +++++++++++++++++++++
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 32 +++++++++++++++++++-
2 files changed, 64 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 3f8ee61c894eb..c2ba421dc4e9d 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -18,6 +18,7 @@
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/IntegerSet.h"
@@ -1453,6 +1454,31 @@ unsigned mlir::affine::permuteLoops(ArrayRef<AffineForOp> input,
return invPermMap[0].second;
}
+// Returns true when affine accesses use different views of the same storage
+// and at least one access writes. The affine dependence analysis cannot
+// compare access maps expressed in different view coordinate systems, so a
+// schedule-changing transform must conservatively preserve this ordering.
+static bool hasConflictingAffineStorageAccesses(
+ ArrayRef<Operation *> affineAccesses) {
+ for (unsigned i = 0, e = affineAccesses.size(); i < e; ++i) {
+ MemRefAccess firstAccess(affineAccesses[i]);
+ Value firstStorage = memref::skipViewLikeOps(
+ cast<MemrefValue>(firstAccess.memref));
+ for (unsigned j = i + 1; j < e; ++j) {
+ MemRefAccess secondAccess(affineAccesses[j]);
+ if (firstAccess.memref == secondAccess.memref)
+ continue;
+ Value secondStorage = memref::skipViewLikeOps(
+ cast<MemrefValue>(secondAccess.memref));
+ if (firstStorage != secondStorage)
+ continue;
+ if (firstAccess.isStore() || secondAccess.isStore())
+ return true;
+ }
+ }
+ return false;
+}
+
// Sinks all sequential loops to the innermost levels (while preserving
// relative order among them) and moves all parallel loops to the
// outermost (while again preserving relative order among them).
@@ -1470,6 +1496,13 @@ AffineForOp mlir::affine::sinkSequentialLoops(AffineForOp forOp) {
collector.collect(loops[0]);
if (collector.hasUnmodeledMemoryEffects)
return forOp;
+ SmallVector<Operation *, 8> affineAccesses;
+ for (Operation *op : collector.loadOpInsts)
+ affineAccesses.push_back(op);
+ for (Operation *op : collector.storeOpInsts)
+ affineAccesses.push_back(op);
+ if (hasConflictingAffineStorageAccesses(affineAccesses))
+ return forOp;
if (loops[0]
->walk([&](Operation *op) {
if (isa<AffineForOp, AffineReadOpInterface, AffineWriteOpInterface>(
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 2ab31dde01df0..b747e8e8b7b8b 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -1689,6 +1689,7 @@ func.func @non_affine_effect_prevents_loop_sinking() {
%out = memref.alloc() : memref<4x3xi64>
%c0 = arith.constant 0 : index
%zero = arith.constant 0 : i64
+ %ten = arith.constant 10 : i64
memref.store %zero, %state[%c0] : memref<1xi64>
affine.for %p = 0 to 5 {
@@ -1703,10 +1704,39 @@ func.func @non_affine_effect_prevents_loop_sinking() {
%linear = affine.apply #linear(%i, %j)
%linear_i64 = arith.index_cast %linear : index to i64
%old = memref.load %state[%c0] : memref<1xi64>
- %next = arith.addi %old, %linear_i64 : i64
+ %scaled = arith.muli %old, %ten : i64
+ %next = arith.addi %scaled, %linear_i64 : i64
memref.store %next, %state[%c0] : memref<1xi64>
affine.store %v, %out[%i, %j] : memref<4x3xi64>
}
}
return
}
+
+// Affine accesses through different views of the same storage are not
+// represented by the raw-view dependence components. They must still prevent
+// sinking when another affine dependence would otherwise select interchange.
+
+// NO-REORDER-LABEL: func.func @affine_alias_prevents_loop_sinking(
+// NO-REORDER: affine.for %{{.*}} = 0 to 1 {
+// NO-REORDER: affine.store
+// NO-REORDER: affine.for %{{.*}} = 1 to 5 {
+// NO-REORDER-NEXT: affine.for %{{.*}} = 1 to 4 {
+func.func @affine_alias_prevents_loop_sinking(
+ %seed: memref<1xi64>, %dep: memref<6x4xi64>, %root: memref<5x5xi64>) {
+ %view = memref.cast %root : memref<5x5xi64> to memref<?x?xi64>
+ %zero = arith.constant 0 : i64
+ affine.for %s = 0 to 1 {
+ affine.store %zero, %seed[%s] : memref<1xi64>
+ }
+ affine.for %i = 1 to 5 {
+ affine.for %j = 1 to 4 {
+ %v = affine.load %dep[%i, %j] : memref<6x4xi64>
+ affine.store %v, %dep[%i + 1, %j] : memref<6x4xi64>
+ %old = affine.load %root[%i - 1, %j + 1] : memref<5x5xi64>
+ %next = arith.addi %old, %v : i64
+ affine.store %next, %view[%i, %j] : memref<?x?xi64>
+ }
+ }
+ return
+}
More information about the Mlir-commits
mailing list