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

Inserting clones requires a lot of assumptions to hold on the input IR, e.g.,
all writes to a buffer need to dominate all reads. This is not guaranteed by
one-shot bufferization and isn't easy to verify, thus it could quickly lead to
incorrect results that are hard to debug. This commit changes the mechanism of
how an ownership indicator is materialized when there is not already a unique
ownership present. Additionally, we don't create copies of returned memrefs
anymore when we don't have ownership. Instead, we insert assert operations to
make sure we have ownership at runtime, or otherwise report to the user that
correctness could not be guaranteed.

Author: maerhart

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());