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