[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.

239

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

239

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();
 }
 
+/// Checks if 'memref' may or must alias a MemRef in 'memrefList'. It is often a
+/// requirement of optimization patterns that there cannot be any aliasing
+/// memref in order to perform the desired simplification. The 'allowSelfAlias'
+/// argument indicates whether 'memref' may be present in 'memrefList' which
+/// makes this helper function applicable to situations where we already know
+/// that 'memref' is in the list but also when we don't want it in the list.
+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 values from the `memref` operand list that are also present in the
 /// `retained` list since they will always alias and thus never actually be

 /// 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)) {
+      if (potentiallyAliasesMemref(aliasAnalysis, deallocOp.getMemrefs(),
+                                   retainedMemref, false)) {
         newRetainedMemrefs.push_back(retainedMemref);
         replacements.push_back({});
         continue;
       }
 
       replacements.push_back(getOrCreateFalse());
     }
 
LogicalResult matchAndRewrite(DeallocOp deallocOp,
     rewriter.replaceOp(deallocOp, replacements);
     return success();
   }
 
 private:
   AliasAnalysis &aliasAnalysis;
 };
 
+/// Split off memrefs to separate 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]] :
+// COM: there is only one value in the retained list because the
+// COM: RemoveRetainedMemrefsGuaranteedToNotAlias pattern also applies here and
+// COM: removes %arg1 from the list. In the second dealloc, this does not apply
+// COM: because function arguments are assumed potentially alias (even if the
+// COM: types don't exactly match).
+//  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]] :