Depends on D131121

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

if (cloneOp.use_empty()) {
       return success();
     }
 
     Value source = cloneOp.getInput();
 
     // We aim to check dealloc op for aliases below, e.g. ReshapeOp
     // which otherwise could prevent removal of unnecessary allocs.
     Operation *sourceDeallocOp;
-    Value source = cloneOp.getInput();
     Operation *source_op = cloneOp.getInput().getDefiningOp();
     while (auto iface = dyn_cast<ViewLikeOpInterface>(source_op))
       source_op = iface.getViewSource().getDefiningOp();
-    if (isa<memref::AllocOp>(source_op)) {
-      llvm::Optional<Operation *> tmpDeallocOp =
-          memref::findDealloc(source_op->getResult(0));
-      if (!tmpDeallocOp.has_value())
-        return failure();
-      Operation *tmpOp = *tmpDeallocOp;
-      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();
     Operation *cloneDeallocOp = *maybeCloneDeallocOp;
 
     // If both are deallocated in the same block, their in-block lifetimes

mlir/test/Dialect/Bufferization/canonicalize.mlir

func.func @alloc_tensor_canonicalize() -> (tensor<4x5x?xf32>) {
   %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 @Atanh_platform_elem_type_output_type(%arg0: memref<*xf32>) -> memref<*xf32> {
-  %0 = memref.rank %arg0 : memref<*xf32>
-  %1 = memref.alloca(%0) : memref<?xindex>
+// -----
+
+func.func @dealloc_canonicalize_clone_removal(%arg0: memref<?xindex>) -> memref<*xf32> {
   %c1 = arith.constant 1 : index
-  %2 = memref.alloc(%c1) : memref<?xf32>
-  %3 = memref.reshape %2(%1) : (memref<?xf32>, memref<?xindex>) -> memref<*xf32>
-  %4 = bufferization.clone %3 : memref<*xf32> to memref<*xf32>
-  memref.dealloc %2 : memref<?xf32>
-  return %4 : memref<*xf32>
+  %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: @Atanh_platform_elem_type_output_type
+// CHECK-LABEL: @dealloc_canonicalize_clone_removal
 //   CHECK-NOT:   bufferization.clone
 //   CHECK-NOT:   memref.dealloc
 //       CHECK:   return {{.*}}