[mlir][bufferization] BufferDeallocationOpInterface: support custom ownership update logic

[maerhart]

Summary

Add a method to the BufferDeallocationOpInterface that allows operations to
implement the interface and provide custom logic to compute the ownership
indicators of values it defines. As a demonstrating example, this new method is
implemented by the arith.select operation.

Depends on D158756

[springerm]

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

46–47

the given MemRef typed value

50

and avoid

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

43

materializes

54

I think you still have to call the other handlers in handleAllInterfaces. They are skipped at the moment. A nullptr return value could indicate that the other interface handlers should be skipped. (A non-null op means, they should not be skipped.)

[maerhart]

rebase

[maerhart]

address comments

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

54

I'm not quite sure if I understand. Which handlers should I call here? arith.select does not implement any interfaces that the handlers care about. According to my understanding, the current implementation is correct.

[maerhart]

rebase

[springerm]

Note: This has landed as 942ce3198591811aabac6265d3292e92f89c786a.

mlir/include/mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h

+//===- BufferDeallocationOpInterfaceImpl.h --------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARITH_TRANSFORMS_BUFFERDEALLOCATIONOPINTERFACEIMPL_H
+#define MLIR_DIALECT_ARITH_TRANSFORMS_BUFFERDEALLOCATIONOPINTERFACEIMPL_H
+
+namespace mlir {
+
+class DialectRegistry;
+
+namespace arith {
+void registerBufferDeallocationOpInterfaceExternalModels(
+    DialectRegistry &registry);
+} // namespace arith
+} // namespace mlir
+
+#endif // MLIR_DIALECT_ARITH_TRANSFORMS_BUFFERDEALLOCATIONOPINTERFACEIMPL_H

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

public:
   /// given block.
   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);
+  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
   /// the 'fromBlock' to not run into a use-after-free situation.
   /// This list consists of the MemRefs in the successor operand list of the
   /// terminator and the MemRefs in the 'out' set of the liveness analysis
   /// intersected with the 'in' set of the destination block.
   ///

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

let methods = [
           the newly inserted operations is returned from which on the driver
           should continue the processing. On failure, the deallocation pass
           will terminate. It is recommended to emit a useful error message in
           that case.
         }],
         /*retType=*/"FailureOr<Operation *>",
         /*methodName=*/"process",
         /*args=*/(ins "DeallocationState &":$state,
-                      "const DeallocationOptions &":$options)>
+                      "const DeallocationOptions &":$options)>,
+      InterfaceMethod<
+        /*desc=*/[{
+          This method allows the implementing operation to specify custom logic
+          to materialize an ownership indicator value for the given MemRef typed
+          value it defines (including block arguments of nested regions). Since
+          the operation itself has more information about its semantics the
+          materialized IR can be more efficient compared to the default
+          implementation and avoid cloning MemRefs and/or doing alias checking
+          at runtime.
+          Note that the same logic could also be implemented in the 'process'
+          method above, however, the IR is always materialized then. If
+          it's desirable to only materialize the IR to compute an updated
+          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).
+        }],
+        /*retType=*/"std::pair<Value, Value>",
+        /*methodName=*/"materializeUniqueOwnershipForMemref",
+        /*args=*/(ins "DeallocationState &":$state,
+                      "const DeallocationOptions &":$options,
+                      "OpBuilder &":$builder,
+                      "Value":$memref),
+        /*methodBody=*/[{}],
+        /*defaultImplementation=*/[{
+          return state.getMemrefWithUniqueOwnership(
+            builder, memref, memref.getParentBlock());
+        }]>,
   ];
 }
 
 #endif  // BUFFER_DEALLOCATION_OP_INTERFACE

mlir/include/mlir/InitAllDialects.h

 #define MLIR_INITALLDIALECTS_H_
 
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h"
 #include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSME/IR/ArmSME.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
registry.insert<acc::OpenACCDialect,
                   transform::TransformDialect,
                   ub::UBDialect,
                   vector::VectorDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on
 
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
+  arith::registerBufferDeallocationOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   builtin::registerCastOpInterfaceExternalModels(registry);
   cf::registerBufferizableOpInterfaceExternalModels(registry);
   cf::registerBufferDeallocationOpInterfaceExternalModels(registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);

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

+//===- BufferDeallocationOpInterfaceImpl.cpp ------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+
+static bool isMemref(Value v) { return v.getType().isa<BaseMemRefType>(); }
+
+namespace {
+/// Provides custom logic to materialize ownership indicator values for the
+/// result value of 'arith.select'. Instead of cloning or runtime alias
+/// checking, this implementation inserts another `arith.select` to choose the
+/// ownership indicator of the operand in the same way the original
+/// `arith.select` chooses the MemRef operand. If at least one of the operand's
+/// ownerships is 'Unknown', fall back to the default implementation.
+///
+/// Example:
+/// ```mlir
+/// // let ownership(%m0) := %o0
+/// // let ownership(%m1) := %o1
+/// %res = arith.select %cond, %m0, %m1
+/// ```
+/// The default implementation would insert a clone and replace all uses of the
+/// result of `arith.select` with that clone:
+/// ```mlir
+/// %res = arith.select %cond, %m0, %m1
+/// %clone = bufferization.clone %res
+/// // let ownership(%res) := 'Unknown'
+/// // let ownership(%clone) := %true
+/// // replace all uses of %res with %clone
+/// ```
+/// This implementation, on the other hand, materializes the following:
+/// ```mlir
+/// %res = arith.select %cond, %m0, %m1
+/// %res_ownership = arith.select %cond, %o0, %o1
+/// // let ownership(%res) := %res_ownership
+/// ```
+struct SelectOpInterface
+    : public BufferDeallocationOpInterface::ExternalModel<SelectOpInterface,
+                                                          arith::SelectOp> {
+  FailureOr<Operation *> process(Operation *op, DeallocationState &state,
+                                 const DeallocationOptions &options) const {
+    return op; // nothing to do
+  }
+
+  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 {selectOp.getResult(), ownership};
+  }
+};
+
+} // namespace
+
+void mlir::arith::registerBufferDeallocationOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, ArithDialect *dialect) {
+    SelectOp::attachInterface<SelectOpInterface>(*ctx);
+  });
+}

mlir/lib/Dialect/Arith/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRArithTransforms
+  BufferDeallocationOpInterfaceImpl.cpp
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   EmulateUnsupportedFloats.cpp
   EmulateWideInt.cpp
   EmulateNarrowType.cpp
   ExpandOps.cpp
   IntNarrowing.cpp
   IntRangeOptimizations.cpp

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

void DeallocationState::getLiveMemrefsIn(Block *block,
   SmallVector<Value> liveMemrefs(
       llvm::make_filter_range(liveness.getLiveIn(block), isMemref));
   llvm::sort(liveMemrefs, ValueComparator());
   memrefs.append(liveMemrefs);
 }
 
 std::pair<Value, Value>
 DeallocationState::getMemrefWithUniqueOwnership(OpBuilder &builder,
-                                                Value memref) {
-  auto iter = ownershipMap.find({memref, memref.getParentBlock()});
+                                                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

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

private:
   /// to be updated according to the IR modifications this pass performs (e.g.,
   /// re-building operations to have more result values, inserting clone
   /// operations, etc.).
   void populateRemainingOwnerships(Operation *op);
 
   /// Given an SSA value of MemRef type, returns the same of a new SSA value
   /// which has 'Unique' ownership where the ownership indicator is guaranteed
   /// to be always 'true'.
-  Value getMemrefWithGuaranteedOwnership(OpBuilder &builder, Value memref);
+  Value materializeMemrefWithGuaranteedOwnership(OpBuilder &builder,
+                                                 Value memref, Block *block);
 
   /// Returns whether the given operation implements FunctionOpInterface, has
   /// private visibility, and the private-function-dynamic-ownership pass option
   /// is enabled.
   bool isFunctionWithoutDynamicOwnership(Operation *op);
 
+  /// Given an SSA value of MemRef type, this function queries the
+  /// BufferDeallocationOpInterface of the defining operation of 'memref' for a
+  /// materialized ownership indicator for 'memref'.  If the op does not
+  /// implement the interface or if the block for which the materialized value
+  /// is requested does not match the block in which 'memref' is defined, the
+  /// default implementation in
+  /// `DeallocationState::getMemrefWithUniqueOwnership` is queried instead.
+  std::pair<Value, Value>
+  materializeUniqueOwnership(OpBuilder &builder, Value memref, Block *block);
+
   /// Checks all the preconditions for operations implementing the
   /// FunctionOpInterface that have to hold for the deallocation to be
   /// applicable:
   /// (1) Checks that there are not explicit control flow loops.
   static LogicalResult verifyFunctionPreconditions(FunctionOpInterface op);
 
   /// Checks all the preconditions for operations inside the region of
   /// operations implementing the FunctionOpInterface that have to hold for the

 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocation Implementation
 //===----------------------------------------------------------------------===//
 
+std::pair<Value, Value>
+BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,
+                                               Block *block) {
+  // The interface can only materialize ownership indicators in the same block
+  // as the defining op.
+  if (memref.getParentBlock() != block)
+    return state.getMemrefWithUniqueOwnership(builder, memref, block);
+
+  Operation *owner = memref.getDefiningOp();
+  if (!owner)
+    owner = memref.getParentBlock()->getParentOp();
+
+  // If the op implements the interface, query it for a materialized ownership
+  // value.
+  if (auto deallocOpInterface = dyn_cast<BufferDeallocationOpInterface>(owner))
+    return deallocOpInterface.materializeUniqueOwnershipForMemref(
+        state, options, builder, memref);
+
+  // Otherwise use the default implementation.
+  return state.getMemrefWithUniqueOwnership(builder, memref, block);
+}
+
 static bool regionOperatesOnMemrefValues(Region &region) {
   WalkResult result = region.walk([](Block *block) {
     if (llvm::any_of(block->getArguments(), isMemref))
       return WalkResult::interrupt();
     for (Operation &op : *block) {
       if (llvm::any_of(op.getOperands(), isMemref))
         return WalkResult::interrupt();
       if (llvm::any_of(op.getResults(), isMemref))
BufferDeallocation::handleInterface(RegionBranchOpInterface op) {
   for (auto result : llvm::make_filter_range(newOp->getResults(), isMemref)) {
     state.updateOwnership(result, newOp->getResult(counter++));
     state.addMemrefToDeallocate(result, newOp->getBlock());
   }
 
   return newOp.getOperation();
 }
 
-Value BufferDeallocation::getMemrefWithGuaranteedOwnership(OpBuilder &builder,
-                                                           Value memref) {
+Value BufferDeallocation::materializeMemrefWithGuaranteedOwnership(
+    OpBuilder &builder, Value memref, Block *block) {
   // First, make sure we at least have 'Unique' ownership already.
   std::pair<Value, Value> newMemrefAndOnwership =
-      state.getMemrefWithUniqueOwnership(builder, memref);
+      materializeUniqueOwnership(builder, memref, block);
   Value newMemref = newMemrefAndOnwership.first;
   Value condition = newMemrefAndOnwership.second;
 
   // Avoid inserting additional IR if ownership is already guaranteed. In
   // particular, this is already the case when we had 'Unknown' ownership
   // initially and a clone was inserted to get to 'Unique' ownership.
   if (matchPattern(condition, m_One()))
     return newMemref;
FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {
   if (options.privateFuncDynamicOwnership && isPrivate) {
     SmallVector<Value> newOperands, ownershipIndicatorsToAdd;
     for (Value operand : op.getArgOperands()) {
       if (!isMemref(operand)) {
         newOperands.push_back(operand);
         continue;
       }
       auto [memref, condition] =
-          state.getMemrefWithUniqueOwnership(builder, operand);
+          materializeUniqueOwnership(builder, operand, op->getBlock());
       newOperands.push_back(memref);
       ownershipIndicatorsToAdd.push_back(condition);
     }
     newOperands.append(ownershipIndicatorsToAdd.begin(),
                        ownershipIndicatorsToAdd.end());
     op.getArgOperandsMutable().assign(newOperands);
 
     unsigned numMemrefs = llvm::count_if(op->getResults(), isMemref);
BufferDeallocation::handleInterface(RegionBranchTerminatorOpInterface op) {
   // function boundary ABI.
   bool funcWithoutDynamicOwnership =
       isFunctionWithoutDynamicOwnership(op->getParentOp());
   if (funcWithoutDynamicOwnership) {
     for (OpOperand &val : op->getOpOperands()) {
       if (!isMemref(val.get()))
         continue;
 
-      val.set(getMemrefWithGuaranteedOwnership(builder, val.get()));
+      val.set(materializeMemrefWithGuaranteedOwnership(builder, val.get(),
+                                                       op->getBlock()));
     }
   }
 
   // TODO: getSuccessorRegions is not implemented by all operations we care
   // about, but we would need to check how many successors there are and under
   // which condition they are taken, etc.
 
   MutableOperandRange operands =

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-callop-interface.mlir

 //  CHECK-NEXT:   test.copy
 //  CHECK-NEXT:   [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]
 //  CHECK-NEXT:   bufferization.dealloc ([[ALLOC0]], [[BASE]] :
 //  CHECK-SAME:     if (%true{{[0-9_]*}}, %true{{[0-9_]*}})
 //   CHECK-NOT:     retain
 //  CHECK-NEXT:   return
 
 // CHECK-DYNAMIC-LABEL: func @function_call_requries_merged_ownership_mid_block
+//  CHECK-DYNAMIC-SAME: ([[ARG0:%.+]]: i1)
 //       CHECK-DYNAMIC:   [[ALLOC0:%.+]] = memref.alloc(
 //  CHECK-DYNAMIC-NEXT:   [[ALLOC1:%.+]] = memref.alloca(
-//  CHECK-DYNAMIC-NEXT:   [[SELECT:%.+]] = arith.select{{.*}}[[ALLOC0]], [[ALLOC1]]
-//  CHECK-DYNAMIC-NEXT:   [[CLONE:%.+]] = bufferization.clone [[SELECT]]
-//  CHECK-DYNAMIC-NEXT:   [[RET:%.+]]:2 = call @f([[CLONE]], %true{{[0-9_]*}})
+//  CHECK-DYNAMIC-NEXT:   [[SELECT:%.+]] = arith.select [[ARG0]], [[ALLOC0]], [[ALLOC1]]
+//  CHECK-DYNAMIC-NEXT:   [[RET:%.+]]:2 = call @f([[SELECT]], [[ARG0]])
 //  CHECK-DYNAMIC-NEXT:   test.copy
 //  CHECK-DYNAMIC-NEXT:   [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
-//  CHECK-DYNAMIC-NEXT:   bufferization.dealloc ([[ALLOC0]], [[CLONE]], [[BASE]] :
-//  CHECK-DYNAMIC-SAME:     if (%true{{[0-9_]*}}, %true{{[0-9_]*}}, [[RET]]#1)
+//  CHECK-DYNAMIC-NEXT:   bufferization.dealloc ([[ALLOC0]], [[BASE]] :
+//  CHECK-DYNAMIC-SAME:     if (%true{{[0-9_]*}}, [[RET]]#1)
 //   CHECK-DYNAMIC-NOT:     retain
 //  CHECK-DYNAMIC-NEXT:   return
 
 // TODO: the inserted clone is not necessary, we just have to know which of the
 // two allocations was selected, either by checking aliasing of the result at
 // runtime or by extracting the select condition using an OpInterface or by
 // hardcoding the select op