[Mlir-commits] [mlir] [mlir][affine] Fix replaceAllMemRefUsesWith failing on memref.load/store (PR #217833)

[Longsheng Mou]

https://github.com/CoTinker updated https://github.com/llvm/llvm-project/pull/217833

>From 4d9f5d98c510d82cd84353fc35142ba100d2f4df Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Fri, 21 Aug 2026 15:33:06 +0800
Subject: [PATCH 1/3] [mlir][affine] Fix replaceAllMemRefUsesWith failing on
 memref.load/store

memref.load/memref.store are dereferencing ops but don't implement AffineMapAccessInterface. The check `!isa<AffineMapAccessInterface>(user)` in `replaceAllMemRefUsesWith` incorrectly treats them as non-dereferencing and returns failure, causing the entire memref replacement to be skipped. This leaves fast buffers uninitialized in affine-data-copy-generate.
Use isDereferencingOp(user) instead, which already covers memref::LoadOp/memref::StoreOp.
---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7043083298615..7bb67e848f7b8 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1356,8 +1356,9 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
 
     // Check if the memref was used in a non-dereferencing context. It is fine
     // for the memref to be used in a non-dereferencing way outside of the
-    // region where this replacement is happening.
-    if (!isa<AffineMapAccessInterface>(*user)) {
+    // region where this replacement is happening. Note: memref.load/store are
+    // dereferencing ops even though they don't implement AffineMapAccessInterface.
+    if (!isDereferencingOp(user)) {
       if (!allowNonDereferencingOps) {
         LLVM_DEBUG(
             llvm::dbgs()

>From 11f3fe5a20e50aed9026f1f6accb97bc211c352a Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Fri, 21 Aug 2026 15:34:34 +0800
Subject: [PATCH 2/3] add test

---
 .../test/Dialect/Affine/affine-data-copy.mlir | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/mlir/test/Dialect/Affine/affine-data-copy.mlir b/mlir/test/Dialect/Affine/affine-data-copy.mlir
index f6d7186100fa6..c3c30925c94d1 100644
--- a/mlir/test/Dialect/Affine/affine-data-copy.mlir
+++ b/mlir/test/Dialect/Affine/affine-data-copy.mlir
@@ -495,3 +495,32 @@ func.func @multiple_blocks(%arg0: index) -> memref<1x2x1xi32> {
 ^bb3:  // pred: ^bb1
   return %1 : memref<1x2x1xi32>
 }
+
+// -----
+
+// Test that memref.store on the same memref as affine.load/store doesn't
+// cause replaceAllMemRefUsesWith to fail. Before the fix, memref.store
+// (which doesn't implement AffineMapAccessInterface) caused the entire
+// memref replacement to fail silently, leaving the fast buffer uninitialized.
+
+memref.global "private" @__constant_4xi32 : memref<4xi32> = dense<[3, -7, 11, -13]>
+
+// CHECK-LABEL:   func.func @memref_store_with_affine_access() -> i32 {
+// CHECK:           %[[ALLOC:.*]] = memref.alloc() : memref<1xi32>
+// CHECK:           affine.for
+// CHECK:             affine.store %{{.*}}, %[[ALLOC]]
+// CHECK:             memref.store %{{.*}}, %[[ALLOC]]
+// CHECK:           affine.load %[[ALLOC]]
+// CHECK:           memref.dealloc %[[ALLOC]]
+// CHECK:         }
+func.func @memref_store_with_affine_access() -> i32 {
+  %c0 = arith.constant 0 : index
+  %c100_i32 = arith.constant 100 : i32
+  %0 = memref.get_global @__constant_4xi32 : memref<4xi32>
+  affine.for %arg0 = 0 to 4 {
+    affine.store %c100_i32, %0[0] : memref<4xi32>
+    memref.store %c100_i32, %0[%c0] : memref<4xi32>
+  }
+  %1 = affine.load %0[0] : memref<4xi32>
+  return %1 : i32
+}

>From a9bfb772ae6c9184a2af9f7cdd74b38645393092 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Fri, 21 Aug 2026 15:56:48 +0800
Subject: [PATCH 3/3] code format

---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7bb67e848f7b8..983150fefb719 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1357,7 +1357,8 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
     // Check if the memref was used in a non-dereferencing context. It is fine
     // for the memref to be used in a non-dereferencing way outside of the
     // region where this replacement is happening. Note: memref.load/store are
-    // dereferencing ops even though they don't implement AffineMapAccessInterface.
+    // dereferencing ops even though they don't implement
+    // AffineMapAccessInterface.
     if (!isDereferencingOp(user)) {
       if (!allowNonDereferencingOps) {
         LLVM_DEBUG(