Revert "[mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI (#66626)"

This reverts commit aa9eb47.
It introduced a double free in a test case. Reverting to have some time
for fixing this and relanding later.

Author: maerhart

mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h

@@ -95,20 +95,7 @@ struct DeallocationOptions {
   // A pass option indicating whether private functions should be modified to
   // pass the ownership of MemRef values instead of adhering to the function
   // boundary ABI.
-  bool privateFuncDynamicOwnership = true;
-
-  /// Inserts `cf.assert` operations to verify the function boundary ABI at
-  /// runtime. Currently, it is only checked that the ownership of returned
-  /// MemRefs is 'true'. This also ensures that the returned memref does not
-  /// originate from the same allocation as a function argument.
-  /// Note: The function boundary ABI is disabled for non-external private
-  /// functions if `privateFuncDynamicOwnership` is enabled and thus this option
-  /// does not apply to them.
-  /// TODO: check that returned MemRefs don't alias each other.
-  /// If it can be determined statically that the ABI is not adhered
-  /// to, an error will already be emitted at compile time. This cannot be
-  /// changed with this option.
-  bool verifyFunctionBoundaryABI = true;
+  bool privateFuncDynamicOwnership = false;
 };
 
 /// This class collects all the state that we need to perform the buffer
@@ -151,12 +138,12 @@ class DeallocationState {
   void getLiveMemrefsIn(Block *block, SmallVectorImpl<Value> &memrefs);
 
   /// Given an SSA value of MemRef type, this function queries the ownership and
-  /// if it is not already in the 'Unique' state, potentially inserts IR to
-  /// determine the ownership (which might involve expensive aliasing checks at
-  /// runtime).
-  Value materializeMemRefOwnership(const DeallocationOptions &options,
-                                   OpBuilder &builder, Value memref,
-                                   Block *block);
+  /// if it is not already in the 'Unique' state, potentially inserts IR to get
+  /// a new SSA value, returned as the first element of the pair, which has
+  /// 'Unique' ownership and can be used instead of the passed Value with the
+  /// the ownership indicator returned as the second element of the pair.
+  std::pair<Value, Value>
+  getMemrefWithUniqueOwnership(OpBuilder &builder, Value memref, Block *block);
 
   /// Given two basic blocks and the values passed via block arguments to the
   /// destination block, compute the list of MemRefs that have to be retained in
@@ -233,16 +220,6 @@ FailureOr<Operation *>
 insertDeallocOpForReturnLike(DeallocationState &state, Operation *op,
                              ValueRange operands,
                              SmallVectorImpl<Value> &updatedOperandOwnerships);
-
-/// Materializes IR that extracts the allocated pointers of the MemRef operands
-/// of the defining operation of `memref` as indices and compares them. The
-/// ownership of the first one that matches is returned and intended to be
-/// assigned to `memref`.
-Value defaultComputeMemRefOwnership(const DeallocationOptions &options,
-                                    DeallocationState &state,
-                                    OpBuilder &builder, Value memref,
-                                    Block *block);
-
 } // namespace deallocation_impl
 
 } // namespace bufferization

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

@@ -55,47 +55,17 @@ def BufferDeallocationOpInterface :
           ownership indicator when needed, it should be implemented using this
           method (which is especially important if operations are created that
           cannot be easily canonicalized away anymore).
-          Ownership indicators have to be materialized when
-            * needed for the condition operands of a `bufferization.dealloc` op
-            * passing them as additional operands to nested regions (e.g.,
-              init_args of `scf.for`)
-            * passing them as additional operands to a call operation when
-              `private-function-dynamic-ownership` is enabled
-
-          In the following example, the deallocation pass would add an
-          additional block argument to `^bb1` for passing the ownership of `%0`
-          along and thus the ownership indicator has to be materialized before
-          the `cf.br` operation and added as a forwarded operand.
-          ```mlir
-            %0 = arith.select %cond, %m1, %m2 : memref<f32>
-            cf.br ^bb1(%0 : memref<f32>)
-          ^bb1(%arg0: memref<f32>)
-            ...
-          ```
-          The `arith.select` operation could implement this interface method to
-          materialize another `arith.select` operation to select the
-          corresponding ownership indicator.
-          ```mlir
-            %0 = arith.select %cond, %m1, %m2 : memref<f32>
-            %0_ownership = arith.select %cond, %m1_ownership, %m2_ownership : i1
-            cf.br ^bb1(%0, %0_ownership : memref<f32>, i1)
-          ^bb1(%arg0: memref<f32>, %arg1: i1)
-            ...
-          ```
-
-          The default implementation assumes that all MemRef operands already
-          have 'Unique' ownership.
         }],
-        /*retType=*/"Value",
+        /*retType=*/"std::pair<Value, Value>",
         /*methodName=*/"materializeUniqueOwnershipForMemref",
         /*args=*/(ins "DeallocationState &":$state,
                       "const DeallocationOptions &":$options,
                       "OpBuilder &":$builder,
                       "Value":$memref),
         /*methodBody=*/[{}],
         /*defaultImplementation=*/[{
-          return deallocation_impl::defaultComputeMemRefOwnership(
-            options, state, builder, memref, memref.getParentBlock());
+          return state.getMemrefWithUniqueOwnership(
+            builder, memref, memref.getParentBlock());
         }]>,
   ];
 }

mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h

@@ -17,7 +17,6 @@
 
 namespace mlir {
 namespace bufferization {
-struct DeallocationOptions;
 
 /// Options for the buffer deallocation pipeline.
 struct BufferDeallocationPipelineOptions
@@ -28,23 +27,7 @@ struct BufferDeallocationPipelineOptions
           "Allows to add additional arguments to private functions to "
           "dynamically pass ownership of memrefs to callees. This can enable "
           "earlier deallocations."),
-      llvm::cl::init(true)};
-  PassOptions::Option<bool> verifyFunctionBoundaryABI{
-      *this, "verify-function-boundary-abi",
-      llvm::cl::desc(
-          "Inserts `cf.assert` operations to verify the function boundary ABI "
-          "at runtime. Currently, it is only checked that the ownership of "
-          "returned MemRefs is 'true'. This makes sure that ownership is "
-          "yielded and the returned MemRef does not originate from the same "
-          "allocation as a function argument. If it can be determined "
-          "statically that the ABI is not adhered to, an error will already be "
-          "emitted at compile time. This cannot be changed with this option."),
-      llvm::cl::init(true)};
-
-  /// Convert this BufferDeallocationPipelineOptions struct to a
-  /// DeallocationOptions struct to be passed to the
-  /// OwnershipBasedBufferDeallocationPass.
-  DeallocationOptions asDeallocationOptions() const;
+      llvm::cl::init(false)};
 };
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h

@@ -1,7 +1,6 @@
 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 
-#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
@@ -32,7 +31,7 @@ std::unique_ptr<Pass> createBufferDeallocationPass();
 /// Creates an instance of the OwnershipBasedBufferDeallocation pass to free all
 /// allocated buffers.
 std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass(
-    const DeallocationOptions &options = DeallocationOptions());
+    bool privateFuncDynamicOwnership = false);
 
 /// Creates a pass that optimizes `bufferization.dealloc` operations. For
 /// example, it reduces the number of alias checks needed at runtime using
@@ -135,9 +134,8 @@ func::FuncOp buildDeallocationLibraryFunction(OpBuilder &builder, Location loc,
 LogicalResult deallocateBuffers(Operation *op);
 
 /// Run ownership basedbuffer deallocation.
-LogicalResult deallocateBuffersOwnershipBased(
-    FunctionOpInterface op,
-    const DeallocationOptions &options = DeallocationOptions());
+LogicalResult deallocateBuffersOwnershipBased(FunctionOpInterface op,
+                                              bool privateFuncDynamicOwnership);
 
 /// Creates a pass that moves allocations upwards to reduce the number of
 /// required copies that are inserted during the BufferDeallocation pass.

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

@@ -219,20 +219,10 @@ def OwnershipBasedBufferDeallocation : Pass<
   }];
   let options = [
     Option<"privateFuncDynamicOwnership", "private-function-dynamic-ownership",
-           "bool", /*default=*/"true",
+           "bool", /*default=*/"false",
            "Allows to add additional arguments to private functions to "
            "dynamically pass ownership of memrefs to callees. This can enable "
            "earlier deallocations.">,
-    Option<"verifyFunctionBoundaryABI", "verify-function-boundary-abi",
-           "bool", /*default=*/"true",
-           "Inserts `cf.assert` operations to verify the function boundary ABI "
-           "at runtime. Currently, it is only checked that the ownership of "
-           "returned MemRefs is 'true'. This makes sure that ownership is "
-           "yielded and the returned MemRef does not originate from the same "
-           "allocation as a function argument. "
-           "If it can be determined statically that the ABI is not adhered "
-           "to, an error will already be emitted at compile time. This cannot "
-           "be changed with this option.">,
   ];
   let constructor = "mlir::bufferization::createOwnershipBasedBufferDeallocationPass()";
 

mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp

@@ -53,21 +53,25 @@ struct SelectOpInterface
     return op; // nothing to do
   }
 
-  Value materializeUniqueOwnershipForMemref(Operation *op,
-                                            DeallocationState &state,
-                                            const DeallocationOptions &options,
-                                            OpBuilder &builder,
-                                            Value value) const {
+  std::pair<Value, Value>
+  materializeUniqueOwnershipForMemref(Operation *op, DeallocationState &state,
+                                      const DeallocationOptions &options,
+                                      OpBuilder &builder, Value value) const {
     auto selectOp = cast<arith::SelectOp>(op);
     assert(value == selectOp.getResult() &&
            "Value not defined by this operation");
 
     Block *block = value.getParentBlock();
+    if (!state.getOwnership(selectOp.getTrueValue(), block).isUnique() ||
+        !state.getOwnership(selectOp.getFalseValue(), block).isUnique())
+      return state.getMemrefWithUniqueOwnership(builder, value,
+                                                value.getParentBlock());
+
     Value ownership = builder.create<arith::SelectOp>(
         op->getLoc(), selectOp.getCondition(),
         state.getOwnership(selectOp.getTrueValue(), block).getIndicator(),
         state.getOwnership(selectOp.getFalseValue(), block).getIndicator());
-    return ownership;
+    return {selectOp.getResult(), ownership};
   }
 };
 

mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp

@@ -132,83 +132,30 @@ void DeallocationState::getLiveMemrefsIn(Block *block,
   memrefs.append(liveMemrefs);
 }
 
-Value DeallocationState::materializeMemRefOwnership(
-    const DeallocationOptions &options, OpBuilder &builder, Value memref,
-    Block *block) {
-  // NOTE: Starts at the operation defining `memref` and performs a DFS along
-  // the reverse def/use chain until MemRef values with 'Unique' ownership are
-  // found. For the operation being currently processed:
-  // * if none of the operands have the same allocated pointer (i.e., originate
-  //   from the same allocation), a new memref was allocated and thus the
-  //   operation should have the allocate side-effect defined on that result
-  //   value and thus the correct unique ownership is pre-populated by the
-  //   ownership pass (unless an interface implementation is incorrect). Note
-  //   that this is problematic for operations of unregistered dialects because
-  //   the allocation side-effect cannot be represented in the assembly format.
-  // * if exactly one operand has the same allocated pointer, this returnes the
-  //   ownership of exactly that operand
-  // * if multiple operands match the allocated pointer of the result, the
-  //   ownership indicators of all of them always have to evaluate to the same
-  //   value because no dealloc operations may be present and because of the
-  //   rules they are passed to nested regions and successor blocks.  This could
-  //   be verified at runtime by inserting `cf.assert` operations, but would
-  //   require O(|operands|^2) additional operations to check and is thus not
-  //   implemented yet (would need to insert a library function to avoid
-  //   code-size explosion which would make the deallocation pass a module pass)
-  auto ipSave = builder.saveInsertionPoint();
-  SmallVector<Value> worklist;
-  worklist.push_back(memref);
-
-  while (!worklist.empty()) {
-    Value curr = worklist.back();
-
-    // If the value already has unique ownership, we don't have to process it
-    // anymore.
-    Ownership ownership = getOwnership(curr, block);
-    if (ownership.isUnique()) {
-      worklist.pop_back();
-      continue;
-    }
-
-    // Check if all operands of MemRef type have unique ownership.
-    Operation *defOp = curr.getDefiningOp();
-    assert(defOp &&
-           "the ownership-based deallocation pass should be written in a way "
-           "that pre-populates ownership for block arguments");
-
-    bool allKnown = true;
-    for (Value val : llvm::make_filter_range(defOp->getOperands(), isMemref)) {
-      Ownership ownership = getOwnership(val, block);
-      if (ownership.isUnique())
-        continue;
-
-      worklist.push_back(val);
-      allKnown = false;
-    }
-
-    // If all MemRef operands have unique ownership, we can check if the op
-    // implements the BufferDeallocationOpInterface and call that or, otherwise,
-    // we call the generic implementation manually here.
-    if (allKnown) {
-      builder.setInsertionPointAfter(defOp);
-      if (auto deallocInterface =
-              dyn_cast<BufferDeallocationOpInterface>(defOp);
-          deallocInterface && curr.getParentBlock() == block)
-        ownership = deallocInterface.materializeUniqueOwnershipForMemref(
-            *this, options, builder, curr);
-      else
-        ownership = deallocation_impl::defaultComputeMemRefOwnership(
-            options, *this, builder, curr, block);
-
-      // Ownership is already 'Unknown', so we need to override instead of
-      // joining.
-      resetOwnerships(curr, block);
-      updateOwnership(curr, ownership, block);
-    }
-  }
-
-  builder.restoreInsertionPoint(ipSave);
-  return getOwnership(memref, block).getIndicator();
+std::pair<Value, Value>
+DeallocationState::getMemrefWithUniqueOwnership(OpBuilder &builder,
+                                                Value memref, Block *block) {
+  auto iter = ownershipMap.find({memref, block});
+  assert(iter != ownershipMap.end() &&
+         "Value must already have been registered in the ownership map");
+
+  Ownership ownership = iter->second;
+  if (ownership.isUnique())
+    return {memref, ownership.getIndicator()};
+
+  // Instead of inserting a clone operation we could also insert a dealloc
+  // operation earlier in the block and use the updated ownerships returned by
+  // the op for the retained values. Alternatively, we could insert code to
+  // check aliasing at runtime and use this information to combine two unique
+  // ownerships more intelligently to not end up with an 'Unknown' ownership in
+  // the first place.
+  auto cloneOp =
+      builder.create<bufferization::CloneOp>(memref.getLoc(), memref);
+  Value condition = buildBoolValue(builder, memref.getLoc(), true);
+  Value newMemref = cloneOp.getResult();
+  updateOwnership(newMemref, condition);
+  memrefsToDeallocatePerBlock[newMemref.getParentBlock()].push_back(newMemref);
+  return {newMemref, condition};
 }
 
 void DeallocationState::getMemrefsToRetain(
@@ -366,25 +313,3 @@ FailureOr<Operation *> deallocation_impl::insertDeallocOpForReturnLike(
 
   return op;
 }
-
-Value deallocation_impl::defaultComputeMemRefOwnership(
-    const DeallocationOptions &options, DeallocationState &state,
-    OpBuilder &builder, Value memref, Block *block) {
-  Operation *defOp = memref.getDefiningOp();
-  SmallVector<Value> operands(
-      llvm::make_filter_range(defOp->getOperands(), isMemref));
-  Value resultPtr = builder.create<memref::ExtractAlignedPointerAsIndexOp>(
-      defOp->getLoc(), memref);
-  Value ownership = state.getOwnership(operands.front(), block).getIndicator();
-
-  for (Value val : ArrayRef(operands).drop_front()) {
-    Value operandPtr = builder.create<memref::ExtractAlignedPointerAsIndexOp>(
-        defOp->getLoc(), val);
-    Value isSameBuffer = builder.create<arith::CmpIOp>(
-        defOp->getLoc(), arith::CmpIPredicate::eq, resultPtr, operandPtr);
-    Value newOwnership = state.getOwnership(val, block).getIndicator();
-    ownership = builder.create<arith::SelectOp>(defOp->getLoc(), isSameBuffer,
-                                                newOwnership, ownership);
-  }
-  return ownership;
-}

mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp

@@ -8,34 +8,22 @@
 
 #include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
 
-#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
-using namespace mlir;
-using namespace bufferization;
-
 //===----------------------------------------------------------------------===//
 // Pipeline implementation.
 //===----------------------------------------------------------------------===//
 
-DeallocationOptions
-BufferDeallocationPipelineOptions::asDeallocationOptions() const {
-  DeallocationOptions opts;
-  opts.privateFuncDynamicOwnership = privateFunctionDynamicOwnership.getValue();
-  opts.verifyFunctionBoundaryABI = verifyFunctionBoundaryABI.getValue();
-  return opts;
-}
-
 void mlir::bufferization::buildBufferDeallocationPipeline(
     OpPassManager &pm, const BufferDeallocationPipelineOptions &options) {
   pm.addPass(memref::createExpandReallocPass(/*emitDeallocs=*/false));
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createOwnershipBasedBufferDeallocationPass(
-      options.asDeallocationOptions()));
+      options.privateFunctionDynamicOwnership.getValue()));
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createBufferDeallocationSimplificationPass());
   pm.addPass(createLowerDeallocationsPass());