[mlir][bufferization] Add DeallocOp canonicalizer to remove memrefs also present in the retained list

[maerhart]

Summary

Since memrefs in the retained list will never be deallocated, we can remove them from the list of memrefs to be deallocated. If the list of memrefs to deallocate becomes empty, we can just delete the dealloc operation.

Depends on D156157

[springerm]

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

786

Should this be memref?

825

nit: This should not be necessary.

mlir/test/Dialect/Bufferization/canonicalize.mlir

301

Some of these results should have a use, so that we can check if the result was replaced correctly.

[maerhart]

Address comments

[maerhart]

Sry, forgot to submit these two comments earlier.

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

786

No, the operation results represent the updated conditions/ownerships. If a value to be deallocated is in the list of values to be retained, it is guaranteed to not be deallocated, thus the condition/ownership stays the same.

825

I changed this a bit, but we need a check here for the case where the result was replaced with the condition operand instead of one of the new results.

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

 namespace {
 
 struct DeallocRemoveDuplicates : public OpRewritePattern<DeallocOp> {
   using OpRewritePattern<DeallocOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(DeallocOp deallocOp,
                                 PatternRewriter &rewriter) const override {
     // Unique memrefs to be deallocated.
+    DenseSet<Value> retained(deallocOp.getRetained().begin(),
+                             deallocOp.getRetained().end());
     DenseMap<Value, unsigned> memrefToCondition;
     SmallVector<Value> newMemrefs, newConditions, newRetained;
     SmallVector<unsigned> resultIndices;
-    for (auto [memref, cond] :
-         llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+    for (auto [i, memref, cond] :
+         llvm::enumerate(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+      if (retained.contains(memref)) {
+        rewriter.replaceAllUsesWith(deallocOp.getResult(i),
+                                    deallocOp.getConditions()[i]);
+        continue;
+      }
+
       if (memrefToCondition.count(memref)) {
         // If the dealloc conditions don't match, we need to make sure that the
         // dealloc happens on the union of cases.
         auto &newCond = newConditions[memrefToCondition[memref]];
         if (newCond != cond)
           newCond =
               rewriter.create<arith::OrIOp>(deallocOp.getLoc(), newCond, cond);
       } else {
if (newConditions.size() == deallocOp.getConditions().size() &&
         newRetained.size() == deallocOp.getRetained().size())
       return failure();
 
     // We need to create a new op because the number of results is always the
     // same as the number of condition operands.
     auto newDealloc = rewriter.create<DeallocOp>(deallocOp.getLoc(), newMemrefs,
                                                  newConditions, newRetained);
     for (auto [i, newIdx] : llvm::enumerate(resultIndices))
+      if (!deallocOp.getResult(i).use_empty())
         rewriter.replaceAllUsesWith(deallocOp.getResult(i),
                                     newDealloc.getResult(newIdx));
 
     rewriter.eraseOp(deallocOp);
     return success();
   }
 };
 
+struct EraseEmptyDealloc : public OpRewritePattern<DeallocOp> {
+  using OpRewritePattern<DeallocOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
+    if (deallocOp.getMemrefs().empty()) {
+      rewriter.eraseOp(deallocOp);
+      return success();
+    }
+    return failure();
+  }
+};
+
 } // anonymous namespace
 
 void DeallocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results.add<DeallocRemoveDuplicates>(context);
+  results.add<DeallocRemoveDuplicates, EraseEmptyDealloc>(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: @dealloc_canonicalize_clone_removal
 //   CHECK-NOT:   bufferization.clone
 //   CHECK-NOT:   memref.dealloc
 //       CHECK:   return {{.*}}
 
 // -----
 
-func.func @dealloc_canonicalize_duplicates(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>, %arg4: memref<2xi32>) {
-  %0:3 = bufferization.dealloc (%arg4, %arg0, %arg0 : memref<2xi32>, memref<2xi32>, memref<2xi32>) if (%arg1, %arg1, %arg1) retain (%arg3, %arg4, %arg3 : memref<2xi32>, memref<2xi32>, memref<2xi32>)
+func.func @dealloc_canonicalize_duplicates(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>, %arg4: memref<2xi32>, %arg5: memref<2xi32>) {
+  %0:3 = bufferization.dealloc (%arg4, %arg0, %arg0 : memref<2xi32>, memref<2xi32>, memref<2xi32>) if (%arg1, %arg1, %arg1) retain (%arg3, %arg5, %arg3 : memref<2xi32>, memref<2xi32>, memref<2xi32>)
   %1:2 = bufferization.dealloc (%arg0, %arg0 : memref<2xi32>, memref<2xi32>) if (%arg1, %arg2)
   return
 }
 
 // CHECK-LABEL: func @dealloc_canonicalize_duplicates
-//  CHECK-SAME:  ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xi32>, [[ARG4:%.+]]: memref<2xi32>)
-//  CHECK-NEXT:   bufferization.dealloc ([[ARG4]], [[ARG0]] : memref<2xi32>, memref<2xi32>) if ([[ARG1]], [[ARG1]]) retain ([[ARG3]], [[ARG4]] : memref<2xi32>, memref<2xi32>)
+//  CHECK-SAME:  ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xi32>, [[ARG4:%.+]]: memref<2xi32>, [[ARG5:%.+]]: memref<2xi32>)
+//  CHECK-NEXT:   bufferization.dealloc ([[ARG4]], [[ARG0]] : memref<2xi32>, memref<2xi32>) if ([[ARG1]], [[ARG1]]) retain ([[ARG3]], [[ARG5]] : memref<2xi32>, memref<2xi32>)
 //  CHECK-NEXT:   [[NEW_COND:%.+]] = arith.ori [[ARG1]], [[ARG2]] : i1
 //  CHECK-NEXT:   bufferization.dealloc ([[ARG0]] : memref<2xi32>) if ([[NEW_COND]])
+
+// -----
+
+func.func @dealloc_canonicalize_retained_and_deallocated(%arg0: memref<2xi32>, %arg1: i1, %arg2: memref<2xi32>) {
+  %0 = bufferization.dealloc (%arg0 : memref<2xi32>) if (%arg1) retain (%arg0 : memref<2xi32>)
+  %1:2 = bufferization.dealloc (%arg0, %arg2 : memref<2xi32>, memref<2xi32>) if (%arg1, %arg1) retain (%arg0 : memref<2xi32>)
+  bufferization.dealloc
+  bufferization.dealloc retain (%arg0 : memref<2xi32>)
+  return
+}
+
+// CHECK-LABEL: func @dealloc_canonicalize_retained_and_deallocated
+//  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: memref<2xi32>)
+//  CHECK-NEXT: bufferization.dealloc ([[ARG2]] : memref<2xi32>) if ([[ARG1]]) retain ([[ARG0]] : memref<2xi32>)
+//  CHECK-NEXT: return