[mlir][bufferization] Add pattern to BufferDeallocationSimplification pass

[maerhart]

Summary

Add a pattern that splits one dealloc operation into multiple dealloc operation
depending on static aliasing information of the values in the memref operand
list. This reduces the total number of aliasing checks required at runtime and
can enable futher canonicalizations of the new and simplified dealloc
operations.

Depends on D157407

[springerm]

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

51

Add comment. Also describe why allowSelfAlias is useful.

213

does it matter?

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

62

Does the analysis make any assumptions about the aliasing of two function bbArgs?

71

Can you describe in the comment why the retain list has only one element.

[maerhart]

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

213

No, just thought it makes it easier to understand that this adds additional dealloc operations, but I can also remove it.

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

62

It's may-alias by default (even if the types don't match). There is a restricted attribute that can be added to functions though, which indicates that no aliasing should be assumed.

71

This is the result of another pattern which removes unnecessary retained memrefs and because of that I did not add this simplification to the example above the pattern in the CPP file. I can add a comment to the test here though to make maintenance easier.
Might be worth it though to split the file first, then have a separate run command for each pattern such that adding a new pattern does not affect already present patterns.

[springerm]

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

71

Just add a comment to the test case that describes the two optimizations that are happening here. That's good enough.

[maerhart]

Address comments.

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

static LogicalResult updateDeallocIfChanged(DeallocOp deallocOp,
 
   rewriter.updateRootInPlace(deallocOp, [&]() {
     deallocOp.getMemrefsMutable().assign(memrefs);
     deallocOp.getConditionsMutable().assign(conditions);
   });
   return success();
 }
 
+static bool potentiallyAliasesMemref(AliasAnalysis &analysis,
+                                     ValueRange memrefList, Value memref,
+                                     bool allowSelfAlias) {
+  for (auto mr : memrefList) {
+    if (allowSelfAlias && mr == memref)
+      continue;
+    if (!analysis.alias(mr, memref).isNo())
+      return true;
+  }
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // Patterns
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Remove memrefs to be deallocated that are also present in the retained list
 /// since they will always alias and thus never actually be deallocated.

 /// 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;
     for (auto retainedMemref : deallocOp.getRetained()) {
-      if (potentiallyAliasesMemref(deallocOp, retainedMemref)) {
+      if (potentiallyAliasesMemref(aliasAnalysis, deallocOp.getMemrefs(),
+                                   retainedMemref, false)) {
         newRetainedMemrefs.push_back(retainedMemref);
         replacements.push_back({});
         continue;
       }
 
       if (!falseValue)
         falseValue = rewriter.create<arith::ConstantOp>(
             deallocOp.getLoc(), rewriter.getBoolAttr(false));
LogicalResult matchAndRewrite(DeallocOp deallocOp,
     rewriter.replaceOp(deallocOp, replacements);
     return success();
   }
 
 private:
   AliasAnalysis &aliasAnalysis;
 };
 
+/// Split off memrefs to separate (newly created) dealloc operations to reduce
+/// the number of runtime checks required and enable further canonicalization of
+/// the new and simpler dealloc operations. A memref can be split off if it is
+/// guaranteed to not alias with any other memref in the `memref` operand list.
+/// The results of the old and the new dealloc operation have to be combined by
+/// computing the element-wise disjunction of them.
+///
+/// Example:
+/// ```mlir
+/// %0:2 = bufferization.dealloc (%m0, %m1 : memref<2xi32>, memref<2xi32>)
+///                           if (%cond0, %cond1)
+///                       retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
+/// return %0#0, %0#1
+/// ```
+/// Given that `%m0` is guaranteed to never alias with `%m1`, the above IR is
+/// canonicalized to the following, thus reducing the number of runtime alias
+/// checks by 1 and potentially enabling further canonicalization of the new
+/// split-up dealloc operations.
+/// ```mlir
+/// %0:2 = bufferization.dealloc (%m0 : memref<2xi32>) if (%cond0)
+///                       retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
+/// %1:2 = bufferization.dealloc (%m1 : memref<2xi32>) if (%cond1)
+///                       retain (%r0, %r1 : memref<2xi32>, memref<2xi32>)
+/// %2 = arith.ori %0#0, %1#0
+/// %3 = arith.ori %0#1, %1#1
+/// return %2, %3
+/// ```
+struct SplitDeallocWhenNotAliasingAnyOther
+    : public OpRewritePattern<DeallocOp> {
+  SplitDeallocWhenNotAliasingAnyOther(MLIRContext *context,
+                                      AliasAnalysis &aliasAnalysis)
+      : OpRewritePattern<DeallocOp>(context), aliasAnalysis(aliasAnalysis) {}
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
+    if (deallocOp.getMemrefs().size() <= 1)
+      return failure();
+
+    SmallVector<Value> newMemrefs, newConditions, replacements;
+    DenseSet<Operation *> exceptedUsers;
+    replacements = deallocOp.getUpdatedConditions();
+    for (auto [memref, cond] :
+         llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+      if (potentiallyAliasesMemref(aliasAnalysis, deallocOp.getMemrefs(),
+                                   memref, true)) {
+        newMemrefs.push_back(memref);
+        newConditions.push_back(cond);
+        continue;
+      }
+
+      auto newDeallocOp = rewriter.create<DeallocOp>(
+          deallocOp.getLoc(), memref, cond, deallocOp.getRetained());
+      replacements = SmallVector<Value>(llvm::map_range(
+          llvm::zip(replacements, newDeallocOp.getUpdatedConditions()),
+          [&](auto replAndNew) -> Value {
+            auto orOp = rewriter.create<arith::OrIOp>(deallocOp.getLoc(),
+                                                      std::get<0>(replAndNew),
+                                                      std::get<1>(replAndNew));
+            exceptedUsers.insert(orOp);
+            return orOp.getResult();
+          }));
+    }
+
+    if (newMemrefs.size() == deallocOp.getMemrefs().size())
+      return failure();
+
+    rewriter.replaceUsesWithIf(deallocOp.getUpdatedConditions(), replacements,
+                               [&](OpOperand &operand) {
+                                 return !exceptedUsers.contains(
+                                     operand.getOwner());
+                               });
+    return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
+                                  rewriter);
+  }
+
+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,
-                 RemoveRetainedMemrefsGuaranteedToNotAlias>(&getContext(),
-                                                            aliasAnalysis);
+                 RemoveRetainedMemrefsGuaranteedToNotAlias,
+                 SplitDeallocWhenNotAliasingAnyOther>(&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

 
 // 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]] :
+
+// -----
+
+func.func @dealloc_split_when_no_other_aliasing(%arg0: i1, %arg1: memref<2xi32>, %arg2: memref<2xi32>, %arg3: i1) -> (i1, i1) {
+  %alloc = memref.alloc() : memref<2xi32>
+  %alloc0 = memref.alloc() : memref<2xi32>
+  %0 = arith.select %arg0, %alloc, %alloc0 : memref<2xi32>
+  %1:2 = bufferization.dealloc (%alloc, %arg2 : memref<2xi32>, memref<2xi32>) if (%arg0, %arg3) retain (%arg1, %0 : memref<2xi32>, memref<2xi32>)
+  return %1#0, %1#1 : i1, i1
+}
+
+// CHECK-LABEL: func @dealloc_split_when_no_other_aliasing
+//  CHECK-SAME: ([[ARG0:%.+]]: i1, [[ARG1:%.+]]: memref<2xi32>, [[ARG2:%.+]]: memref<2xi32>, [[ARG3:%.+]]: i1)
+//  CHECK-NEXT:   [[ALLOC0:%.+]] = memref.alloc(
+//  CHECK-NEXT:   [[ALLOC1:%.+]] = memref.alloc(
+//  CHECK-NEXT:   [[V0:%.+]] = arith.select{{.*}}[[ALLOC0]], [[ALLOC1]] :
+//  CHECK-NEXT:   [[V1:%.+]] = bufferization.dealloc ([[ALLOC0]] : memref<2xi32>) if ([[ARG0]]) retain ([[V0]] : memref<2xi32>)
+//  CHECK-NEXT:   [[V2:%.+]]:2 = bufferization.dealloc ([[ARG2]] : memref<2xi32>) if ([[ARG3]]) retain ([[ARG1]], [[V0]] : memref<2xi32>, memref<2xi32>)
+//  CHECK-NEXT:   [[V3:%.+]] = arith.ori [[V1]], [[V2]]#1
+//  CHECK-NEXT:   bufferization.dealloc
+//  CHECK-NEXT:   return [[V2]]#0, [[V3]] :