add dealloc alias check to bufferization

[kushanam]

Summary

Traverse the cloneOp for aliases to find the alloc op

[frgossen]

Nice! Thank you!

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

455

LLVM coding style is sourceOp. Please fix the other variable names also.

Maybe a better name here is allocOp, which is what you're trying to find.

455

Instead of cloneOp.getInput() you can reuse source from above

456–457

What if it is an alias of an argument, e.g. (reshape(%buffer_arg)), then you will run into a nullptr here.

You can use dyn_cast_or_null. Then the surrounding if statement is not needed also.

465

This is a canonicalSource, possibly the allocation value.

[kushanam]

This comment was removed by kushanam.

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 *source_op = cloneOp.getInput().getDefiningOp();
+    if (source_op)
+      while (auto iface = dyn_cast<ViewLikeOpInterface>(source_op))
+        source_op = iface.getViewSource().getDefiningOp();
+    
     llvm::Optional<Operation *> maybeCloneDeallocOp =
         memref::findDealloc(cloneOp.getOutput());
     // Skip if either of them has > 1 deallocate operations.
     if (!maybeCloneDeallocOp.has_value())
       return failure();
+    Value source_val = source_op ? source_op->getResult(0) : source;
     llvm::Optional<Operation *> maybeSourceDeallocOp =
-        memref::findDealloc(source);
+        memref::findDealloc(source_val);
     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 &&

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 {{.*}}