Add a canonicalization pattern to remove Dealloc operations if the memref is an AllocOp that is only used by Dealloc operations.

PiperOrigin-RevId: 229606558

mlir/lib/StandardOps/StandardOps.cpp

@@ -671,6 +671,37 @@ Attribute CmpIOp::constantFold(ArrayRef<Attribute> operands,
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
+namespace {
+/// Fold Dealloc instructions that are deallocating an AllocOp that is only used
+/// by other Dealloc operations.
+struct SimplifyDeadDealloc : public RewritePattern {
+  SimplifyDeadDealloc(MLIRContext *context)
+      : RewritePattern(DeallocOp::getOperationName(), 1, context) {}
+
+  PatternMatchResult match(OperationInst *op) const override {
+    auto dealloc = op->cast<DeallocOp>();
+
+    // Check that the memref operand's defining instruction is an AllocOp.
+    Value *memref = dealloc->getMemRef();
+    OperationInst *defOp = memref->getDefiningInst();
+    if (!defOp || !defOp->isa<AllocOp>())
+      return matchFailure();
+
+    // Check that all of the uses of the AllocOp are other DeallocOps.
+    for (auto &use : memref->getUses()) {
+      auto *user = dyn_cast<OperationInst>(use.getOwner());
+      if (!user || !user->isa<DeallocOp>())
+        return matchFailure();
+    }
+    return matchSuccess();
+  }
+
+  void rewrite(OperationInst *op, PatternRewriter &rewriter) const override {
+    // Erase the dealloc operation.
+    op->erase();
+  }
+};
+} // end anonymous namespace.
 
 void DeallocOp::build(Builder *builder, OperationState *result, Value *memref) {
   result->addOperands(memref);
@@ -699,6 +730,7 @@ void DeallocOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
   /// dealloc(memrefcast) -> dealloc
   results.push_back(
       std::make_unique<MemRefCastFolder>(getOperationName(), context));
+  results.push_back(std::make_unique<SimplifyDeadDealloc>(context));
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Transforms/canonicalize.mlir

@@ -174,6 +174,33 @@ func @dead_alloc_fold() {
   return
 }
 
+// CHECK-LABEL: func @dead_dealloc_fold
+func @dead_dealloc_fold() {
+  // CHECK-NEXT: return
+  %a = alloc() : memref<4xf32>
+  dealloc %a: memref<4xf32>
+  return
+}
+
+// CHECK-LABEL: func @dead_dealloc_fold_multi_use
+func @dead_dealloc_fold_multi_use(%cond : i1) {
+  // CHECK-NEXT: cond_br
+  %a = alloc() : memref<4xf32>
+  cond_br %cond, ^bb1, ^bb2
+
+  // CHECK-LABEL: bb1:
+^bb1:
+  // CHECK-NEXT: return
+  dealloc %a: memref<4xf32>
+  return
+
+  // CHECK-LABEL: bb2:
+^bb2:
+  // CHECK-NEXT: return
+  dealloc %a: memref<4xf32>
+  return
+}
+
 // CHECK-LABEL: func @dyn_shape_fold(%arg0: index, %arg1: index)
 func @dyn_shape_fold(%L : index, %M : index) -> (memref<? x ? x i32>, memref<? x ? x f32>) {
   // CHECK: %c0 = constant 0 : index