[Mlir-commits] [mlir] [mlir][affine] Handle generic memory ops in loop fusion dependence analysis (PR #224657)

[purnima shrivastava]

https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/224657

>From 35b72a6d97445d508a9aecaadfbc9f16ac479aaa Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Wed, 23 Sep 2026 16:22:15 +0530
Subject: [PATCH] [mlir][affine] Handle generic memory ops in loop fusion
 dependence analysis

---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 31 +++++++++-------------
 mlir/test/Dialect/Affine/loop-fusion.mlir  | 22 +++++++++++++++
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 67885513dcd3b..c99a88592a269 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -243,17 +243,12 @@ 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 `op` may read from or write to `memref`.
+static bool mayAccessMemRef(Operation *op, Value memref) {
+  SmallVector<Value> effectedValues;
+  getEffectedValues<MemoryEffects::Read, MemoryEffects::Write>(op,
+                                                               effectedValues);
+  return llvm::is_contained(effectedValues, memref);
 }
 
 /// Returns true if there may be a dependence on `memref` from srcNode's
@@ -272,17 +267,15 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
   // true if there exists a conflicting read/write access involving such.
 
   // Check whether there is a dependence from a source read/write op to a
-  // destination read/write one; all expected to be memref/affine load/store.
+  // destination read/write op on `memref`.
   auto hasNonAffineDep = [&](ArrayRef<Operation *> srcMemOps,
                              ArrayRef<Operation *> dstMemOps) {
     return llvm::any_of(srcMemOps, [&](Operation *srcOp) {
-      Value srcMemref = getMemRef(srcOp);
-      if (srcMemref != memref)
-        return false;
-      return llvm::find_if(dstMemOps, [&](Operation *dstOp) {
-               return srcMemref == getMemRef(dstOp);
-             }) != dstMemOps.end();
-    });
+             return mayAccessMemRef(srcOp, memref);
+           }) &&
+           llvm::any_of(dstMemOps, [&](Operation *dstOp) {
+             return mayAccessMemRef(dstOp, memref);
+           });
   };
 
   SmallVector<Operation *> dstOps;
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 1ea42517988c3..9dffe5c931866 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1576,3 +1576,25 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
 
 // Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
 
+// -----
+
+// CHECK-LABEL: func @vector_transfer_read
+// CHECK:      affine.for
+// CHECK-NEXT:   affine.store
+// CHECK-NEXT: }
+// CHECK-NEXT: affine.for
+// CHECK-NEXT:   vector.transfer_read
+func.func @vector_transfer_read(%arg0: memref<4xf32>) {
+  %c0 = arith.constant 0.0 : f32
+
+  affine.for %i = 0 to 4 {
+    affine.store %c0, %arg0[%i] : memref<4xf32>
+  }
+
+  affine.for %i = 0 to 4 {
+    %0 = vector.transfer_read %arg0[%i], %c0
+        : memref<4xf32>, vector<4xf32>
+  }
+
+  return
+}