[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

802

Should this be memref?

842–844

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

802

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.

842–844

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

 ///                     retain (%arg3 : memref<2xi32>)
 /// ```
 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())) {
+    SmallVector<int32_t> resultIndices(deallocOp.getMemrefs().size(), -1);
+    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.
         Value &newCond = newConditions[memrefToCondition[memref]];
         if (newCond != cond)
           newCond =
               rewriter.create<arith::OrIOp>(deallocOp.getLoc(), newCond, cond);
       } else {
         memrefToCondition.insert({memref, newConditions.size()});
         newMemrefs.push_back(memref);
         newConditions.push_back(cond);
       }
-      resultIndices.push_back(memrefToCondition[memref]);
+      resultIndices[i] = memrefToCondition[memref];
     }
 
     // Unique retained values
     DenseSet<Value> seen;
     for (auto retained : deallocOp.getRetained()) {
       if (!seen.contains(retained)) {
         seen.insert(retained);
         newRetained.push_back(retained);
       }
     }
 
     // Return failure if we don't change anything such that we don't run into an
     // infinite loop of pattern applications.
     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 (newIdx != -1)
         rewriter.replaceAllUsesWith(deallocOp.getResult(i),
                                     newDealloc.getResult(newIdx));
 
     rewriter.eraseOp(deallocOp);
     return success();
   }
 };
 
+/// Erase deallocation operations where the variadic list of memrefs to
+/// deallocate is emtpy. Example:
+/// ```mlir
+/// bufferization.dealloc retain (%arg0: memref<2xi32>)
+/// ```
+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>) -> (i1, i1, i1, i1, i1) {
-  %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>) -> (i1, i1, i1, i1, i1) {
+  %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 %0#0, %0#1, %0#2, %1#0, %1#1 : i1, i1, i1, i1, i1
 }
 
 // CHECK-LABEL: func @dealloc_canonicalize_duplicates
-//  CHECK-SAME:  ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xi32>, [[ARG4:%.+]]: memref<2xi32>)
-//  CHECK-NEXT:   [[V0:%.+]]:2 = 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:   [[V0:%.+]]:2 = 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:   [[V1:%.+]] = bufferization.dealloc ([[ARG0]] : memref<2xi32>) if ([[NEW_COND]])
 //  CHECK-NEXT:   return [[V0]]#0, [[V0]]#1, [[V0]]#1, [[V1]], [[V1]] :
+
+// -----
+
+func.func @dealloc_canonicalize_retained_and_deallocated(%arg0: memref<2xi32>, %arg1: i1, %arg2: memref<2xi32>) -> (i1, i1, i1) {
+  %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 %0, %1#0, %1#1 : i1, i1, i1
+}
+
+// CHECK-LABEL: func @dealloc_canonicalize_retained_and_deallocated
+//  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: memref<2xi32>)
+//  CHECK-NEXT: [[V0:%.+]] = bufferization.dealloc ([[ARG2]] : memref<2xi32>) if ([[ARG1]]) retain ([[ARG0]] : memref<2xi32>)
+//  CHECK-NEXT: return [[ARG1]], [[ARG1]], [[V0]] :