add dealloc alias check to bufferization

[kushanam]

CC @frgossen for review.

[frgossen]

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

456

is currently unused

460–469

Starting with if (isa<memref::..., I think this can still be the old logic. The loop above is all you need, I think.

mlir/test/Dialect/Bufferization/canonicalize.mlir

258–273

Can you add the // ----- between the test cases?

260

Can you give the test function a meaningful name?

262

This could be passed in as an arg to keep the test simpler.

[kushanam]

Add dealloc alias check to bufferization
Traverse the cloneOp for aliases to find the alloc op

[frgossen]

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

457

How about only updateing the above source value/op here and using the previous implementation with the maybeSourceDeallocOp?
That would make this patch very simple to understand if we break something.

[kushanam]

add dealloc alias check to bufferization
Traverse the cloneOp for aliases to find the alloc op

[kushanam]

Add dealloc alias check to bufferization
Traverse the cloneOp for aliases to find the alloc op

[frgossen]

Nice!

Now the test cases:

Failed Tests (3):
  MLIR :: Dialect/Bufferization/canonicalize.mlir
  MLIR :: Dialect/Bufferization/inlining.mlir
  MLIR :: Transforms/canonicalize.mlir

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

455

LLVM naming convention wants it to be named sourceOp.

nit: source.getDefiningOp()

[kushanam]

Moved to D131740

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

LogicalResult matchAndRewrite(CloneOp cloneOp,
                                 PatternRewriter &rewriter) const override {
     if (cloneOp.use_empty()) {
       rewriter.eraseOp(cloneOp);
       return success();
     }
 
     Value source = cloneOp.getInput();
 
-    // This only finds dealloc operations for the immediate value. It should
-    // also consider aliases. That would also make the safety check below
-    // redundant.
+    // We aim to check dealloc op for aliases below, e.g. ReshapeOp
+    // which otherwise could prevent removal of unnecessary allocs.
+    Operation *sourceDeallocOp;
+    Operation *source_op = cloneOp.getInput().getDefiningOp();
+    while (auto iface = dyn_cast<ViewLikeOpInterface>(source_op))
+      source_op = iface.getViewSource().getDefiningOp();
+    llvm::Optional<Operation *> maybeSourceDeallocOp =
+        memref::findDealloc(source_op->getResult(0));
+    if (!maybeSourceDeallocOp.has_value())
+      return failure();
+    Operation *tmpOp = *maybeSourceDeallocOp;
+    if (tmpOp) {
+      sourceDeallocOp = tmpOp;
+    }
+    
     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);
-    if (!maybeSourceDeallocOp.has_value())
-      return failure();
     Operation *cloneDeallocOp = *maybeCloneDeallocOp;
-    Operation *sourceDeallocOp = *maybeSourceDeallocOp;
 
     // If both are deallocated in the same block, their in-block lifetimes
     // might not fully overlap, so we cannot decide which one to drop.
     if (cloneDeallocOp && sourceDeallocOp &&
         cloneDeallocOp->getBlock() == sourceDeallocOp->getBlock())
       return failure();
 
     Block *currentBlock = cloneOp->getBlock();

mlir/test/Dialect/Bufferization/canonicalize.mlir

 func.func @alloc_tensor_canonicalize() -> (tensor<4x5x?xf32>) {
   %c6 = arith.constant 6 : index
   %0 = bufferization.alloc_tensor(%c6) : tensor<4x5x?xf32>
   return %0 : tensor<4x5x?xf32>
 }
 // CHECK: func @alloc_tensor_canonicalize
 // 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 {{.*}}