[Mlir-commits] [mlir] cd99637 - [mlir][bufferization] Fix stale BufferOriginAnalysis in BufferDeallocationSimplification (#210105)

[llvmlistbot at llvm.org]

Author: Vito Secona
Date: 2026-07-19T14:46:58+02:00
New Revision: cd99637821097c239b665793243d8e6a332fc719

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

LOG: [mlir][bufferization] Fix stale BufferOriginAnalysis in BufferDeallocationSimplification (#210105)

Fixes #205228

This change fixes a use-after-free bug in the
BufferDeallocationSimplification pipeline caused by the greedy pattern
rewriter deleting operation tracked by the BufferOriginAnalysis.
BufferViewFlowAnalysis (internally used by BufferOriginAnalysis) creates
the dependencies map once during initialization. In this specific case,
the folder removed the `memref.cast` ops making later uses a
use-after-free bug.

Below is the snippet of running the pass using
`-debug-only=greedy-rewriter`. It erases a `memref.cast` and crashes in
RemoveDeallocMemrefsContainedInRetained.

```
...
[greedy-rewriter:1] //===-------------------------------------------===//
[greedy-rewriter:1] Processing operation : 'memref.cast'(0x654d43d18670) {
[greedy-rewriter:1]   %4 = "memref.cast"(%3) : (memref<?x?xf32, 1>) -> memref<4x4xf32, 1>
[greedy-rewriter:1] 
[greedy-rewriter:1] } -> success : operation was folded
[greedy-rewriter:1] //===-------------------------------------------===//
[greedy-rewriter:1] ** Replace : 'memref.cast'(0x654d43d18670)
[greedy-rewriter:1] ** Modified: 'scf.yield'(0x654d43c97bb0)
[greedy-rewriter:1] ** Erase   : 'memref.cast'(0x654d43d18670)
...
[greedy-rewriter:1] Processing operation : 'bufferization.dealloc'(0x654d43cbd7d0) {
[greedy-rewriter:1]   %5 = "bufferization.dealloc"(%4#0, %1, %3) <{operandSegmentSizes = array<i32: 1, 1, 1>}> : (memref<f32, 1>, i1, memref<4x4xf32, 1>) -> i1
[greedy-rewriter:1] 
[greedy-rewriter:1] 
[greedy-rewriter:1]   * Pattern (anonymous namespace)::RemoveDeallocMemrefsContainedInRetained : 'bufferization.dealloc -> ()' {
```

The fix is to disable folding and split the canonicalization patterns to
run separately. This way, the patterns that rely on BufferOriginAnalysis
still work because folding was disabled and BufferOriginAnalysis never
invalidated. Then, the canonicalization patterns still benefit from
folding.

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
    mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
index a465c957d063e..06e90333fe469 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
@@ -469,12 +469,24 @@ struct BufferDeallocationSimplificationPass
                  RetainedMemrefAliasingAlwaysDeallocatedMemref>(&getContext(),
                                                                 analysis);
 
-    populateDeallocOpCanonicalizationPatterns(patterns, &getContext());
     // We don't want that the block structure changes invalidating the
     // `BufferOriginAnalysis` so we apply the rewrites with `Normal` level of
-    // region simplification
+    // region simplification and disable folding.
     if (failed(applyPatternsGreedily(
             getOperation(), std::move(patterns),
+            GreedyRewriteConfig()
+                .setRegionSimplificationLevel(GreedySimplifyRegionLevel::Normal)
+                .enableFolding(false))))
+      signalPassFailure();
+
+    // We run canonicalization separately so it can benefit from
+    // folding, which was disabled previously to avoid invalidating
+    // `BufferOriginAnalysis`.
+    RewritePatternSet canonicalizationPatterns(&getContext());
+    populateDeallocOpCanonicalizationPatterns(canonicalizationPatterns,
+                                              &getContext());
+    if (failed(applyPatternsGreedily(
+            getOperation(), std::move(canonicalizationPatterns),
             GreedyRewriteConfig().setRegionSimplificationLevel(
                 GreedySimplifyRegionLevel::Normal))))
       signalPassFailure();

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
index b40a17cf800bf..3acc09d23b59e 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
@@ -169,3 +169,32 @@ func.func @duplicate_memref(%arg0: memref<5xf32>, %arg1: memref<6xf32>, %c: i1)
 // CHECK-LABEL: func @duplicate_memref(
 //       CHECK:   %[[r:.*]] = bufferization.dealloc (%{{.*}} : memref<5xf32>) if (%{{.*}}) retain (%{{.*}} : memref<6xf32>)
 //       CHECK:   return %[[r]]
+
+// -----
+
+module {
+  func.func @memref_cast_folding(%arg0: memref<4x4xf32>) -> memref<4x4xf32, 1> {
+    %cfalse = arith.constant false
+    %ctrue = arith.constant true
+
+    %memspacecast = memref.memory_space_cast %arg0 : memref<4x4xf32> to memref<4x4xf32, 1>
+    %cast = memref.cast %memspacecast : memref<4x4xf32, 1> to memref<?x?xf32, 1>
+    %cast_1 = memref.cast %cast : memref<?x?xf32, 1> to memref<4x4xf32, 1>
+
+    %5 = scf.if %cfalse -> (memref<4x4xf32, 1>) {
+      scf.yield %cast_1 : memref<4x4xf32, 1>
+    } else {
+      %7 = bufferization.clone %cast_1 : memref<4x4xf32, 1> to memref<4x4xf32, 1>
+      scf.yield %7 : memref<4x4xf32, 1>
+    }
+
+    %base_buffer, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<4x4xf32> -> memref<f32>, index, index, index, index, index
+    %base_buffer_5, %offset_6, %sizes_7:2, %strides_8:2 = memref.extract_strided_metadata %memspacecast : memref<4x4xf32, 1> -> memref<f32, 1>, index, index, index, index, index
+
+    %6 = bufferization.dealloc (%base_buffer, %base_buffer_5 : memref<f32>, memref<f32, 1>) if (%cfalse, %ctrue) retain (%5 : memref<4x4xf32, 1>)
+    return %memspacecast : memref<4x4xf32, 1>
+  }
+}
+
+// CHECK-LABEL: func @memref_cast_folding
+//       CHECK: %[[r:.*]] = bufferization.dealloc (%{{.*}} : memref<f32, 1>) if (%{{.*}}) retain (%{{.*}} : memref<4x4xf32, 1>)