llvm · maerhart · Sep 27, 2023 · Sep 27, 2023 · matthias-springer · Sep 28, 2023
@@ -10,6 +10,8 @@
 #define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERDEALLOCATIONOPINTERFACE_H_
 
 #include "mlir/Analysis/Liveness.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Support/LLVM.h"
@@ -92,10 +94,65 @@ class Ownership {
 
 /// Options for BufferDeallocationOpInterface-based buffer deallocation.
 struct DeallocationOptions {
+  using DetectionFn = std::function<bool(Operation *)>;
+  using ReplaceDeallocFn = std::function<FailureOr<ValueRange>(Operation *)>;
+
   // 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 makes sure that ownership is yielded and the
+  /// returned MemRef does not originate from the same allocation as a function
+  /// argument. 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;
+
+  /// Given an allocation side-effect on the passed operation, determine whether
+  /// this allocation operation is of relevance (i.e., should assign ownership
+  /// to the allocated value). If it is determined to not be relevant,
+  /// ownership will be set to 'false', i.e., it will be leaked. This is useful
+  /// to support deallocation of multiple different kinds of allocation ops.
+  DetectionFn isRelevantAllocOp = [](Operation *op) {
+    return isa<memref::MemRefDialect, bufferization::BufferizationDialect>(
+        op->getDialect());
+  };
+
+  /// Given a free side-effect on the passed operation, determine whether this
+  /// deallocation operation is of relevance (i.e., should be removed if the
+  /// `removeExistingDeallocations` option is enabled or otherwise an error
+  /// should be emitted because existing deallocation operations are not
+  /// supported without that flag). If it is determined to not be relevant,
+  /// the operation will be ignored. This is useful to support deallocation of
+  /// multiple different kinds of allocation ops where deallocations for some of
+  /// them are already present in the IR.
+  DetectionFn isRelevantDeallocOp = [](Operation *op) {
+    return isa<memref::MemRefDialect, bufferization::BufferizationDialect>(
+        op->getDialect());
+  };
+
+  /// When enabled, remove deallocation operations determined to be relevant
+  /// according to `isRelevantDeallocOp`. If the operation has result values,
+  /// `getDeallocReplacement` will be called to determine the SSA values that
+  /// should be used as replacements.
+  bool removeExistingDeallocations = false;
+
+  /// Provides SSA values for deallocation operations when
+  /// `removeExistingDeallocations` is enabled. May return a failure when the
+  /// given deallocation operation is not supported (e.g., because no
+  /// replacement for a result value can be determined). A failure will directly
+  /// lead to a failure emitted by the deallocation pass.
+  ReplaceDeallocFn getDeallocReplacement =
+      [](Operation *op) -> FailureOr<ValueRange> {
+    if (isa<memref::DeallocOp>(op))
+      return ValueRange{};
+    // ReallocOp has to be expanded before running the dealloc pass.
+    return failure();
+  };
 };
 
 /// This class collects all the state that we need to perform the buffer
@@ -138,12 +195,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 getMemrefWithUniqueOwnership(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

@@ -55,8 +55,38 @@ 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
+            * passed along MemRefs to successor blocks via additional forwarded
+              operands of terminator ops
+            * 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
+            * a copy is made conditionally on the current ownership, etc.
+
+          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)
+            ...
+          ```
         }],
-        /*retType=*/"std::pair<Value, Value>",
+        /*retType=*/"Value",
         /*methodName=*/"materializeUniqueOwnershipForMemref",
         /*args=*/(ins "DeallocationState &":$state,
                       "const DeallocationOptions &":$options,
@@ -65,7 +95,7 @@ def BufferDeallocationOpInterface :
         /*methodBody=*/[{}],
         /*defaultImplementation=*/[{
           return state.getMemrefWithUniqueOwnership(
-            builder, memref, memref.getParentBlock());
+            options, builder, memref, memref.getParentBlock());
         }]>,
   ];
 }

@@ -17,6 +17,7 @@
 
 namespace mlir {
 namespace bufferization {
+struct DeallocationOptions;
 
 /// Options for the buffer deallocation pipeline.
 struct BufferDeallocationPipelineOptions
@@ -27,7 +28,29 @@ 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)};
+  PassOptions::Option<bool> removeExistingDeallocations{
+      *this, "remove-existing-deallocations",
+      llvm::cl::desc("Removes all pre-existing memref.dealloc operations and "
+                     "insert all deallocations according to the buffer "
+                     "deallocation rules."),
       llvm::cl::init(false)};
+
+  /// Convert this BufferDeallocationPipelineOptions struct to a
+  /// DeallocationOptions struct to be passed to the
+  /// OwnershipBasedBufferDeallocationPass.
+  DeallocationOptions asDeallocationOptions() const;
 };
 
 //===----------------------------------------------------------------------===//

@@ -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.

@@ -219,10 +219,25 @@ 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.">,
+    Option<"removeExistingDeallocations", "remove-existing-deallocations",
+           "bool", /*default=*/"false",
+           "Remove already existing MemRef deallocation operations and let the "
+           "deallocation pass insert the deallocation operations according to "
+           "its rules.">,
   ];
   let constructor = "mlir::bufferization::createOwnershipBasedBufferDeallocationPass()";
 

diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -53,25 +53,26 @@ 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,
+      return state.getMemrefWithUniqueOwnership(options, 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;
   }
 };
 

@@ -132,30 +132,79 @@ 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::getMemrefWithUniqueOwnership(
+    const DeallocationOptions &options, OpBuilder &builder, Value memref,
+    Block *block) {
+  // NOTE: * if none of the operands have the same allocated pointer, 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).
+  //       * if exactly one operand has the same allocated pointer, this retunes
+  //       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();
+    Ownership ownership = getOwnership(curr, block);
+    if (ownership.isUnique()) {
+      worklist.pop_back();
+      continue;
+    }
+
+    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 (allKnown) {
+      builder.setInsertionPointAfter(defOp);
+      SmallVector<Value> operands(
+          llvm::make_filter_range(defOp->getOperands(), isMemref));
+      Value resultPtr = builder.create<memref::ExtractAlignedPointerAsIndexOp>(
+          defOp->getLoc(), curr);
+      Value ownership = 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 = getOwnership(val, block).getIndicator();
+        ownership = builder.create<arith::SelectOp>(
+            defOp->getLoc(), isSameBuffer, newOwnership, ownership);
+      }
+      // 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(