[mlir][bufferization] Add pattern to BufferDeallocationSimplification pass that removes unnecessary retain values

[maerhart]

Summary

Adds a pattern that removes memrefs from the retained list which are
guaranteed to not alias any memref in the memrefs list. The corresponding
result value can be replaced with false in that case according to the
operation description.

When applied after BufferDeallocation, this can considerably reduce the
overhead that needs to be added during the lowering of the dealloc operation to
check for aliasing (especially when there is only one element in the memref
list and all retained values can be removed).

Depends on D157398

[springerm]

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp

188–190

Nit: I would put a lambda just after falseValue:

auto getOrCreateFalse = [&]() { if ... rewriter.create<artith::ConstantOp>(...) ...};

mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir

12

will this fold away with -canonicalize?

[maerhart]

mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir

12

Yes exactly, the canonicalizer removes such dealloc operations.

[maerhart]

Add lambda.

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp

LogicalResult matchAndRewrite(DeallocOp deallocOp,
     return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
                                   rewriter);
   }
 
 private:
   AliasAnalysis &aliasAnalysis;
 };
 
+/// Remove memrefs from the `retained` list which are guaranteed to not alias
+/// any memref in the `memrefs` list. The corresponding result value can be
+/// replaced with `false` in that case according to the operation description.
+///
+/// Example:
+/// ```mlir
+/// %0:2 = bufferization.dealloc (%m : memref<2xi32>) if (%cond)
+///                       retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
+/// return %0#0, %0#1
+/// ```
+/// can be canonicalized to the following given that `%r0` and `%r1` do not
+/// alias `%m`:
+/// ```mlir
+/// bufferization.dealloc (%m : memref<2xi32>) if (%cond)
+/// return %false, %false
+/// ```
+struct RemoveRetainedMemrefsGuaranteedToNotAlias
+    : public OpRewritePattern<DeallocOp> {
+  RemoveRetainedMemrefsGuaranteedToNotAlias(MLIRContext *context,
+                                            AliasAnalysis &aliasAnalysis)
+      : OpRewritePattern<DeallocOp>(context), aliasAnalysis(aliasAnalysis) {}
+
+  bool potentiallyAliasesMemref(DeallocOp deallocOp,
+                                Value retainedMemref) const {
+    for (auto memref : deallocOp.getMemrefs()) {
+      if (!aliasAnalysis.alias(memref, retainedMemref).isNo())
+        return true;
+    }
+    return false;
+  }
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> newRetainedMemrefs, replacements;
+    Value falseValue;
+    auto getOrCreateFalse = [&]() -> Value {
+      if (!falseValue)
+        falseValue = rewriter.create<arith::ConstantOp>(
+            deallocOp.getLoc(), rewriter.getBoolAttr(false));
+      return falseValue;
+    };
+
+    for (auto retainedMemref : deallocOp.getRetained()) {
+      if (potentiallyAliasesMemref(deallocOp, retainedMemref)) {
+        newRetainedMemrefs.push_back(retainedMemref);
+        replacements.push_back({});
+        continue;
+      }
+
+      replacements.push_back(getOrCreateFalse());
+    }
+
+    if (newRetainedMemrefs.size() == deallocOp.getRetained().size())
+      return failure();
+
+    auto newDeallocOp = rewriter.create<DeallocOp>(
+        deallocOp.getLoc(), deallocOp.getMemrefs(), deallocOp.getConditions(),
+        newRetainedMemrefs);
+    int i = 0;
+    for (auto &repl : replacements) {
+      if (!repl)
+        repl = newDeallocOp.getUpdatedConditions()[i++];
+    }
+
+    rewriter.replaceOp(deallocOp, replacements);
+    return success();
+  }
+
+private:
+  AliasAnalysis &aliasAnalysis;
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationSimplificationPass
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// The actual buffer deallocation pass that inserts and moves dealloc nodes
 /// into the right positions. Furthermore, it inserts additional clones if
 /// necessary. It uses the algorithm described at the top of the file.
 struct BufferDeallocationSimplificationPass
     : public bufferization::impl::BufferDeallocationSimplificationBase<
           BufferDeallocationSimplificationPass> {
   void runOnOperation() override {
     AliasAnalysis &aliasAnalysis = getAnalysis<AliasAnalysis>();
     RewritePatternSet patterns(&getContext());
-    patterns.add<DeallocRemoveDeallocMemrefsContainedInRetained>(&getContext(),
-                                                                 aliasAnalysis);
+    patterns.add<DeallocRemoveDeallocMemrefsContainedInRetained,
+                 RemoveRetainedMemrefsGuaranteedToNotAlias>(&getContext(),
 
     if (failed(
             applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
 std::unique_ptr<Pass>
 mlir::bufferization::createBufferDeallocationSimplificationPass() {
   return std::make_unique<BufferDeallocationSimplificationPass>();
 }

mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir

 // RUN: mlir-opt %s --buffer-deallocation-simplification --split-input-file | FileCheck %s
 
 func.func @dealloc_deallocated_in_retained(%arg0: memref<2xi32>, %arg1: i1, %arg2: memref<2xi32>) -> (i1, i1, i1, i1) {
   %0 = bufferization.dealloc (%arg0 : memref<2xi32>) if (%arg1) retain (%arg0 : memref<2xi32>)
   %1 = bufferization.dealloc (%arg0, %arg2 : memref<2xi32>, memref<2xi32>) if (%arg1, %arg1) retain (%arg0 : memref<2xi32>)
   %2:2 = bufferization.dealloc (%arg0 : memref<2xi32>) if (%arg1) retain (%arg0, %arg2 : memref<2xi32>, memref<2xi32>)
   return %0, %1, %2#0, %2#1 : i1, i1, i1, i1
 }
 
 // CHECK-LABEL: func @dealloc_deallocated_in_retained
 //  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: memref<2xi32>)
-//  CHECK-NEXT: [[V0:%.+]] = bufferization.dealloc retain ([[ARG0]] : memref<2xi32>)
-//  CHECK-NEXT: [[O0:%.+]] = arith.ori [[V0]], [[ARG1]]
+//  CHECK-NEXT: bufferization.dealloc
 //  CHECK-NEXT: [[V1:%.+]] = bufferization.dealloc ([[ARG2]] : memref<2xi32>) if ([[ARG1]]) retain ([[ARG0]] : memref<2xi32>)
 //  CHECK-NEXT: [[O1:%.+]] = arith.ori [[V1]], [[ARG1]]
 //  CHECK-NEXT: [[V2:%.+]]:2 = bufferization.dealloc ([[ARG0]] : memref<2xi32>) if ([[ARG1]]) retain ([[ARG0]], [[ARG2]] : memref<2xi32>, memref<2xi32>)
-//  CHECK-NEXT: return [[O0]], [[O1]], [[V2]]#0, [[V2]]#1 :
+//  CHECK-NEXT: return [[ARG1]], [[O1]], [[V2]]#0, [[V2]]#1 :
 
 // -----
 
 func.func @dealloc_deallocated_in_retained_extract_base_memref(%arg0: memref<2xi32>, %arg1: i1, %arg2: memref<2xi32>) -> (i1, i1, i1, i1) {
   %base_buffer, %offset, %size, %stride = memref.extract_strided_metadata %arg0 : memref<2xi32> -> memref<i32>, index, index, index
   %base_buffer0, %offset0, %size0, %stride0 = memref.extract_strided_metadata %arg2 : memref<2xi32> -> memref<i32>, index, index, index
   %0 = bufferization.dealloc (%base_buffer : memref<i32>) if (%arg1) retain (%arg0 : memref<2xi32>)
   %1 = bufferization.dealloc (%base_buffer, %base_buffer0 : memref<i32>, memref<i32>) if (%arg1, %arg1) retain (%arg0 : memref<2xi32>)
   %2:2 = bufferization.dealloc (%base_buffer : memref<i32>) if (%arg1) retain (%arg0, %arg2 : memref<2xi32>, memref<2xi32>)
   return %0, %1, %2#0, %2#1 : i1, i1, i1, i1
 }
 
 // CHECK-LABEL: func @dealloc_deallocated_in_retained_extract_base_memref
 //  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: memref<2xi32>)
 //  CHECK-NEXT: [[BASE0:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[ARG0]] :
 //  CHECK-NEXT: [[BASE1:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[ARG2]] :
-//  CHECK-NEXT: [[V0:%.+]] = bufferization.dealloc retain ([[ARG0]] : memref<2xi32>)
-//  CHECK-NEXT: [[O0:%.+]] = arith.ori [[V0]], [[ARG1]]
+//  CHECK-NEXT: bufferization.dealloc
 //  CHECK-NEXT: [[V1:%.+]] = bufferization.dealloc ([[BASE1]] : memref<i32>) if ([[ARG1]]) retain ([[ARG0]] : memref<2xi32>)
 //  CHECK-NEXT: [[O1:%.+]] = arith.ori [[V1]], [[ARG1]]
 //  CHECK-NEXT: [[V2:%.+]]:2 = bufferization.dealloc ([[BASE0]] : memref<i32>) if ([[ARG1]]) retain ([[ARG0]], [[ARG2]] : memref<2xi32>, memref<2xi32>)
-//  CHECK-NEXT: return [[O0]], [[O1]], [[V2]]#0, [[V2]]#1 :
+//  CHECK-NEXT: return [[ARG1]], [[O1]], [[V2]]#0, [[V2]]#1 :
+
+// -----
+
+func.func @dealloc_deallocated_in_retained(%arg0: i1, %arg1: memref<2xi32>) -> (i1, i1) {
+  %alloc = memref.alloc() : memref<2xi32>
+  %alloc0 = memref.alloc() : memref<2xi32>
+  %0:2 = bufferization.dealloc (%alloc : memref<2xi32>) if (%arg0) retain (%alloc0, %arg1 : memref<2xi32>, memref<2xi32>)
+  return %0#0, %0#1 : i1, i1
+}
+
+// CHECK-LABEL: func @dealloc_deallocated_in_retained
+//  CHECK-SAME: ([[ARG0:%.+]]: i1, [[ARG1:%.+]]: memref<2xi32>)
+//  CHECK-NEXT: [[FALSE:%.+]] = arith.constant false
+//  CHECK-NEXT: [[ALLOC:%.+]] = memref.alloc(
+//  CHECK-NEXT: bufferization.dealloc ([[ALLOC]] : memref<2xi32>) if ([[ARG0]])
+//  CHECK-NOT: retain
+//  CHECK-NEXT: return [[FALSE]], [[FALSE]] :