[Mlir-commits] [mlir] 894e8a5 - [MLIR] Add dealloc alias check to bufferization

[Frederik Gossen]

Author: root
Date: 2022-08-17T19:11:59-04:00
New Revision: 894e8a5446dcee1ab7a2b36471a27e9025d27607

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

LOG: [MLIR] Add dealloc alias check to bufferization

Traverse the cloneOp for aliases to find the alloc op

Reviewed By: frgossen, bondhugula

Differential Revision: https://reviews.llvm.org/D131797

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 9ca51f8beadd6..5fc8e6abb0a25 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -449,17 +449,20 @@ struct SimplifyClones : public OpRewritePattern<CloneOp> {
     }
 
     Value source = cloneOp.getInput();
+    // Aims to find the dealloc op for the canonical source
+    // which otherwise could prevent removal of unnecessary allocs.
+    Value canonicalSource = source;
+    while (auto iface = dyn_cast_or_null<ViewLikeOpInterface>(
+               canonicalSource.getDefiningOp()))
+      canonicalSource = iface.getViewSource();
 
-    // This only finds dealloc operations for the immediate value. It should
-    // also consider aliases. That would also make the safety check below
-    // redundant.
     llvm::Optional<Operation *> maybeCloneDeallocOp =
         memref::findDealloc(cloneOp.getOutput());
     // Skip if either of them has > 1 deallocate operations.
     if (!maybeCloneDeallocOp.has_value())
       return failure();
     llvm::Optional<Operation *> maybeSourceDeallocOp =
-        memref::findDealloc(source);
+        memref::findDealloc(canonicalSource);
     if (!maybeSourceDeallocOp.has_value())
       return failure();
     Operation *cloneDeallocOp = *maybeCloneDeallocOp;

diff  --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index f76f040937aaf..fa18b1502f40d 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -256,3 +256,18 @@ func.func @alloc_tensor_canonicalize() -> (tensor<4x5x?xf32>) {
 // CHECK:   %[[T0:.+]] = bufferization.alloc_tensor() : tensor<4x5x6xf32>
 // CHECK:   %[[T1:.+]] = tensor.cast %[[T0]] : tensor<4x5x6xf32> to tensor<4x5x?xf32>
 // CHECK:   return %[[T1]]
+
+// -----
+
+func.func @dealloc_canonicalize_clone_removal(%arg0: memref<?xindex>) -> memref<*xf32> {
+  %c1 = arith.constant 1 : index
+  %0 = memref.alloc(%c1) : memref<?xf32>
+  %1 = memref.reshape %0(%arg0) : (memref<?xf32>, memref<?xindex>) -> memref<*xf32>
+  %2 = bufferization.clone %1 : memref<*xf32> to memref<*xf32>
+  memref.dealloc %0 : memref<?xf32>
+  return %2 : memref<*xf32>
+}
+// CHECK-LABEL: @dealloc_canonicalize_clone_removal
+//   CHECK-NOT:   bufferization.clone
+//   CHECK-NOT:   memref.dealloc
+//       CHECK:   return {{.*}}