[mlir][bufferization] Simplify helper potentiallyAliasesMemref (#78690)

This commit simplifies a helper function in the ownership-based buffer deallocation pass. Fixes a potential double-free (depending on the scheduling of patterns).

Author: matthias-springer

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

@@ -77,18 +77,12 @@ static bool distinctAllocAndBlockArgument(Value v1, Value v2) {
  return areDistinct(v1Base, v2Base) || areDistinct(v2Base, v1Base);
 }
 
-/// Checks if `memref` may or must alias a MemRef in `otherList`. 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 `otherList` 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.
+/// Checks if `memref` may potentially alias a MemRef in `otherList`. It is
+/// often a requirement of optimization patterns that there cannot be any
+/// aliasing memref in order to perform the desired simplification.
 static bool potentiallyAliasesMemref(AliasAnalysis &analysis,
- ValueRange otherList, Value memref,
- bool allowSelfAlias) {
+ ValueRange otherList, Value memref) {
  for (auto other : otherList) {
- if (allowSelfAlias && other == memref)
- continue;
  if (distinctAllocAndBlockArgument(other, memref))
  continue;
  if (!analysis.alias(other, memref).isNo())
@@ -243,7 +237,7 @@ struct RemoveRetainedMemrefsGuaranteedToNotAlias
 
  for (auto retainedMemref : deallocOp.getRetained()) {
  if (potentiallyAliasesMemref(aliasAnalysis, deallocOp.getMemrefs(),
- retainedMemref, false)) {
+ retainedMemref)) {
  newRetainedMemrefs.push_back(retainedMemref);
  replacements.push_back({});
  continue;
@@ -314,11 +308,13 @@ struct SplitDeallocWhenNotAliasingAnyOther
 
  SmallVector<Value> remainingMemrefs, remainingConditions;
  SmallVector<SmallVector<Value>> updatedConditions;
- for (auto [memref, cond] :
- llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+ for (int64_t i = 0, e = deallocOp.getMemrefs().size(); i < e; ++i) {
+ Value memref = deallocOp.getMemrefs()[i];
+ Value cond = deallocOp.getConditions()[i];
+ SmallVector<Value> otherMemrefs(deallocOp.getMemrefs());
+ otherMemrefs.erase(otherMemrefs.begin() + i);
  // Check if `memref` can split off into a separate bufferization.dealloc.
- if (potentiallyAliasesMemref(aliasAnalysis, deallocOp.getMemrefs(),
- memref, true)) {
+ if (potentiallyAliasesMemref(aliasAnalysis, otherMemrefs, memref)) {
  // `memref` alias with other memrefs, do not split off.
  remainingMemrefs.push_back(memref);
  remainingConditions.push_back(cond);

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

@@ -160,3 +160,14 @@ func.func @alloc_and_bbarg(%arg0: memref<5xf32>, %arg1: index, %arg2: index, %ar
 // CHECK: bufferization.dealloc (%[[view]] : memref<f32>)
 // CHECK-NOT: retain
 // CHECK: scf.yield %[[alloc]], %[[true]]
+
+// -----
+
+func.func @duplicate_memref(%arg0: memref<5xf32>, %arg1: memref<6xf32>, %c: i1) -> i1 {
+ %0 = bufferization.dealloc (%arg0, %arg0 : memref<5xf32>, memref<5xf32>) if (%c, %c) retain (%arg1 : memref<6xf32>)
+ return %0 : i1
+}
+
+// CHECK-LABEL: func @duplicate_memref(
+// CHECK: %[[r:.*]] = bufferization.dealloc (%{{.*}} : memref<5xf32>) if (%{{.*}}) retain (%{{.*}} : memref<6xf32>)
+// CHECK: return %[[r]]