[Mlir-commits] [mlir] [mlir][Affine] Handle memref.copy in loop fusion dependence checks (PR #203811)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Peter Chen J. (peter941221)

<details>
<summary>Changes</summary>

Fixes #<!-- -->203762

`mayDependence()` treated every non-affine memory op as if it had one memref operand. `memref.copy` lands in both the load and store effect buckets during loop-fusion dependence checks, so that path called `getMemRef(memref.copy)` and hit `llvm_unreachable("unexpected op")`.

This patch removes that single-memref assumption from the non-affine path. It checks whether each candidate op reads or writes the specific memref under analysis instead, so `memref.copy` works without changing the existing dependence walk for affine load/store pairs.

The regression covers the reported crash shape and keeps both memrefs carried by `memref.copy` in play: the source memref feeds a later `affine.load`, and the destination memref feeds a later `affine.store`.

Validation
build-mlir/bin/llvm-lit -sv -j1 mlir/test/Dialect/Affine/loop-fusion.mlir
build-mlir/bin/mlir-opt --affine-loop-fusion repro.mlir -o /dev/null

AI tool disclosure: I used OpenAI Codex to help inspect the failing code path, draft the patch, and draft the PR text. I reviewed and tested the final change myself before sending it for review.

---
Full diff: https://github.com/llvm/llvm-project/pull/203811.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Analysis/Utils.cpp (+6-16) 
- (modified) mlir/test/Dialect/Affine/loop-fusion.mlir (+27-1) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index ebe932a14694a..a63f19834832a 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -243,19 +243,6 @@ 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 there may be a dependence on `memref` from srcNode's
 /// memory ops to dstNode's memory ops, while using the affine memory
 /// dependence analysis checks. The method assumes that there is at least one
@@ -275,12 +262,15 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
   // destination read/write one; all expected to be memref/affine load/store.
   auto hasNonAffineDep = [&](ArrayRef<Operation *> srcMemOps,
                              ArrayRef<Operation *> dstMemOps) {
+    auto accessesMemref = [&](Operation *op) {
+      return hasEffect<MemoryEffects::Read>(op, memref) ||
+             hasEffect<MemoryEffects::Write>(op, memref);
+    };
     return llvm::any_of(srcMemOps, [&](Operation *srcOp) {
-      Value srcMemref = getMemRef(srcOp);
-      if (srcMemref != memref)
+      if (!accessesMemref(srcOp))
         return false;
       return llvm::find_if(dstMemOps, [&](Operation *dstOp) {
-               return srcMemref == getMemRef(dstOp);
+               return accessesMemref(dstOp);
              }) != dstMemOps.end();
     });
   };
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 1ea42517988c3..5f72bf79e0ac7 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -13,6 +13,33 @@
 
 // -----
 
+// CHECK-LABEL: func @memref_copy_in_loop_nest_does_not_crash() {
+func.func @memref_copy_in_loop_nest_does_not_crash() {
+  %mem = memref.alloca() : memref<10xi32>
+  %mem2 = memref.alloca() : memref<10xi32>
+  %c0 = arith.constant 0 : i32
+
+  affine.for %i = 0 to 10 {
+    memref.copy %mem, %mem2 : memref<10xi32> to memref<10xi32>
+    affine.store %c0, %mem[%i] : memref<10xi32>
+  }
+  affine.for %j = 0 to 10 {
+    %v = affine.load %mem[%j] : memref<10xi32>
+  }
+  affine.for %k = 0 to 10 {
+    affine.store %c0, %mem2[%k] : memref<10xi32>
+  }
+
+  // CHECK:      memref.copy
+  // CHECK:      affine.store
+  // CHECK:      affine.load
+  // CHECK:      affine.store
+  // 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 +1602,3 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
 }
 
 // Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
-

``````````

</details>


https://github.com/llvm/llvm-project/pull/203811