[mlir][memref] Fix canonicalization of memref.clone

[herhut]

Summary

The previous code was a bit too eager in canonicalizing.

[herhut]

I think this should become a pass instead, so that it can also consider aliases and
remove more clones. Block local aliases would suffice to catch situations like

%931 = memref.reshape %930(%107) : (memref<*xf16>, memref<?xindex>) -> memref<*xf16>
%932 = memref.clone %931 : memref<*xf16> to memref<*xf16>
memref.dealloc %930 : memref<*xf16>

[herhut]

Sent the wrong state of the change

[herhut]

Hey @dfki-jugr. I debugged this a bit today and here is how far I got. We can either land this as a cleanup or you can fold it into some changes you want to do anyway.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

537

This is a dangerous folding, as it means that a clone can change a memref's type. One then has to take care of this in all patterns that want to drop a clone. I'd rather just keep the casting behavior separate.

mlir/lib/Transforms/BufferDeallocation.cpp

154

This is a unrelated minor bug.

[dfki-jugr]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

508

Can we drop the last one in the block in this case?

527

nit: The last sentence is broken: information, of = if, drop on ...?

[herhut]

Fix comments and re-enable cast handling

[herhut]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

508

We should drop the earlier one, right? I won't add this now, maybe as a follow up.

[bkramer]

lgtm

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

LogicalResult matchAndRewrite(CloneOp cloneOp,
                                 PatternRewriter &rewriter) const override {
     if (cloneOp.use_empty()) {
       rewriter.eraseOp(cloneOp);
       return success();
     }
 
     Value source = cloneOp.input();
 
-    // Removes the clone operation and the corresponding dealloc and alloc
-    // operation (if any).
-    auto tryRemoveClone = [&](Operation *sourceOp, Operation *dealloc,
-                              Operation *alloc) {
-      if (!sourceOp || !dealloc || !alloc ||
-          alloc->getBlock() != dealloc->getBlock())
-        return false;
-      rewriter.replaceOp(cloneOp, source);
-      rewriter.eraseOp(dealloc);
-      return true;
-    };
-
-    // Removes unnecessary clones that are derived from the result of the clone
-    // op.
-    Operation *deallocOp = findDealloc(cloneOp.output());
-    Operation *sourceOp = source.getDefiningOp();
-    if (tryRemoveClone(sourceOp, deallocOp, sourceOp))
-      return success();
+    // This only finds dealloc operations for the immediate value. It should
+    // also consider aliases. That would also make the safety check below
+    // redundant.
+    Operation *cloneDeallocOp = findDealloc(cloneOp.output());
+    Operation *sourceDeallocOp = findDealloc(source);
+
+    // 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();
+    Operation *redundantDealloc = nullptr;
+    if (cloneDeallocOp && cloneDeallocOp->getBlock() == currentBlock) {
+      redundantDealloc = cloneDeallocOp;
+    } else if (sourceDeallocOp && sourceDeallocOp->getBlock() == currentBlock) {
+      redundantDealloc = sourceDeallocOp;
+    }
+
+    if (!redundantDealloc)
+      return failure();
+
+    // Safety check that there are no other deallocations inbetween
+    // cloneOp and redundantDealloc, as otherwise we might deallocate an alias
+    // of source before the uses of the clone. With alias information, we could
+    // restrict this to only fail of the dealloc's operand is an alias
+    // of the source.
+    for (Operation *pos = cloneOp->getNextNode(); pos != redundantDealloc;
+         pos = pos->getNextNode()) {
+      auto effectInterface = dyn_cast<MemoryEffectOpInterface>(pos);
+      if (!effectInterface)
+        continue;
+      if (effectInterface.hasEffect<MemoryEffects::Free>())
+        return failure();
+    }
 
-    // Removes unnecessary clones that are derived from the source of the clone
-    // op.
-    deallocOp = findDealloc(source);
-    if (tryRemoveClone(sourceOp, deallocOp, cloneOp))
-      return success();
-
-    return failure();
+    rewriter.replaceOpWithNewOp<memref::CastOp>(cloneOp, cloneOp.getType(),
+                                                source);
+    rewriter.eraseOp(redundantDealloc);
   }
 };
 
 } // end anonymous namespace.
 
 void CloneOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                           MLIRContext *context) {
   results.insert<SimplifyClones>(context);
 }
 
 OpFoldResult CloneOp::fold(ArrayRef<Attribute> operands) {
   return succeeded(foldMemRefCast(*this)) ? getResult() : Value();
 }
 
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult DeallocOp::fold(ArrayRef<Attribute> cstOperands,

mlir/lib/Transforms/BufferDeallocation.cpp

void recurse(Operation *op, Block *predecessor) {
     // If the current op implements the `BranchOpInterface`, there can be
     // cycles in the scope of all successor blocks.
     if (isa<BranchOpInterface>(op)) {
       for (Block *succ : current->getSuccessors())
         recurse(*succ, current);
     }
     // Recurse into all distinct regions and check for explicit control-flow
     // loops.
-    for (Region &region : op->getRegions())
+    for (Region &region : op->getRegions()) {
+      if (!region.empty())
         recurse(region.front(), current);
     }
+  }
 
   /// Recurses into explicit control-flow structures that are given by
   /// the successor relation defined on the block level.
   void recurse(Block &block, Block *predecessor) {
     // Try to enter the current block. If this is not possible, we are
     // currently processing this block and can safely return here.
     if (!enter(block, predecessor))
       return;

mlir/test/Dialect/MemRef/canonicalize.mlir

-// RUN: mlir-opt %s -canonicalize --split-input-file | FileCheck %s
+// RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
 
 // Test case: Basic folding of memref.tensor_load(memref.buffer_cast(t)) -> t
 // CHECK-LABEL: func @tensor_load_of_buffer_cast(
 //  CHECK-SAME:   %[[TENSOR:.*]]: tensor<?xf32>) -> tensor<?xf32> {
 //       CHECK: return %[[TENSOR]]
 func @tensor_load_of_buffer_cast(%arg0: tensor<?xf32>) -> tensor<?xf32> {
   %0 = memref.buffer_cast %arg0 : memref<?xf32>
   %1 = memref.tensor_load %0 : memref<?xf32>
func @rank_reducing_subview_canonicalize(%arg0 : memref<?x?x?xf32>, %arg1 : index,
   return %0 : memref<?x?xf32, #map0>
 }
 // CHECK-LABEL: func @rank_reducing_subview_canonicalize
 //  CHECK-SAME:   %[[ARG0:.+]]: memref<?x?x?xf32>
 //       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1]
 //  CHECK-SAME:      [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1]
 //  CHECK-SAME:      : memref<?x?x?xf32> to memref<4x?xf32
 //       CHECK:   %[[RESULT:.+]] = memref.cast %[[SUBVIEW]]
-//       CHEKC:   return %[[RESULT]]
+//       CHECK:   return %[[RESULT]]
+
+// -----
+
+// CHECK-LABEL: @clone_before_dealloc
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+func @clone_before_dealloc(%arg0: memref<?xf32>) -> memref<?xf32> {
+  // CHECK-NEXT: return %[[ARG]]
+  %0 = memref.clone %arg0 : memref<?xf32> to memref<?xf32>
+  memref.dealloc %arg0 : memref<?xf32>
+  return %0 : memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @clone_before_dealloc
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+func @clone_before_dealloc(%arg0: memref<?xf32>) -> memref<?xf32> {
+  // CHECK-NEXT: "use"(%arg0)
+  // CHECK-NEXT: return %[[ARG]]
+  %0 = memref.clone %arg0 : memref<?xf32> to memref<?xf32>
+  "use"(%0) : (memref<?xf32>) -> ()
+  memref.dealloc %0 : memref<?xf32>
+  return %arg0 : memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @clone_after_cast
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+func @clone_after_cast(%arg0: memref<?xf32>) -> memref<32xf32> {
+  // CHECK-NEXT: memref.clone %[[ARG]] : memref<?xf32> to memref<32xf32>
+  // CHECK-NOT: memref.cast
+  %0 = memref.cast %arg0 : memref<?xf32> to memref<32xf32>
+  %1 = memref.clone %0 : memref<32xf32> to memref<32xf32>
+  return %1 : memref<32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @clone_and_cast
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+func @clone_and_cast(%arg0: memref<?xf32>) -> memref<32xf32> {
+  // CHECK-NEXT: %[[RES:.*]] = memref.cast %[[ARG]] : memref<?xf32> to memref<32xf32>
+  %0 = memref.clone %arg0 : memref<?xf32> to memref<32xf32>
+  // CHECK-NEXT: return %[[RES]]
+  memref.dealloc %arg0 : memref<?xf32>
+  return %0 : memref<32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @alias_is_freed
+func @alias_is_freed(%arg0 : memref<?xf32>) {
+  // CHECK: memref.clone
+  // CHECK: memref.dealloc
+  // CHECK: memref.dealloc
+  %0 = memref.cast %arg0 : memref<?xf32> to memref<32xf32>
+  %1 = memref.clone %0 : memref<32xf32> to memref<32xf32>
+  memref.dealloc %arg0 : memref<?xf32>
+  "use"(%1) : (memref<32xf32>) -> ()
+  memref.dealloc %1 : memref<32xf32>
+  return
+}