[Mlir-commits] [mlir] [MLIR][Affine] Handle cast aliases in fusion dependencies (PR #213413)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 31 23:13:44 PDT 2026
https://github.com/1sgtpepper updated https://github.com/llvm/llvm-project/pull/213413
>From 8de53cd7a3c3064ea89e5c454b869e6a660410e2 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 13:59:03 +0800
Subject: [PATCH 1/2] [MLIR][Affine] Handle cast aliases in fusion dependencies
---
.../Dialect/Affine/Analysis/CMakeLists.txt | 1 +
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 85 +++++++++++--------
mlir/test/Dialect/Affine/loop-fusion-4.mlir | 41 ++++++++-
3 files changed, 91 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt b/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt
index 3a1996349dbed..6c52f7751f1f0 100644
--- a/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Analysis/CMakeLists.txt
@@ -18,6 +18,7 @@ add_mlir_dialect_library(MLIRAffineAnalysis
MLIRControlFlowInterfaces
MLIRDialectUtils
MLIRInferTypeOpInterface
+ MLIRMemRefUtils
MLIRSideEffectInterfaces
MLIRPresburger
MLIRSCFDialect
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 321c8e34d907c..11e412f8d8704 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -19,6 +19,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IntegerSet.h"
#include "llvm/ADT/SetVector.h"
@@ -38,6 +39,49 @@ using llvm::SmallDenseMap;
using Node = MemRefDependenceGraph::Node;
+/// Returns the values that `op` may have a memref effect of type `EffectTys`
+/// on, not considering recursive effects. An op with unknown memory effects
+/// (e.g. a call to an external function without a memory-effect interface) is
+/// conservatively assumed to affect all its memref operands. Fully aliasing
+/// views are canonicalized so the MDG uses one key for the view and its source.
+template <typename... EffectTys>
+static void getMayEffectedValues(Operation *op,
+ SmallVectorImpl<Value> &values) {
+ auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!memOp) {
+ if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
+ // No effects.
+ return;
+ // Memref operands have to be considered as being affected.
+ for (Value operand : op->getOperands()) {
+ if (isa<MemRefType>(operand.getType()))
+ values.push_back(memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(operand)));
+ }
+ return;
+ }
+ SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
+ memOp.getEffects(effects);
+ for (auto &effect : effects) {
+ Value effectVal = effect.getValue();
+ if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
+ isa<MemRefType>(effectVal.getType()))
+ values.push_back(memref::skipFullyAliasingOperations(
+ cast<MemrefValue>(effectVal)));
+ };
+}
+
+/// Returns true if `op` may have a memory effect of type `EffectTys` on
+/// `memref`, i.e., whether `memref` is among the values returned by
+/// `getMayEffectedValues` for `op`.
+template <typename... EffectTys>
+static bool mayHaveEffect(Operation *op, Value memref) {
+ SmallVector<Value> values;
+ getMayEffectedValues<EffectTys...>(op, values);
+ return llvm::is_contained(
+ values, memref::skipFullyAliasingOperations(cast<MemrefValue>(memref)));
+}
+
// LoopNestStateCollector walks loop nests and collects load and store
// operations, and whether or not a region holding op other than ForOp and IfOp
// was encountered in the loop nest.
@@ -83,7 +127,7 @@ unsigned Node::getLoadOpCount(Value memref) const {
if (auto affineLoad = dyn_cast<AffineReadOpInterface>(loadOp)) {
if (memref == affineLoad.getMemRef())
++loadOpCount;
- } else if (hasEffect<MemoryEffects::Read>(loadOp, memref)) {
+ } else if (mayHaveEffect<MemoryEffects::Read>(loadOp, memref)) {
++loadOpCount;
}
}
@@ -98,8 +142,7 @@ unsigned Node::getStoreOpCount(Value memref) const {
if (auto affineStore = dyn_cast<AffineWriteOpInterface>(storeOp)) {
if (memref == affineStore.getMemRef())
++storeOpCount;
- } else if (hasEffect<MemoryEffects::Write>(const_cast<Operation *>(storeOp),
- memref)) {
+ } else if (mayHaveEffect<MemoryEffects::Write>(storeOp, memref)) {
++storeOpCount;
}
}
@@ -114,7 +157,7 @@ unsigned Node::hasStore(Value memref) const {
if (auto affineStore = dyn_cast<AffineWriteOpInterface>(storeOp)) {
if (memref == affineStore.getMemRef())
return true;
- } else if (hasEffect<MemoryEffects::Write>(storeOp, memref)) {
+ } else if (mayHaveEffect<MemoryEffects::Write>(storeOp, memref)) {
return true;
}
return false;
@@ -123,7 +166,7 @@ unsigned Node::hasStore(Value memref) const {
unsigned Node::hasFree(Value memref) const {
return llvm::any_of(memrefFrees, [&](Operation *freeOp) {
- return hasEffect<MemoryEffects::Free>(freeOp, memref);
+ return mayHaveEffect<MemoryEffects::Free>(freeOp, memref);
});
}
@@ -160,32 +203,6 @@ void Node::getLoadAndStoreMemrefSet(
}
}
-/// Returns the values that this op has a memref effect of type `EffectTys` on,
-/// not considering recursive effects.
-template <typename... EffectTys>
-static void getEffectedValues(Operation *op, SmallVectorImpl<Value> &values) {
- auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
- if (!memOp) {
- if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
- // No effects.
- return;
- // Memref operands have to be considered as being affected.
- for (Value operand : op->getOperands()) {
- if (isa<MemRefType>(operand.getType()))
- values.push_back(operand);
- }
- return;
- }
- SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
- memOp.getEffects(effects);
- for (auto &effect : effects) {
- Value effectVal = effect.getValue();
- if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
- isa<MemRefType>(effectVal.getType()))
- values.push_back(effectVal);
- };
-}
-
/// Add `op` to MDG creating a new node and adding its memory accesses (affine
/// or non-affine to memrefAccesses (memref -> list of nodes with accesses) map.
static Node *
@@ -210,7 +227,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefLoads) {
SmallVector<Value> effectedValues;
- getEffectedValues<MemoryEffects::Read>(op, effectedValues);
+ getMayEffectedValues<MemoryEffects::Read>(op, effectedValues);
if (llvm::any_of(((ValueRange)effectedValues).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
// We do not know the interaction here.
@@ -221,7 +238,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefStores) {
SmallVector<Value> effectedValues;
- getEffectedValues<MemoryEffects::Write>(op, effectedValues);
+ getMayEffectedValues<MemoryEffects::Write>(op, effectedValues);
if (llvm::any_of((ValueRange(effectedValues)).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
return nullptr;
@@ -231,7 +248,7 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
}
for (Operation *op : collector.memrefFrees) {
SmallVector<Value> effectedValues;
- getEffectedValues<MemoryEffects::Free>(op, effectedValues);
+ getMayEffectedValues<MemoryEffects::Free>(op, effectedValues);
if (llvm::any_of((ValueRange(effectedValues)).getTypes(),
[](Type type) { return !isa<MemRefType>(type); }))
return nullptr;
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index cf530016c201a..980e68e6b9cfa 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -583,11 +583,15 @@ func.func @zero_tolerance(%arg0: memref<65536xcomplex<f64>>, %arg1: memref<30x13
affine.store %18, %2[%arg2] : memref<131072xi128>
affine.store %13, %1[%arg2] : memref<131072xi1>
}
- // The next two nests are fused.
+ // The next nest cannot fuse with the one following it across the opaque
+ // external call below, which may write to its memref operand in place.
// ZERO-TOLERANCE: affine.for %{{.*}} = 0 to 30
// ZERO-TOLERANCE-NEXT: affine.for %{{.*}} = 0 to 131072
// ZERO-TOLERANCE: func.call @__external_reduce_barrett
// ZERO-TOLERANCE: affine.store
+ // ZERO-TOLERANCE: call @__external_levelwise_forward_ntt
+ // ZERO-TOLERANCE-NEXT: affine.for %{{.*}} = 0 to 30
+ // ZERO-TOLERANCE-NEXT: affine.for %{{.*}} = 0 to 131072
// ZERO-TOLERANCE: affine.load
// ZERO-TOLERANCE-NEXT: affine.store
affine.for %arg2 = 0 to 30 {
@@ -611,9 +615,14 @@ func.func @zero_tolerance(%arg0: memref<65536xcomplex<f64>>, %arg1: memref<30x13
affine.store %7, %arg1[%arg2, %arg3] : memref<30x131072xi64>
}
}
- // Under maximal fusion, just one nest.
+ // Under maximal fusion, the first two nests fuse, but the last nest cannot
+ // fuse into them across the opaque external call, which may write to its
+ // memref operand in place.
// PRODUCER-CONSUMER-MAXIMAL: affine.for %{{.*}} = 0 to 30
// PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 131072
+ // PRODUCER-CONSUMER-MAXIMAL: call @__external_levelwise_forward_ntt
+ // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 30
+ // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 131072
// PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for %{{.*}}
memref.dealloc %2 : memref<131072xi128>
memref.dealloc %1 : memref<131072xi1>
@@ -884,3 +893,31 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
}
return %alloc : memref<1024x8192xf32>
}
+
+// -----
+
+// The external call receives a fully aliasing cast of the producer's memref.
+// Fusion must preserve the call between the producer and consumer loops.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @cast_alias_external_call
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+// PRODUCER-CONSUMER-MAXIMAL: memref.cast
+// PRODUCER-CONSUMER-MAXIMAL: call @escape
+// PRODUCER-CONSUMER-MAXIMAL: affine.for
+func.func @cast_alias_external_call(
+ %in: memref<32xf64>, %comm: memref<32xf64>, %out: memref<32xf64>) {
+ affine.for %i = 0 to 16 {
+ %a = affine.load %in[%i] : memref<32xf64>
+ %b = arith.addf %a, %a : f64
+ affine.store %b, %comm[%i] : memref<32xf64>
+ }
+ %view = memref.cast %comm : memref<32xf64> to memref<?xf64>
+ func.call @escape(%view) : (memref<?xf64>) -> ()
+ affine.for %j = 0 to 16 {
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ return
+}
+func.func private @escape(memref<?xf64>)
>From ba43cc3744d3fee08e720ca8fa97ace576c368b7 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sat, 1 Aug 2026 14:13:28 +0800
Subject: [PATCH 2/2] 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)));
};
}
More information about the Mlir-commits
mailing list