[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 30548f97a6572a0c16d0619a3b9345837960e5f8 Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Fri, 18 Sep 2026 19:51:08 +0530
Subject: [PATCH 1/3] [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  | 17 ++++++++++++
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 67885513dcd3be..2b0338276024c8 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -243,17 +243,15 @@ 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);
+  for (Value effectedValue : effectedValues)
+    if (effectedValue == memref)
+      return true;
+  return false;
 }
 
 /// Returns true if there may be a dependence on `memref` from srcNode's
@@ -272,16 +270,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)
+      if (!mayAccessMemRef(srcOp, memref))
         return false;
-      return llvm::find_if(dstMemOps, [&](Operation *dstOp) {
-               return srcMemref == getMemRef(dstOp);
-             }) != dstMemOps.end();
+      return llvm::any_of(dstMemOps, [&](Operation *dstOp) {
+        return mayAccessMemRef(dstOp, memref);
+      });
     });
   };
 
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 1ea42517988c3c..4f72b2deab50cf 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1576,3 +1576,20 @@ 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
+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
+}

>From d791fc396a8612f59eab1172c080e6ac0ebc9e84 Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Sat, 19 Sep 2026 20:09:20 +0530
Subject: [PATCH 2/3] [mlir][affine] Check vector transfer loops remain unfused

---
 mlir/test/Dialect/Affine/loop-fusion.mlir | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 4f72b2deab50cf..9dffe5c9318666 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1579,6 +1579,11 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
 // -----
 
 // 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
 

>From 8a37ac541545e65b29df06e6e1da5db3dfa89d8e Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Wed, 23 Sep 2026 13:53:13 +0530
Subject: [PATCH 3/3] [mlir][affine] Use llvm::is_contained for memref access
 check

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

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 2b0338276024c8..e8314584225149 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -248,10 +248,7 @@ static bool mayAccessMemRef(Operation *op, Value memref) {
   SmallVector<Value> effectedValues;
   getEffectedValues<MemoryEffects::Read, MemoryEffects::Write>(op,
                                                                effectedValues);
-  for (Value effectedValue : effectedValues)
-    if (effectedValue == memref)
-      return true;
-  return false;
+  return llvm::is_contained(effectedValues, memref);
 }
 
 /// Returns true if there may be a dependence on `memref` from srcNode's