[mlir][bufferization] Add bufferization.dealloc canonicalizer to remove unused alloc-dealloc pairs

[maerhart]

Summary

Deallocation operations where the allocated value is the 'memref' and
'retained' list are currently not supported. This is because when values
are in the retained list, they typically have a use-site at a later
point and another deallocation op exists at that later point to free the
memref then. There alrady exists a canonicalization pattern in the
buffer deallocation simplification pass that removes the allocated value
from the earlier dealloc because it will never be actually deallocated
in that case and thus does not have to be considered in this new
pattern.

[springerm]

memref.alloc ops have a side effect. Do they fold away if if they have no uses? (without this pattern)

[maerhart]

In D158740#4614040, @springerm wrote:

memref.alloc ops have a side effect. Do they fold away if if they have no uses? (without this pattern)

Yes they do fold away when running -canonicalize because they only have the allocation side-effect. Should we rely on that instead of manually deleting the alloc op in the pattern?

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

SmallVector<Value> newMemrefs(
           return memref;
         }));
 
     return updateDeallocIfChanged(deallocOp, newMemrefs,
                                   deallocOp.getConditions(), rewriter);
   }
 };
 
+/// Removes pairs of `bufferization.dealloc` and alloc operations if there is no
+/// other user of the allocated value and the allocating operation can be safely
+/// removed. If the same value is present multiple times, this pattern relies on
+/// other canonicalization patterns to remove the duplicate first.
+///
+/// Example:
+/// ```mlir
+/// %alloc = memref.alloc() : memref<2xi32>
+/// bufferization.dealloc (%alloc, %arg0, : ...) if (%true, %true)
+/// ```
+/// is canonicalized to
+/// ```mlir
+/// bufferization.dealloc (%arg0 : ...) if (%true)
+/// ```
+struct RemoveAllocDeallocPairWhenNoOtherUsers
+    : public OpRewritePattern<DeallocOp> {
+  using OpRewritePattern<DeallocOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> newMemrefs, newConditions;
+    SmallVector<Operation *> toDelete;
+    for (auto [memref, cond] :
+         llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+      if (auto allocOp = memref.getDefiningOp<MemoryEffectOpInterface>()) {
+        // Check that it is indeed an allocate effect, that the op has no other
+        // side effects (which would not allow us to remove the op), and that
+        // there are no other users.
+        if (allocOp.getEffectOnValue<MemoryEffects::Allocate>(memref) &&
+            hasSingleEffect<MemoryEffects::Allocate>(allocOp, memref) &&
+            memref.hasOneUse()) {
+          toDelete.push_back(allocOp);
+          continue;
+        }
+      }
+
+      newMemrefs.push_back(memref);
+      newConditions.push_back(cond);
+    }
+
+    if (failed(updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
+                                      rewriter)))
+      return failure();
+
+    for (Operation *op : toDelete)
+      rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+
 } // anonymous namespace
 
 void DeallocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results.add<DeallocRemoveDuplicateDeallocMemrefs,
               DeallocRemoveDuplicateRetainedMemrefs, EraseEmptyDealloc,
-              EraseAlwaysFalseDealloc, SkipExtractMetadataOfAlloc>(context);
+              EraseAlwaysFalseDealloc, SkipExtractMetadataOfAlloc,
+              RemoveAllocDeallocPairWhenNoOtherUsers>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"

mlir/test/Dialect/Bufferization/canonicalize.mlir

 
 // CHECK-LABEL: func @dealloc_always_false_condition
 //  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: memref<2xi32>, [[ARG2:%.+]]: i1)
 //  CHECK-NEXT: bufferization.dealloc ([[ARG1]] : {{.*}}) if ([[ARG2]])
 //  CHECK-NEXT: return
 
 // -----
 
-func.func @dealloc_base_memref_extract_of_alloc(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>) {
+func.func @dealloc_base_memref_extract_of_alloc(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>) -> memref<2xi32> {
   %alloc = memref.alloc() : memref<2xi32>
   %base0, %size0, %stride0, %offset0 = memref.extract_strided_metadata %alloc : memref<2xi32> -> memref<i32>, index, index, index
   %base1, %size1, %stride1, %offset1 = memref.extract_strided_metadata %arg3 : memref<2xi32> -> memref<i32>, index, index, index
   bufferization.dealloc (%base0, %arg0, %base1 : memref<i32>, memref<2xi32>, memref<i32>) if (%arg1, %arg2, %arg2)
-  return
+  return %alloc : memref<2xi32>
 }
 
 // CHECK-LABEL: func @dealloc_base_memref_extract_of_alloc
 //  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xi32>)
 //  CHECK-NEXT: [[ALLOC:%.+]] = memref.alloc() : memref<2xi32>
 //  CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[ARG3]] :
 //  CHECK-NEXT: bufferization.dealloc ([[ALLOC]], [[ARG0]], [[BASE]] : memref<2xi32>, memref<2xi32>, memref<i32>) if ([[ARG1]], [[ARG2]], [[ARG2]])
 //  CHECK-NEXT: return
+
+// -----
+
+func.func @dealloc_base_memref_extract_of_alloc(%arg0: memref<2xi32>) {
+  %true = arith.constant true
+  %alloc = memref.alloc() : memref<2xi32>
+  bufferization.dealloc (%alloc, %arg0 : memref<2xi32>, memref<2xi32>) if (%true, %true)
+  return
+}
+
+// CHECK-LABEL: func @dealloc_base_memref_extract_of_alloc
+//  CHECK-SAME:([[ARG0:%.+]]: memref<2xi32>)
+//   CHECK-NOT: memref.alloc(
+//       CHECK: bufferization.dealloc ([[ARG0]] : memref<2xi32>) if (%true