[Mlir-commits] [mlir] [MLIR][Affine] Avoid crash on non-affine memory effects (PR #205865)

[llvmlistbot at llvm.org]

https://github.com/mygitljf updated https://github.com/llvm/llvm-project/pull/205865

>From da3c4ce0e725f5da156340d9802d2b918c7acca8 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 26 Jun 2026 00:57:17 +0000
Subject: [PATCH 1/2] [MLIR][Affine] Handle memory effects

---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 28 +++++++++++-----------
 mlir/test/Dialect/Affine/loop-fusion.mlir  | 25 ++++++++++++++++++-
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index cac305df8ba75..4aa419e5e81bc 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -243,17 +243,18 @@ 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");
+static bool mayAccessMemRef(Operation *op, Value memref) {
+  if (auto affineRead = dyn_cast<AffineReadOpInterface>(op))
+    return affineRead.getMemRef() == memref;
+  if (auto affineWrite = dyn_cast<AffineWriteOpInterface>(op))
+    return affineWrite.getMemRef() == memref;
+  if (hasEffect<MemoryEffects::Read>(op, memref) ||
+      hasEffect<MemoryEffects::Write>(op, memref))
+    return true;
+  if (isa<MemoryEffectOpInterface>(op) ||
+      op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
+    return false;
+  return llvm::is_contained(op->getOperands(), memref);
 }
 
 /// Returns true if there may be a dependence on `memref` from srcNode's
@@ -276,11 +277,10 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
   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);
+               return mayAccessMemRef(dstOp, memref);
              }) != dstMemOps.end();
     });
   };
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 1ea42517988c3..3eb3e09bb61fb 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -13,6 +13,30 @@
 
 // -----
 
+// CHECK-LABEL: func @non_affine_memory_effect_in_loop_does_not_crash(
+func.func @non_affine_memory_effect_in_loop_does_not_crash(%arg0: memref<8xf32>) {
+  %0 = builtin.unrealized_conversion_cast %arg0 : memref<8xf32> to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+  %c0 = arith.constant 0 : index
+  affine.for %i = 0 to 8 {
+    %1 = builtin.unrealized_conversion_cast %i : index to i64
+    %2 = llvm.extractvalue %0[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %3 = llvm.getelementptr %2[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+    %4 = llvm.load %3 : !llvm.ptr -> f32
+    affine.store %4, %arg0[%c0] : memref<8xf32>
+  }
+  affine.for %i = 0 to 8 {
+    %1 = builtin.unrealized_conversion_cast %i : index to i64
+    %2 = llvm.extractvalue %0[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %3 = llvm.getelementptr %2[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+    %4 = llvm.load %3 : !llvm.ptr -> f32
+    affine.store %4, %arg0[%c0] : memref<8xf32>
+  }
+  // CHECK: return
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func @should_fuse_raw_dep_for_locality() {
 func.func @should_fuse_raw_dep_for_locality() {
   %m = memref.alloc() : memref<10xf32>
@@ -1575,4 +1599,3 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
 }
 
 // Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
-

>From a1636e345275746ac2838334df47421863dba84a Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 26 Jun 2026 12:13:17 +0000
Subject: [PATCH 2/2] [MLIR][Affine] Document memref access check

---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 4aa419e5e81bc..9620c9d4693cc 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -243,6 +243,10 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
   return &node;
 }
 
+/// Returns whether `op` may read from or write to `memref`.
+/// Memory effects are used when available; recursively-effectful ops are
+/// handled through their nested operations. Otherwise, conservatively treat a
+/// memref operand as a possible access, which may add dependence edges.
 static bool mayAccessMemRef(Operation *op, Value memref) {
   if (auto affineRead = dyn_cast<AffineReadOpInterface>(op))
     return affineRead.getMemRef() == memref;