[mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI

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

@@ -95,7 +95,20 @@ 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 = false;
+  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;
 };
 
 /// This class collects all the state that we need to perform the buffer
@@ -138,12 +151,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 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);
+  /// 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);
 
   /// 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
@@ -220,6 +233,16 @@ 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,17 +55,47 @@ 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=*/"std::pair<Value, Value>",
+        /*retType=*/"Value",
         /*methodName=*/"materializeUniqueOwnershipForMemref",
         /*args=*/(ins "DeallocationState &":$state,
                       "const DeallocationOptions &":$options,
                       "OpBuilder &":$builder,
                       "Value":$memref),
         /*methodBody=*/[{}],
         /*defaultImplementation=*/[{
-          return state.getMemrefWithUniqueOwnership(
-            builder, memref, memref.getParentBlock());
+          return deallocation_impl::defaultComputeMemRefOwnership(
+            options, state, builder, memref, memref.getParentBlock());
         }]>,
   ];
 }

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

@@ -17,6 +17,7 @@
 
 namespace mlir {
 namespace bufferization {
+struct DeallocationOptions;
 
 /// Options for the buffer deallocation pipeline.
 struct BufferDeallocationPipelineOptions
@@ -27,7 +28,23 @@ 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(false)};
+      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;
 };
 
 //===----------------------------------------------------------------------===//

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

@@ -1,6 +1,7 @@
 #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 {
@@ -31,7 +32,7 @@ std::unique_ptr<Pass> createBufferDeallocationPass();
 /// Creates an instance of the OwnershipBasedBufferDeallocation pass to free all
 /// allocated buffers.
 std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass(
-    bool privateFuncDynamicOwnership = false);
+    const DeallocationOptions &options = DeallocationOptions());
 
 /// Creates a pass that optimizes `bufferization.dealloc` operations. For
 /// example, it reduces the number of alias checks needed at runtime using
@@ -134,8 +135,9 @@ func::FuncOp buildDeallocationLibraryFunction(OpBuilder &builder, Location loc,
 LogicalResult deallocateBuffers(Operation *op);
 
 /// Run ownership basedbuffer deallocation.
-LogicalResult deallocateBuffersOwnershipBased(FunctionOpInterface op,
-                                              bool privateFuncDynamicOwnership);
+LogicalResult deallocateBuffersOwnershipBased(
+    FunctionOpInterface op,
+    const DeallocationOptions &options = DeallocationOptions());
 
 /// 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,10 +219,20 @@ def OwnershipBasedBufferDeallocation : Pass<
   }];
   let options = [
     Option<"privateFuncDynamicOwnership", "private-function-dynamic-ownership",
-           "bool", /*default=*/"false",
+           "bool", /*default=*/"true",
            "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,25 +53,21 @@ struct SelectOpInterface
     return op; // nothing to do
   }
 
-  std::pair<Value, Value>
-  materializeUniqueOwnershipForMemref(Operation *op, DeallocationState &state,
-                                      const DeallocationOptions &options,
-                                      OpBuilder &builder, Value value) const {
+  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 {selectOp.getResult(), ownership};
+    return ownership;
   }
 };
 

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

@@ -132,30 +132,83 @@ void DeallocationState::getLiveMemrefsIn(Block *block,
   memrefs.append(liveMemrefs);
 }
 
-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};
+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();
 }
 
 void DeallocationState::getMemrefsToRetain(
@@ -313,3 +366,25 @@ 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,22 +8,34 @@
 
 #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.privateFunctionDynamicOwnership.getValue()));
+      options.asDeallocationOptions()));
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createBufferDeallocationSimplificationPass());
   pm.addPass(createLowerDeallocationsPass());