[Mlir-commits] [mlir] bb3f2bc - [MLIR][Affine] Fix crash in replaceAllMemRefUsesWith when replacement fails (#186282)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-23T13:21:24+01:00
New Revision: bb3f2bc86ca2dcd0c8688edd989055e0025c3554

URL: https://github.com/llvm/llvm-project/commit/bb3f2bc86ca2dcd0c8688edd989055e0025c3554
DIFF: https://github.com/llvm/llvm-project/commit/bb3f2bc86ca2dcd0c8688edd989055e0025c3554.diff

LOG: [MLIR][Affine] Fix crash in replaceAllMemRefUsesWith when replacement fails (#186282)

When the multi-user overload of `replaceAllMemRefUsesWith` iterates over
collected ops and calls the single-op overload, the single-op version
can legitimately fail (e.g., when the op uses the same memref in
multiple incompatible roles, such as both source and tag in
`affine.dma_start`). The code previously called `llvm_unreachable` in
this case, causing a crash.

Fix by propagating the failure via `return failure()` instead.

Fixes #60021

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Utils/Utils.cpp
    mlir/test/Dialect/Affine/pipeline-data-transfer.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index c905162987c39..7b976943e1bb7 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1374,7 +1374,7 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
     if (failed(replaceAllMemRefUsesWith(
             oldMemRef, newMemRef, user, extraIndices, indexRemap, extraOperands,
             symbolOperands, allowNonDereferencingOps)))
-      llvm_unreachable("memref replacement guaranteed to succeed here");
+      return failure();
   }
 
   return success();

diff  --git a/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir b/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
index 9ea282b1dc858..35507c37be79b 100644
--- a/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
+++ b/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
@@ -378,3 +378,21 @@ func.func @escaping_and_indexed_use_mix() {
 // CHECK-NEXT:   "compute"(%{{.*}}) : (memref<32xf32, 1>) -> ()
 // CHECK-NEXT:   [[VAL:%[0-9a-zA-Z_]+]] = affine.load %{{.*}}[%{{.*}}] : memref<32xf32, 1>
 // CHECK-NEXT:   "foo"([[VAL]]) : (f32) -> ()
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/60021.
+// An op that uses the same memref as both source and tag in dma_start should
+// not crash; the pass should gracefully skip replacement.
+// CHECK-LABEL: func @same_memref_source_and_tag
+func.func @same_memref_source_and_tag(%arg0: index, %arg1: index) {
+  %c0 = arith.constant 0 : index
+  %alloc = memref.alloc() : memref<1xf32>
+  %alloc_0 = memref.alloc() : memref<100x100xf32, 2>
+  affine.for %arg2 = 0 to 1 {
+    affine.dma_start %alloc[%arg0], %alloc_0[%arg0, %arg0], %alloc[%c0], %c0 : memref<1xf32>, memref<100x100xf32, 2>, memref<1xf32>
+    affine.dma_wait %alloc[%c0], %c0 : memref<1xf32>
+  }
+  return
+}
+// CHECK: affine.for