[PATCH] D154892: [MLIR][MemRef] Avoid returning ReallocOp from findDealloc

[Alexander Shaposhnikov via Phabricator via llvm-commits]

alexander-shaposhnikov created this revision.
alexander-shaposhnikov added reviewers: pifon2a, jpienaar.
alexander-shaposhnikov created this object with visibility "All Users".
Herald added subscribers: bviyer, Moerafaat, bzcheeseman, sdasgup3, wenzhicui, wrengr, cota, teijeong, rdzhabarov, tatianashp, msifontes, jurahul, Kayjukh, grosul1, Joonsoo, liufengdb, aartbik, mgester, arpith-jacob, antiagainst, shauheen, rriddle, mehdi_amini.
Herald added a reviewer: springerm.
Herald added a project: All.
alexander-shaposhnikov requested review of this revision.
Herald added a reviewer: nicolasvasilache.
Herald added subscribers: stephenneuendorffer, nicolasvasilache.
Herald added a project: MLIR.

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 multiple assumptions)

Test plan:
ninja check-mlir
ninja check-all


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D154892

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


Index: mlir/test/Dialect/Bufferization/canonicalize.mlir
===================================================================
--- mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -209,6 +209,19 @@
 
 // -----
 
+// 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>
Index: mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
===================================================================
--- mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
+++ mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
@@ -49,6 +49,9 @@
   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;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D154892.538814.patch
Type: text/x-patch
Size: 1438 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230710/f1376bc7/attachment.bin>