[mlir][bufferization] Add DeallocOp canonicalizer to remove duplicate values

[maerhart]

Summary

Duplicate values in the retained list can just be removed, however, for duplicates in the list of memrefs to deallocate, we also need to check the conditions and if thhey don't match, we need to compute the OR in order to not miss a case leading to a memory leak.

[maerhart]

Make sure the arith dialect is registered since the canonicalizer builds arith.ori operations in some situations.

[springerm]

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

771

Let's put a comment with a small example here.

785

We usually spell out the type unless there's a cast/dyn_cast in the assignment.

[maerhart]

Address comments

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td

let description = [{
     bufferization passes.
 
     An overview of the bufferization infrastructure and important conceptual
     details related to using the MLIR dialect conversion infrastructure can be
     found in [bufferization](/docs/Bufferization/) and [buffer
     deallocation](/docs/BufferDeallocationInternals/).
   }];
   let dependentDialects = [
-    "affine::AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect"
+    "affine::AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect",
+    "arith::ArithDialect"
   ];
 
   let extraClassDeclaration = [{
     /// Verify an attribute from this dialect on the argument at 'argIndex' for
     /// the region at 'regionIndex' on the given operation. Returns failure if
     /// the verification failed, success otherwise. This hook may optionally be
     /// invoked from any operation containing a region.
     LogicalResult verifyRegionArgAttribute(Operation *,

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

def Bufferization_DeallocOp : Bufferization_Op<"dealloc", [
   let results = (outs Variadic<I1>:$updatedConditions);
 
   let assemblyFormat = [{
     (` ``(` $memrefs^ `:` type($memrefs) `)` `if` ` ` `(` $conditions `)` )?
     (`retain` ` ` `(` $retained^ `:` type($retained) `)` )? attr-dict
   }];
 
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 #endif // BUFFERIZATION_OPS

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

 
 LogicalResult DeallocOp::verify() {
   if (getMemrefs().size() != getConditions().size())
     return emitOpError(
         "must have the same number of conditions as memrefs to deallocate");
   return success();
 }
 
+namespace {
+
+struct DeallocRemoveDuplicates : public OpRewritePattern<DeallocOp> {
+  using OpRewritePattern<DeallocOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
+    // Unique memrefs to be deallocated.
+    DenseMap<Value, unsigned> memrefToCondition;
+    SmallVector<Value> newMemrefs, newConditions, newRetained;
+    SmallVector<unsigned> resultIndices;
+    for (auto [memref, cond] :
+         llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+      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 {
+        memrefToCondition.insert({memref, newConditions.size()});
+        newMemrefs.push_back(memref);
+        newConditions.push_back(cond);
+      }
+      resultIndices.push_back(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))
+      rewriter.replaceAllUsesWith(deallocOp.getResult(i),
+                                  newDealloc.getResult(newIdx));
+
+    rewriter.eraseOp(deallocOp);
+    return success();
+  }
+};
+
+} // anonymous namespace
+
+void DeallocOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                            MLIRContext *context) {
+  results.add<DeallocRemoveDuplicates>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"

mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRBufferizationDialect
   AllocationOpInterface.cpp
   BufferizableOpInterface.cpp
   BufferizationOps.cpp
   BufferizationDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 
   DEPENDS
   MLIRAllocationOpInterfaceIncGen
   MLIRBufferizationOpsIncGen
   MLIRBufferizationEnumsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
+  MLIRArithDialect
   MLIRDestinationStyleOpInterface
   MLIRDialect
   MLIRFuncDialect
   MLIRIR
   MLIRSparseTensorDialect
   MLIRTensorDialect
   MLIRMemRefDialect
   )

mlir/test/Dialect/Bufferization/canonicalize.mlir

func.func @dealloc_canonicalize_clone_removal(%arg0: memref<?xindex>) -> memref<*xf32> {
   %2 = bufferization.clone %1 : memref<*xf32> to memref<*xf32>
   memref.dealloc %0 : memref<?xf32>
   return %2 : memref<*xf32>
 }
 // 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>)
+  %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-NEXT:   [[NEW_COND:%.+]] = arith.ori [[ARG1]], [[ARG2]] : i1
+//  CHECK-NEXT:   bufferization.dealloc ([[ARG0]] : memref<2xi32>) if ([[NEW_COND]])