[Mlir-commits] [mlir] 57609fb - [MLIR][MemRef] Avoid returning ReallocOp from findDealloc

[Alexander Shaposhnikov]

Author: Alexander Shaposhnikov
Date: 2023-07-11T21:33:06Z
New Revision: 57609fb6dde1fce5dfa850816b964c8b5fda65a3

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

LOG: [MLIR][MemRef] Avoid returning ReallocOp from findDealloc

This diff makes findDealloc return nullopt if the list of users contains "realloc".
Fixes https://github.com/llvm/llvm-project/issues/60726
(In particular, in SimplifyClones (uses findDealloc) treating "realloc" as "dealloc"
breaks the assumption that "dealloc" has no users)

Test plan: ninja check-mlir check-all

Differential revision: https://reviews.llvm.org/D154892

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
    mlir/test/Dialect/Bufferization/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
index f5ddc945459858..8304000bbcfd16 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
@@ -49,6 +49,9 @@ std::optional<Operation *> mlir::memref::findDealloc(Value allocValue) {
   for (Operation *user : allocValue.getUsers()) {
     if (!hasEffect<MemoryEffects::Free>(user, allocValue))
       continue;
+    // If we found a realloc instead of dealloc, return std::nullopt.
+    if (isa<memref::ReallocOp>(user))
+      return std::nullopt;
     // If we found > 1 dealloc, return std::nullopt.
     if (dealloc)
       return std::nullopt;

diff  --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index fae4351417648e..07959e29cd788f 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -209,6 +209,19 @@ func.func @clone_multiple_dealloc_of_clone(%arg0: memref<?xf32>) -> memref<?xf32
 
 // -----
 
+// Verify SimplifyClones skips clones followed by realloc.
+// CHECK-LABEL: @clone_and_realloc
+func.func @clone_and_realloc(%arg0: memref<?xf32>) {
+  %0 = bufferization.clone %arg0 : memref<?xf32> to memref<32xf32>
+  "use"(%0) : (memref<32xf32>) -> ()
+  %1 = memref.realloc %0 : memref<32xf32> to memref<64xf32>
+  memref.dealloc %1 : memref<64xf32>
+  return
+}
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+// CHECK-NOT: %cast = memref.cast %[[ARG]]
+
+// -----
 
 // CHECK-LABEL: func @tensor_cast_to_memref
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<4x6x16x32xi8>