[mlir][bufferization] Change semantics of DeallocOp result values

This change allows supporting operations for which we don't get precise aliasing information without the need to insert clone operations. E.g., `arith.select`.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D156992

Author: maerhart

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

@@ -485,22 +485,41 @@ def Bufferization_DeallocOp : Bufferization_Op<"dealloc", [
     deallocating that memref). If two memrefs alias each other, only one will be
     deallocated to avoid double free situations.
 
-    The memrefs to be deallocated must be the originally allocated memrefs,
-    however, the memrefs to be retained may be arbitrary memrefs.
-
-    Returns a list of conditions corresponding to the list of memrefs which
-    indicates the new ownerships, i.e., if the memref was deallocated the
-    ownership was dropped (set to 'false') and otherwise will be the same as the
-    input condition.
+    The number of variadic `memref` operands (the memrefs to be deallocated)
+    must equal the number of variadic `condition` operands and correspond to
+    each other element-wise.
+
+    The `memref` operands must be the originally allocated memrefs, however, the
+    `retained` memref operands may be arbitrary memrefs.
+
+    This operation returns a variadic number of `updatedConditions` operands,
+    one updated condition per retained memref. An updated condition indicates
+    the ownership of the respective retained memref. It is computed as the
+    disjunction of all `conditions` operands where the corresponding to
+    `memrefs` operand aliases with the retained memref. If the retained memref
+    has no aliases among `memrefs`, the resulting updated condition is 'false'.
+    This is because all memrefs that need to be deallocated within one basic
+    block should be added to the same `bufferization.dealloc` operation at the
+    end of the block; if no aliasing memref is present, then it does not have to
+    be deallocated and thus we don't need to claim ownership. If the memrefs to
+    be deallocated are split over multiple dealloc operations (e.g., to avoid
+    aliasing checks at runtime between the `memref` operands), then the results
+    have to be manually combined using an `arith.ori` operation and all of them
+    still require the same list of `retained` memref operands unless the
+    (potentially empty) set of aliasing memrefs can be determined statically. In
+    that case, the `updatedCondition` operand can be replaced accordingly (e.g.,
+    by a canonicalizer).
 
     Example:
     ```mlir
-    %0:2 = bufferization.dealloc %a0, %a1 if %cond0, %cond1 retain %r0, %r1 :
-      memref<2xf32>, memref<4xi32> retain memref<?xf32>, memref<f64>
+    %0:3 = bufferization.dealloc (%a0, %a1 : memref<2xf32>, memref<4xi32>)
+      if (%cond0, %cond1) retain (%r0, %r1, %r2 : memref<?xf32>, memref<f64>,
+      memref<2xi32>)
     ```
-    Deallocation will be called on `%a0` if `%cond0` is 'true' and neither `%r0`
-    or `%r1` are aliases of `%a0`. `%a1` will be deallocated when `%cond1` is
-    set to 'true' and none of `%r0`, %r1` and `%a0` are aliases.
+    Deallocation will be called on `%a0` if `%cond0` is 'true' and neither
+    `%r0`, `%r1`, or `%r2` are aliases of `%a0`. `%a1` will be deallocated when
+    `%cond1` is set to 'true' and none of `%r0`, %r1`, `%r2`, and `%a0` are
+    aliases.
   }];
 
   let arguments = (ins Variadic<AnyRankedOrUnrankedMemRef>:$memrefs,

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

@@ -755,7 +755,8 @@ LogicalResult DeallocOp::inferReturnTypes(
     ValueRange operands, DictionaryAttr attributes, OpaqueProperties properties,
     RegionRange regions, SmallVectorImpl<Type> &inferredReturnTypes) {
   DeallocOpAdaptor adaptor(operands, attributes, properties, regions);
-  inferredReturnTypes = SmallVector<Type>(adaptor.getConditions().getTypes());
+  inferredReturnTypes = SmallVector<Type>(adaptor.getRetained().size(),
+                                          IntegerType::get(context, 1));
   return success();
 }
 
@@ -766,44 +767,46 @@ LogicalResult DeallocOp::verify() {
   return success();
 }
 
+static LogicalResult updateDeallocIfChanged(DeallocOp deallocOp,
+                                            ArrayRef<Value> memrefs,
+                                            ArrayRef<Value> conditions,
+                                            PatternRewriter &rewriter) {
+  if (deallocOp.getMemrefs() == memrefs)
+    return failure();
+
+  rewriter.updateRootInPlace(deallocOp, [&]() {
+    deallocOp.getMemrefsMutable().assign(memrefs);
+    deallocOp.getConditionsMutable().assign(conditions);
+  });
+  return success();
+}
+
 namespace {
 
-/// Remove duplicate values in the list of retained memrefs as well as the list
-/// of memrefs to be deallocated. For the latter, we need to make sure the
-/// corresponding condition values match as well, or otherwise have to combine
-/// them (by computing the disjunction of them).
+/// Remove duplicate values in the list of memrefs to be deallocated. We need to
+/// make sure the corresponding condition value is updated accordingly since
+/// their two conditions might not cover the same set of cases. In that case, we
+/// have to combine them (by computing the disjunction of them).
 /// Example:
 /// ```mlir
-/// %0:2 = bufferization.dealloc (%arg0, %arg0 : ...)
-///                           if (%arg1, %arg2)
-///                       retain (%arg3, %arg3 : ...)
+/// bufferization.dealloc (%arg0, %arg0 : ...) if (%arg1, %arg2)
 /// ```
 /// is canonicalized to
 /// ```mlir
 /// %0 = arith.ori %arg1, %arg2 : i1
-/// %1 = bufferization.dealloc (%arg0 : memref<2xi32>)
-///                         if (%0)
-///                     retain (%arg3 : memref<2xi32>)
+/// bufferization.dealloc (%arg0 : memref<2xi32>) if (%0)
 /// ```
-struct DeallocRemoveDuplicates : public OpRewritePattern<DeallocOp> {
+struct DeallocRemoveDuplicateDeallocMemrefs
+    : public OpRewritePattern<DeallocOp> {
   using OpRewritePattern<DeallocOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(DeallocOp deallocOp,
                                 PatternRewriter &rewriter) const override {
     // Unique memrefs to be deallocated.
-    DenseSet<Value> retained(deallocOp.getRetained().begin(),
-                             deallocOp.getRetained().end());
     DenseMap<Value, unsigned> memrefToCondition;
-    SmallVector<Value> newMemrefs, newConditions, newRetained;
-    SmallVector<int32_t> resultIndices(deallocOp.getMemrefs().size(), -1);
+    SmallVector<Value> newMemrefs, newConditions;
     for (auto [i, memref, cond] :
          llvm::enumerate(deallocOp.getMemrefs(), deallocOp.getConditions())) {
-      if (retained.contains(memref)) {
-        rewriter.replaceAllUsesWith(deallocOp.getResult(i),
-                                    deallocOp.getConditions()[i]);
-        continue;
-      }
-
       if (memrefToCondition.count(memref)) {
         // If the dealloc conditions don't match, we need to make sure that the
         // dealloc happens on the union of cases.
@@ -816,50 +819,133 @@ struct DeallocRemoveDuplicates : public OpRewritePattern<DeallocOp> {
         newMemrefs.push_back(memref);
         newConditions.push_back(cond);
       }
-      resultIndices[i] = memrefToCondition[memref];
     }
 
+    // Return failure if we don't change anything such that we don't run into an
+    // infinite loop of pattern applications.
+    return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
+                                  rewriter);
+  }
+};
+
+/// Remove duplicate values in the list of retained memrefs. We need to make
+/// sure the corresponding result condition value is replaced properly.
+/// Example:
+/// ```mlir
+/// %0:2 = bufferization.dealloc retain (%arg3, %arg3 : ...)
+/// ```
+/// is canonicalized to
+/// ```mlir
+/// %0 = bufferization.dealloc retain (%arg3 : memref<2xi32>)
+/// ```
+struct DeallocRemoveDuplicateRetainedMemrefs
+    : public OpRewritePattern<DeallocOp> {
+  using OpRewritePattern<DeallocOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
     // Unique retained values
-    DenseSet<Value> seen;
+    DenseMap<Value, unsigned> seen;
+    SmallVector<Value> newRetained;
+    SmallVector<unsigned> resultReplacementIdx;
+    unsigned i = 0;
     for (auto retained : deallocOp.getRetained()) {
-      if (!seen.contains(retained)) {
-        seen.insert(retained);
-        newRetained.push_back(retained);
+      if (seen.count(retained)) {
+        resultReplacementIdx.push_back(seen[retained]);
+        continue;
       }
+
+      seen[retained] = i;
+      newRetained.push_back(retained);
+      resultReplacementIdx.push_back(i++);
     }
 
     // Return failure if we don't change anything such that we don't run into an
     // infinite loop of pattern applications.
-    if (newConditions.size() == deallocOp.getConditions().size() &&
-        newRetained.size() == deallocOp.getRetained().size())
+    if (newRetained.size() == deallocOp.getRetained().size())
       return failure();
 
     // We need to create a new op because the number of results is always the
     // same as the number of condition operands.
-    auto newDealloc = rewriter.create<DeallocOp>(deallocOp.getLoc(), newMemrefs,
-                                                 newConditions, newRetained);
-    for (auto [i, newIdx] : llvm::enumerate(resultIndices))
-      if (newIdx != -1)
-        rewriter.replaceAllUsesWith(deallocOp.getResult(i),
-                                    newDealloc.getResult(newIdx));
-
-    rewriter.eraseOp(deallocOp);
+    auto newDeallocOp =
+        rewriter.create<DeallocOp>(deallocOp.getLoc(), deallocOp.getMemrefs(),
+                                   deallocOp.getConditions(), newRetained);
+    SmallVector<Value> replacements(
+        llvm::map_range(resultReplacementIdx, [&](unsigned idx) {
+          return newDeallocOp.getUpdatedConditions()[idx];
+        }));
+    rewriter.replaceOp(deallocOp, replacements);
     return success();
   }
 };
 
+/// Remove memrefs to be deallocated that are also present in the retained list
+/// since they will always alias and thus never actually be deallocated.
+/// Example:
+/// ```mlir
+/// %0 = bufferization.dealloc (%arg0 : ...) if (%arg1) retain (%arg0 : ...)
+/// ```
+/// is canonicalized to
+/// ```mlir
+/// %0 = bufferization.dealloc retain (%arg0 : ...)
+/// ```
+struct DeallocRemoveDeallocMemrefsContainedInRetained
+    : public OpRewritePattern<DeallocOp> {
+  using OpRewritePattern<DeallocOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
+    // Unique memrefs to be deallocated.
+    DenseMap<Value, unsigned> retained;
+    for (auto [i, ret] : llvm::enumerate(deallocOp.getRetained()))
+      retained[ret] = i;
+
+    // There must not be any duplicates in the retain list anymore because we
+    // would miss updating one of the result values otherwise.
+    if (retained.size() != deallocOp.getRetained().size())
+      return failure();
+
+    SmallVector<Value> newMemrefs, newConditions;
+    for (auto [memref, cond] :
+         llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+      if (retained.contains(memref)) {
+        rewriter.setInsertionPointAfter(deallocOp);
+        auto orOp = rewriter.create<arith::OrIOp>(
+            deallocOp.getLoc(),
+            deallocOp.getUpdatedConditions()[retained[memref]], cond);
+        rewriter.replaceAllUsesExcept(
+            deallocOp.getUpdatedConditions()[retained[memref]],
+            orOp.getResult(), orOp);
+        continue;
+      }
+
+      newMemrefs.push_back(memref);
+      newConditions.push_back(cond);
+    }
+
+    // Return failure if we don't change anything such that we don't run into an
+    // infinite loop of pattern applications.
+    return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
+                                  rewriter);
+  }
+};
+
 /// Erase deallocation operations where the variadic list of memrefs to
-/// deallocate is emtpy. Example:
+/// deallocate is empty. Example:
 /// ```mlir
-/// bufferization.dealloc retain (%arg0: memref<2xi32>)
+/// %0 = bufferization.dealloc retain (%arg0: memref<2xi32>)
 /// ```
 struct EraseEmptyDealloc : public OpRewritePattern<DeallocOp> {
   using OpRewritePattern<DeallocOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(DeallocOp deallocOp,
                                 PatternRewriter &rewriter) const override {
     if (deallocOp.getMemrefs().empty()) {
-      rewriter.eraseOp(deallocOp);
+      Value constFalse = rewriter.create<arith::ConstantOp>(
+          deallocOp.getLoc(), rewriter.getBoolAttr(false));
+      rewriter.replaceOp(
+          deallocOp, SmallVector<Value>(deallocOp.getUpdatedConditions().size(),
+                                        constFalse));
       return success();
     }
     return failure();
@@ -871,55 +957,40 @@ struct EraseEmptyDealloc : public OpRewritePattern<DeallocOp> {
 ///
 /// Example:
 /// ```
-/// %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xi32>, memref<2xi32>)
+/// bufferization.dealloc (%arg0, %arg1 : memref<2xi32>, memref<2xi32>)
 ///                           if (%arg2, %false)
 /// ```
 /// becomes
 /// ```
-/// %0 = bufferization.dealloc (%arg0 : memref<2xi32>) if (%arg2)
+/// bufferization.dealloc (%arg0 : memref<2xi32>) if (%arg2)
 /// ```
 struct EraseAlwaysFalseDealloc : public OpRewritePattern<DeallocOp> {
   using OpRewritePattern<DeallocOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(DeallocOp deallocOp,
                                 PatternRewriter &rewriter) const override {
     SmallVector<Value> newMemrefs, newConditions;
-    SmallVector<Value> replacements;
-
-    for (auto [res, memref, cond] :
-         llvm::zip(deallocOp.getUpdatedConditions(), deallocOp.getMemrefs(),
-                   deallocOp.getConditions())) {
-      if (matchPattern(cond, m_Zero())) {
-        replacements.push_back(cond);
-        continue;
+    for (auto [memref, cond] :
+         llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+      if (!matchPattern(cond, m_Zero())) {
+        newMemrefs.push_back(memref);
+        newConditions.push_back(cond);
       }
-      newMemrefs.push_back(memref);
-      newConditions.push_back(cond);
-      replacements.push_back({});
     }
 
-    if (newMemrefs.size() == deallocOp.getMemrefs().size())
-      return failure();
-
-    auto newDeallocOp = rewriter.create<DeallocOp>(
-        deallocOp.getLoc(), newMemrefs, newConditions, deallocOp.getRetained());
-    unsigned i = 0;
-    for (auto &repl : replacements)
-      if (!repl)
-        repl = newDeallocOp.getResult(i++);
-
-    rewriter.replaceOp(deallocOp, replacements);
-    return success();
+    return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
+                                  rewriter);
   }
 };
 
 } // anonymous namespace
 
 void DeallocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results
-      .add<DeallocRemoveDuplicates, EraseEmptyDealloc, EraseAlwaysFalseDealloc>(
-          context);
+  results.add<DeallocRemoveDuplicateDeallocMemrefs,
+              DeallocRemoveDuplicateRetainedMemrefs,
+              DeallocRemoveDeallocMemrefsContainedInRetained, EraseEmptyDealloc,
+              EraseAlwaysFalseDealloc>(context);
 }
 
 //===----------------------------------------------------------------------===//