[Mlir-commits] [mlir] [mlir][affine] Block producer-consumer fusion across opaque memref effects (PR #225250)
Demetrios Chiuratto Agourakis
llvmlistbot at llvm.org
Sun Oct 4 03:31:44 PDT 2026
https://github.com/agourakis82 updated https://github.com/llvm/llvm-project/pull/225250
>From 962d58cc3222e83fa24c9710c621ed3516f3334a 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 across opaque memref 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 effects and non-dereferencing users nested in either
candidate loop.
Restrict the user check to the candidate nests 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.
Treat value-less read/write/free effects as opaque accesses to any PC
memref, and keep value-less allocation alone from blocking fusion.
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 | 68 +++++
.../Dialect/Affine/loop-fusion-opaque-pc.mlir | 237 ++++++++++++++++++
2 files changed, 305 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 77eda4062ec0c..4f9c857d5f433 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,65 @@ static void gatherEscapingMemrefs(unsigned id, const MemRefDependenceGraph &mdg,
}
}
+/// Returns true if the candidate loop nest has opaque effects on any
+/// producer-consumer memref. An effect is treated as opaque when:
+/// 1) the candidate contains non-affine memory effects on the memref, or
+/// 2) the memref has a non-dereferencing user in the candidate loop 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 &node,
+ const DenseSet<Value> &producerConsumerMemrefs) {
+ if (producerConsumerMemrefs.empty())
+ return false;
+
+ // Non-affine memory ops collected on the candidate node.
+ auto touchesPCMemref = [&](Operation *op) {
+ if (auto effectOp = dyn_cast<MemoryEffectOpInterface>(op)) {
+ SmallVector<MemoryEffects::EffectInstance> effects;
+ effectOp.getEffects(effects);
+ return llvm::any_of(effects, [&](const auto &effect) {
+ if (!isa<MemoryEffects::Read, MemoryEffects::Write,
+ MemoryEffects::Free>(effect.getEffect()))
+ return false;
+ // Value-less effects may access any PC memref, just as in the
+ // dependence analysis. They have no proven affine footprint.
+ return !effect.getValue() ||
+ producerConsumerMemrefs.contains(effect.getValue());
+ });
+ }
+ // Ops without an effect model (e.g. func.call) were conservatively
+ // recorded because they take a memref operand.
+ if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
+ return false;
+ return llvm::any_of(op->getOperands(), [&](Value operand) {
+ return producerConsumerMemrefs.contains(operand);
+ });
+ };
+
+ if (llvm::any_of(node.memrefLoads, touchesPCMemref) ||
+ llvm::any_of(node.memrefStores, touchesPCMemref) ||
+ llvm::any_of(node.memrefFrees, touchesPCMemref))
+ return true;
+
+ // Fusion reorders operations in the candidate nests relative to each other.
+ // Users outside the candidates are handled by the dependence graph and must
+ // not blanket-block fusion (for example, intervening reads or returns).
+ for (Value memref : producerConsumerMemrefs) {
+ bool hasOpaqueUser = llvm::any_of(memref.getUsers(), [&](Operation *user) {
+ if (!node.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 +1011,14 @@ 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) ||
+ hasOpaqueEffectOnPCMemrefs(*dstNode, 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 0000000000000..8dbaf40a4b046
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-fusion-opaque-pc.mlir
@@ -0,0 +1,237 @@
+// 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
+}
+
+// Value-less read effects may access any PC memref.
+// CHECK-LABEL: func @fusion_blocked_by_valueless_read
+func.func @fusion_blocked_by_valueless_read(
+ %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>
+ %unused = "test.side_effect_op"() {effects = [{effect = "read"}]} : () -> i32
+ }
+ 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: "test.side_effect_op"() {effects = [{effect = "read"}]}
+ // CHECK: affine.for
+ return
+}
+
+
+// Value-less write effects may access any PC memref.
+// CHECK-LABEL: func @fusion_blocked_by_valueless_write
+func.func @fusion_blocked_by_valueless_write(
+ %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>
+ %unused = "test.side_effect_op"() {effects = [{effect = "write"}]} : () -> i32
+ }
+ 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: "test.side_effect_op"() {effects = [{effect = "write"}]}
+ // CHECK: affine.for
+ return
+}
+
+
+// Value-less free effects may access any PC memref.
+// CHECK-LABEL: func @fusion_blocked_by_valueless_free
+func.func @fusion_blocked_by_valueless_free(
+ %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>
+ %unused = "test.side_effect_op"() {effects = [{effect = "free"}]} : () -> i32
+ }
+ 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: "test.side_effect_op"() {effects = [{effect = "free"}]}
+ // CHECK: affine.for
+ return
+}
+
+
+// Allocation alone is not a read/write/free effect on PC memory.
+// CHECK-LABEL: func @valueless_allocate_does_not_block_fusion
+func.func @valueless_allocate_does_not_block_fusion(
+ %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>
+ %unused = "test.side_effect_op"() {effects = [{effect = "allocate"}]} : () -> i32
+ }
+ 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: "test.side_effect_op"() {effects = [{effect = "allocate"}]}
+ // CHECK-NOT: affine.for
+ return
+}
+
+// Opaque call effects in the consumer also require the full producer to finish.
+// CHECK-LABEL: func @fusion_blocked_by_opaque_consumer_call
+func.func @fusion_blocked_by_opaque_consumer_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>
+ }
+ affine.for %j = 0 to 16 {
+ func.call @opaque(%comm) : (memref<32xf64>) -> ()
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ // CHECK: affine.for
+ // CHECK: affine.for
+ // CHECK: func.call @opaque
+ return
+}
+
+
+// Opaque read effects in the consumer also require the full producer to finish.
+// CHECK-LABEL: func @fusion_blocked_by_opaque_consumer_read
+func.func @fusion_blocked_by_opaque_consumer_read(
+ %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 {
+ %unused = "test.side_effect_op"() {effects = [{effect = "read"}]} : () -> i32
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ // CHECK: affine.for
+ // CHECK: affine.for
+ // CHECK: "test.side_effect_op"() {effects = [{effect = "read"}]}
+ return
+}
+
+
+// Opaque write effects in the consumer also require the full producer to finish.
+// CHECK-LABEL: func @fusion_blocked_by_opaque_consumer_write
+func.func @fusion_blocked_by_opaque_consumer_write(
+ %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 {
+ %unused = "test.side_effect_op"() {effects = [{effect = "write"}]} : () -> i32
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ // CHECK: affine.for
+ // CHECK: affine.for
+ // CHECK: "test.side_effect_op"() {effects = [{effect = "write"}]}
+ return
+}
+
+
+// Opaque free effects in the consumer also require the full producer to finish.
+// CHECK-LABEL: func @fusion_blocked_by_opaque_consumer_free
+func.func @fusion_blocked_by_opaque_consumer_free(
+ %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 {
+ %unused = "test.side_effect_op"() {effects = [{effect = "free"}]} : () -> i32
+ %c = affine.load %comm[%j] : memref<32xf64>
+ %d = arith.addf %c, %c : f64
+ affine.store %d, %out[%j] : memref<32xf64>
+ }
+ // CHECK: affine.for
+ // CHECK: affine.for
+ // CHECK: "test.side_effect_op"() {effects = [{effect = "free"}]}
+ return
+}
More information about the Mlir-commits
mailing list