[Mlir-commits] [mlir] [mlir][affine] Block producer-consumer fusion across opaque memref effects (PR #225250)
Demetrios Chiuratto Agourakis
llvmlistbot at llvm.org
Sat Oct 3 08:29:28 PDT 2026
https://github.com/agourakis82 updated https://github.com/llvm/llvm-project/pull/225250
>From 5e3fbb2c020c7f14353d2b25d35bc4cb74cf579b Mon Sep 17 00:00:00 2001
From: Demetrios Chiuratto Agourakis <demetrios at agourakis.med.br>
Date: Tue, 22 Sep 2026 01:28:35 +0000
Subject: [PATCH] [mlir][affine] Block PC fusion on opaque producer effects
Producer-consumer fusion can reorder opaque memory effects on a shared
memref across consumer iterations. Add hasOpaqueEffectOnPCMemrefs as an
early legality gate for non-affine producer effects and non-dereferencing
users nested in the producer.
Restrict the user check to the producer nest so intervening reads,
returned memrefs, and other users outside the producer continue to be
handled by the existing dependence graph. Preserve ordinary affine
fusion, including calls on unrelated memrefs.
This correctness gate is separate from the generic dependence handling
already merged in #224657; retain that upstream implementation unchanged.
Fixes #211591
---
.../Dialect/Affine/Transforms/LoopFusion.cpp | 60 ++++++++++++++++++
.../Dialect/Affine/loop-fusion-opaque-pc.mlir | 63 +++++++++++++++++++
2 files changed, 123 insertions(+)
create mode 100644 mlir/test/Dialect/Affine/loop-fusion-opaque-pc.mlir
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 77eda4062ec0c1..a113fab921855c 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -22,6 +22,7 @@
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h"
@@ -225,6 +226,58 @@ static void gatherEscapingMemrefs(unsigned id, const MemRefDependenceGraph &mdg,
}
}
+/// Returns true if the producer nest has opaque effects on any
+/// producer-consumer memref. An effect is treated as opaque when:
+/// 1) the producer contains non-affine memory effects on the memref, or
+/// 2) the memref has a non-dereferencing user in the producer nest
+/// (for example `func.call`), which can observe/reorder whole-buffer
+/// effects across consumer iterations.
+///
+/// Note: mere "escaping" via block arguments (used for private-memref
+/// decisions) is NOT sufficient to block fusion.
+static bool
+hasOpaqueEffectOnPCMemrefs(const MemRefDependenceGraph::Node &srcNode,
+ const DenseSet<Value> &producerConsumerMemrefs) {
+ if (producerConsumerMemrefs.empty())
+ return false;
+
+ // Non-affine memory ops collected on the producer node.
+ auto touchesPCMemref = [&](Operation *op) {
+ for (Value memref : producerConsumerMemrefs) {
+ if (hasEffect<MemoryEffects::Read>(op, memref) ||
+ hasEffect<MemoryEffects::Write>(op, memref) ||
+ hasEffect<MemoryEffects::Free>(op, memref))
+ return true;
+ // Ops without an effect model (e.g. func.call) were conservatively
+ // recorded because they take a memref operand.
+ if (!isa<MemoryEffectOpInterface>(op) &&
+ !op->hasTrait<OpTrait::HasRecursiveMemoryEffects>() &&
+ llvm::is_contained(op->getOperands(), memref))
+ return true;
+ }
+ return false;
+ };
+
+ if (llvm::any_of(srcNode.memrefLoads, touchesPCMemref) ||
+ llvm::any_of(srcNode.memrefStores, touchesPCMemref) ||
+ llvm::any_of(srcNode.memrefFrees, touchesPCMemref))
+ return true;
+
+ // Only users in the producer are reordered by inserting it into the
+ // consumer. Other users are handled by the dependence graph and must not
+ // blanket-block fusion (for example, intervening reads or returned memrefs).
+ for (Value memref : producerConsumerMemrefs) {
+ bool hasOpaqueUser = llvm::any_of(memref.getUsers(), [&](Operation *user) {
+ if (!srcNode.op->isProperAncestor(user))
+ return false;
+ return !isa<AffineMapAccessInterface>(*user);
+ });
+ if (hasOpaqueUser)
+ 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).
@@ -951,6 +1004,13 @@ struct GreedyFusion {
DenseSet<Value> srcEscapingMemRefs;
gatherEscapingMemrefs(srcNode->id, *mdg, srcEscapingMemRefs);
+ // Opaque/non-affine effects on a PC memref are a hard legality
+ // barrier: fusion can reorder them across consumer iterations.
+ if (hasOpaqueEffectOnPCMemrefs(*srcNode, producerConsumerMemrefs)) {
+ LDBG() << "Skipping fusion due to opaque effects on PC memref";
+ continue;
+ }
+
// Compute an operation list insertion point for the fused loop
// nest which preserves dependences.
Operation *fusedLoopInsPoint =
diff --git a/mlir/test/Dialect/Affine/loop-fusion-opaque-pc.mlir b/mlir/test/Dialect/Affine/loop-fusion-opaque-pc.mlir
new file mode 100644
index 00000000000000..66eeeffed5fc73
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-fusion-opaque-pc.mlir
@@ -0,0 +1,63 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer maximal}))' | FileCheck %s
+
+// Opaque func.call on a producer-consumer memref must block fusion.
+// CHECK-LABEL: func @fusion_blocked_by_opaque_call_on_pc_memref
+func.func private @opaque(memref<32xf64>)
+func.func @fusion_blocked_by_opaque_call_on_pc_memref(
+ %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>
+ func.call @opaque(%comm) : (memref<32xf64>) -> ()
+ }
+ 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>
+ }
+ // CHECK: affine.for
+ // CHECK: func.call @opaque
+ // CHECK: affine.for
+ return
+}
+
+// CHECK-LABEL: func @fusion_still_happens_without_opaque_effect
+func.func @fusion_still_happens_without_opaque_effect(
+ %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>
+ }
+ 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>
+ }
+ // CHECK: affine.for
+ // CHECK-NOT: affine.for
+ return
+}
+
+// A call on an unrelated memref must not blanket-block PC fusion of %comm.
+// CHECK-LABEL: func @call_on_unrelated_memref_does_not_block_pc_fusion
+func.func @call_on_unrelated_memref_does_not_block_pc_fusion(
+ %in: memref<32xf64>, %comm: memref<32xf64>,
+ %scratch: 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>
+ func.call @opaque(%scratch) : (memref<32xf64>) -> ()
+ }
+ 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>
+ }
+ // CHECK: affine.for
+ // CHECK: func.call @opaque
+ // CHECK-NOT: affine.for
+ return
+}
More information about the Mlir-commits
mailing list