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

[Alessandro Potenza]

================
@@ -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 {
----------------
alepot55 wrote:

This does fix the reported program. On a neighbouring one it trades one wrong answer for another: with the non-affine index at 3 rather than 0, `affineDataCopyGenerate` still sizes the buffer from the affine accesses alone, so the store comes out as `memref.store %c99_i32, %alloc[%4] : memref<1xi32>` with `%4` folded to 3, and the write to `%m[3]` is lost.

```mlir
memref.global "private" @m : memref<4xi32> = dense<[3, -7, 11, -13]>
func.func @main() {
  %s = arith.constant 99 : i32
  %c3 = arith.constant 3 : index
  %m = memref.get_global @m : memref<4xi32>
  affine.for %i = 0 to 4 {
    affine.store %s, %m[0] : memref<4xi32>
    memref.store %s, %m[%c3] : memref<4xi32>
  }
  ...
}
```

Running that under `--affine-data-copy-generate=fast-mem-space=0` I get m[0] 99 and m[3] -13 with this patch, and m[0] garbage with m[3] 99 without it. Would it be safer for the pass to skip a memref that has a non-affine user in range, rather than the predicate letting it through?


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