[Mlir-commits] [mlir] [MLIR][Affine] Fix affine-loop-fusion crash on non-affine memory ops (PR #214995)

[llvmlistbot at llvm.org]

https://github.com/RohithPariki created https://github.com/llvm/llvm-project/pull/214995

Fixes #211591.

### Description
This fixes an issue in the affine loop fusion pass where dependence checks involving non-affine memory ops (such as `func.call` or custom memory ops) triggered an assertion failure. 

The failure occurred in `getMemRef()`, which incorrectly assumed that the operations passed to it were exclusively `memref.load`, `memref.store`, or `affine` read/write ops. When faced with a generic operation (like a call), it threw an `llvm_unreachable` assertion.

This patch updates the dependence check to use `getEffectedValues` from `MemoryEffects`. It can now correctly and safely check if the specific memref in question is accessed by these non-affine operations without crashing.

**Note to reviewers:**
* AI was used to help identify the root cause and generate this patch. 


>From d10d230b9bd199f8d3e03f97bb7090ac728b350f Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Sat, 8 Aug 2026 18:00:56 +0530
Subject: [PATCH] [MLIR][Affine] Fix affine-loop-fusion crash on non-affine
 memory ops

This fixes an issue in affine loop fusion where dependence checks on non-affine memory ops (such as func.call) triggered an assertion in getMemRef because it expected only affine or memref load/store ops. The dependence check now uses getEffectedValues to correctly check if the specific memref is accessed.
---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 35 +++++++++++-----------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 321c8e34d907c..2e22528dc8228 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,13 +262,25 @@ 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) {
+      if (auto memrefLoad = dyn_cast<memref::LoadOp>(op))
+        return memrefLoad.getMemRef() == memref;
+      if (auto affineLoad = dyn_cast<AffineReadOpInterface>(op))
+        return affineLoad.getMemRef() == memref;
+      if (auto memrefStore = dyn_cast<memref::StoreOp>(op))
+        return memrefStore.getMemRef() == memref;
+      if (auto affineStore = dyn_cast<AffineWriteOpInterface>(op))
+        return affineStore.getMemRef() == memref;
+      SmallVector<Value> effected;
+      getEffectedValues<MemoryEffects::Read, MemoryEffects::Write,
+                        MemoryEffects::Free>(op, effected);
+      return llvm::is_contained(effected, 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);
-             }) != dstMemOps.end();
+      return llvm::any_of(dstMemOps, accessesMemref);
     });
   };