[Mlir-commits] [mlir] [mlir][bufferization] Make alloc_tensor's bufferization customizable (PR #215590)
Victor Perez
llvmlistbot at llvm.org
Wed Aug 12 08:29:56 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/215590
>From 7a672a55d6017e9e5afd1e8a859b326cf3f5237b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
<victor.pc.upm at gmail.com>
Date: Tue, 11 Aug 2026 08:02:39 -0700
Subject: [PATCH] [mlir][bufferization] Make the dialect's bufferization
customizable
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The bufferization dialect's ops declared `BufferizableOpInterface` in their
ODS trait lists, so the models were part of the op definitions. The MLIR
interface map keeps the first model attached to an op, and an op-defined
model is always first, so a downstream project could not replace them. For
`bufferization.alloc_tensor` the only remaining hook was
`BufferizationOptions::allocationFn`. That hook cannot change the buffer
type without a contract break: `getBufferType` keeps reporting a static
identity layout, the allocation function returns a different layout, and
every consumer that predicts the type before the buffer exists then
disagrees with the buffer that appears.
Move the models out of the op definitions into external models, the same
way `arith`, `tensor`, `scf`, and the other dialects do it.
`bufferization::registerBufferizableOpInterfaceExternalModels` attaches the
models for `alloc_tensor`, `dealloc_tensor`,
`materialize_in_destination`, `to_buffer` and `to_tensor`, and
`registerAllDialects` calls it. A project that must control how these ops
bufferize attaches its own model first and keeps `getBufferType` and
`bufferize` in agreement.
The model bodies are unchanged, so the default bufferization result is
identical. Note that a tool which builds its own registry and does not call
`registerAllDialects` must now call the new function to bufferize the ops
of the bufferization dialect.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
.../Bufferization/IR/BufferizationOps.td | 126 +------
.../Transforms/BufferizableOpInterfaceImpl.h | 21 ++
.../Bufferization/IR/BufferizationOps.cpp | 202 -----------
.../BufferizableOpInterfaceImpl.cpp | 324 ++++++++++++++++++
.../Bufferization/Transforms/CMakeLists.txt | 1 +
mlir/lib/RegisterAllDialects.cpp | 2 +
.../test/Dialect/Bufferization/bufferize.mlir | 162 +++++++++
7 files changed, 513 insertions(+), 325 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.h
create mode 100644 mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
create mode 100644 mlir/test/Dialect/Bufferization/bufferize.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index a9b2b9f39519d..30317c2d7fd4d 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -11,7 +11,6 @@
include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.td"
include "mlir/Dialect/Bufferization/IR/BufferViewFlowOpInterface.td"
-include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td"
include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td"
include "mlir/Dialect/Bufferization/IR/BufferizationBase.td"
include "mlir/Interfaces/DestinationStyleOpInterface.td"
@@ -27,7 +26,7 @@ class Bufferization_Op<string mnemonic, list<Trait> traits = []>
//===----------------------------------------------------------------------===//
def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
- [AttrSizedOperandSegments, BufferizableOpInterface,
+ [AttrSizedOperandSegments,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface, [
"reifyResultShapes"]>]> {
let summary = "allocate buffer for a tensor";
@@ -93,29 +92,6 @@ def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
let results = (outs AnyTensor:$result);
let extraClassDeclaration = [{
- LogicalResult bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state);
-
- bool resultBufferizesToMemoryWrite(OpResult opResult,
- const AnalysisState &state);
-
- bool bufferizesToAllocation(Value value) { return true; }
-
- bool bufferizesToMemoryRead(OpOperand &opOperand,
- const AnalysisState &state);
-
- bool bufferizesToMemoryWrite(OpOperand &opOperand,
- const AnalysisState &state);
-
- AliasingValueList getAliasingValues(
- OpOperand &opOperand, const AnalysisState &state);
-
- FailureOr<BufferLikeType> getBufferType(
- Value value, const BufferizationOptions &options,
- const BufferizationState &state,
- SmallVector<Value> &invocationStack);
-
RankedTensorType getType() {
return ::llvm::cast<RankedTensorType>(getResult().getType());
}
@@ -219,7 +195,7 @@ def Bufferization_CloneOp : Bufferization_Op<"clone", [
def Bufferization_MaterializeInDestinationOp
: Bufferization_Op<"materialize_in_destination",
[AllElementTypesMatch<["source", "dest"]>,
- BufferizableOpInterface, DestinationStyleOpInterface,
+ DestinationStyleOpInterface,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface, [
"reifyResultShapes"]>,
DeclareOpInterfaceMethods<SubsetOpInterface,
@@ -283,32 +259,11 @@ def Bufferization_MaterializeInDestinationOp
let results = (outs Optional<AnyTensor>:$result);
let extraClassDeclaration = [{
- LogicalResult bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state);
-
- bool bufferizesToMemoryRead(OpOperand &opOperand,
- const AnalysisState &state);
-
- bool bufferizesToMemoryWrite(OpOperand &opOperand,
- const AnalysisState &state);
-
- bool bufferizesToElementwiseAccess(const AnalysisState &state,
- ArrayRef<OpOperand *> opOperands);
-
- bool mustBufferizeInPlace(OpOperand &opOperand,
- const AnalysisState &state);
-
- AliasingValueList getAliasingValues(
- OpOperand &opOperand, const AnalysisState &state);
-
RankedTensorType getType() {
return ::llvm::cast<RankedTensorType>(getResult().getType());
}
MutableOperandRange getDpsInitsMutable();
-
- bool isWritable(Value value, const AnalysisState &state);
}];
let builders = [
@@ -329,8 +284,7 @@ def Bufferization_MaterializeInDestinationOp
// DeallocTensorOp
//===----------------------------------------------------------------------===//
-def Bufferization_DeallocTensorOp : Bufferization_Op<"dealloc_tensor",
- [BufferizableOpInterface]> {
+def Bufferization_DeallocTensorOp : Bufferization_Op<"dealloc_tensor"> {
string summary = "release underlying storage format of given tensor";
string description = [{
`bufferization.dealloc_tensor` is a buffer deallocation in tensor land. This
@@ -360,27 +314,6 @@ def Bufferization_DeallocTensorOp : Bufferization_Op<"dealloc_tensor",
let arguments = (ins AnyTensor:$tensor);
let results = (outs);
let assemblyFormat = "$tensor attr-dict `:` type($tensor)";
-
- let extraClassDeclaration = [{
- bool bufferizesToMemoryRead(OpOperand &opOperand,
- const AnalysisState &state) const {
- return false;
- }
-
- bool bufferizesToMemoryWrite(OpOperand &opOperand,
- const AnalysisState &state) const {
- return false;
- }
-
- AliasingValueList getAliasingValues(
- OpOperand &opOperand, const AnalysisState &state) const {
- return {};
- }
-
- LogicalResult bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state);
- }];
}
//===----------------------------------------------------------------------===//
@@ -396,7 +329,6 @@ class Bufferization_TensorAndBufferMatch<string tensor, string buffer> : PredOpT
>;
def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
- BufferizableOpInterface,
SameOperandsAndResultShape,
SameOperandsAndResultElementType,
Bufferization_TensorAndBufferMatch<"result", "buffer">
@@ -464,25 +396,6 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
::mlir::bufferization::TensorLikeType getType() {
return getResult().getType();
}
-
- //===------------------------------------------------------------------===//
- // BufferizableOpInterface implementation
- //===------------------------------------------------------------------===//
-
- LogicalResult bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state) const {
- // to_tensor/to_buffer pairs fold away after bufferization.
- return success();
- }
-
- bool isWritable(Value value, const AnalysisState &state);
-
- FailureOr<BufferLikeType> getBufferType(
- Value value, const BufferizationOptions &options,
- const BufferizationState &state, SmallVector<Value> &invocationStack) {
- return getBuffer().getType();
- }
}];
let assemblyFormat = [{
@@ -500,7 +413,6 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
//===----------------------------------------------------------------------===//
def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
- BufferizableOpInterface,
SameOperandsAndResultShape,
SameOperandsAndResultElementType,
Pure,
@@ -527,38 +439,6 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
let arguments = (ins Bufferization_TensorLikeTypeInterface:$tensor, UnitAttr:$read_only);
let results = (outs Bufferization_BufferLikeTypeInterface:$buffer);
- let extraClassDeclaration = [{
- //===------------------------------------------------------------------===//
- // BufferizableOpInterface implementation
- //===------------------------------------------------------------------===//
-
- // Note: ToBufferOp / ToTensorOp are temporary ops that are inserted at the
- // bufferization boundary. When One-Shot bufferization is complete, there
- // should be no such ops left over. If `allowUnknownOps` (or after running a
- // partial bufferization pass), such ops may be part of the resulting IR,
- // but such IR may no longer be analyzable by One-Shot analysis.
-
- bool bufferizesToMemoryRead(OpOperand &opOperand,
- const AnalysisState &state) const {
- // It is unknown whether the resulting memref will be read or not.
- return true;
- }
-
- bool bufferizesToMemoryWrite(OpOperand &opOperand,
- const AnalysisState &state) {
- return !getReadOnly();
- }
-
- AliasingValueList getAliasingValues(
- OpOperand &opOperand, const AnalysisState &state) const {
- return {};
- }
-
- LogicalResult bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state);
- }];
-
let assemblyFormat = [{
$tensor (`read_only` $read_only^)? attr-dict `:` type($tensor) `to` type($buffer)
}];
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.h
new file mode 100644
index 0000000000000..5ef272384c082
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.h
@@ -0,0 +1,21 @@
+//===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
+//
+// 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_BUFFERIZATION_TRANSFORMS_BUFFERIZABLEOPINTERFACEIMPL_H
+#define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZABLEOPINTERFACEIMPL_H
+
+namespace mlir {
+
+class DialectRegistry;
+
+namespace bufferization {
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry);
+} // namespace bufferization
+} // namespace mlir
+
+#endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZABLEOPINTERFACEIMPL_H
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 80db89aa1b4bc..9b9e631e7cebd 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -162,109 +162,6 @@ void mlir::bufferization::populateDynamicDimSizes(
// AllocTensorOp
//===----------------------------------------------------------------------===//
-LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state) {
- OpBuilder::InsertionGuard g(rewriter);
- Location loc = getLoc();
-
- // Nothing to do for dead AllocTensorOps.
- if (getOperation()->getUses().empty()) {
- rewriter.eraseOp(getOperation());
- return success();
- }
-
- // Get "copy" buffer.
- Value copyBuffer;
- if (getCopy()) {
- FailureOr<Value> maybeCopyBuffer =
- getBuffer(rewriter, getCopy(), options, state);
- if (failed(maybeCopyBuffer))
- return failure();
- copyBuffer = *maybeCopyBuffer;
- }
-
- // Create memory allocation.
- auto allocType = bufferization::getBufferType(getResult(), options, state);
- if (failed(allocType))
- return failure();
- SmallVector<Value> dynamicDims = getDynamicSizes();
- if (getCopy()) {
- assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
- populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
- }
- FailureOr<Value> alloc =
- options.allocationFn(rewriter, loc, llvm::cast<MemRefType>(*allocType),
- dynamicDims, options.bufferAlignment);
- if (failed(alloc))
- return failure();
-
- // Create memory copy (if any).
- if (getCopy()) {
- if (failed(options.memCpyFn(rewriter, loc, copyBuffer, *alloc)))
- return failure();
- }
-
- // Replace op.
- replaceOpWithBufferizedValues(rewriter, getOperation(), *alloc);
-
- return success();
-}
-
-bool AllocTensorOp::resultBufferizesToMemoryWrite(OpResult opResult,
- const AnalysisState &state) {
- // AllocTensorOps do not write unless they have a `copy` value.
- return static_cast<bool>(getCopy());
-}
-
-bool AllocTensorOp::bufferizesToMemoryRead(OpOperand &opOperand,
- const AnalysisState &state) {
- assert(opOperand.getOperandNumber() == getNumOperands() - 1 &&
- "expected copy operand");
- return true;
-}
-
-bool AllocTensorOp::bufferizesToMemoryWrite(OpOperand &opOperand,
- const AnalysisState &state) {
- assert(opOperand.getOperandNumber() == getNumOperands() - 1 &&
- "expected copy operand");
- return false;
-}
-
-AliasingValueList AllocTensorOp::getAliasingValues(OpOperand &opOperand,
- const AnalysisState &state) {
- // This is a new allocation. It does not alias with any other buffer.
- return {};
-}
-
-FailureOr<BufferLikeType>
-AllocTensorOp::getBufferType(Value value, const BufferizationOptions &options,
- const BufferizationState &state,
- SmallVector<Value> &invocationStack) {
- assert(value == getResult() && "invalid value");
-
- // Compute memory space of this allocation.
- Attribute memorySpace;
- if (getMemorySpace().has_value()) {
- memorySpace = *getMemorySpace();
- } else if (getCopy()) {
- auto copyBufferType =
- bufferization::detail::asMemRefType(bufferization::getBufferType(
- getCopy(), options, state, invocationStack));
- if (failed(copyBufferType))
- return failure();
- memorySpace = copyBufferType->getMemorySpace();
- } else if (auto ms = options.defaultMemorySpaceFn(
- cast<TensorLikeType>(getType()))) {
- memorySpace = *ms;
- } else {
- return getOperation()->emitError("could not infer memory space");
- }
-
- return cast<BufferLikeType>(
- getMemRefTypeWithStaticIdentityLayout(getType(), memorySpace));
-}
-
LogicalResult AllocTensorOp::verify() {
if (getCopy() && !getDynamicSizes().empty())
return emitError("dynamic sizes not needed when copying a tensor");
@@ -549,90 +446,10 @@ void CloneOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<SimplifyClones>(context);
}
-//===----------------------------------------------------------------------===//
-// DeallocTensorOp
-//===----------------------------------------------------------------------===//
-
-LogicalResult DeallocTensorOp::bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state) {
- FailureOr<Value> buffer = getBuffer(rewriter, getTensor(), options, state);
- if (failed(buffer))
- return failure();
- memref::DeallocOp::create(rewriter, getLoc(), *buffer);
- rewriter.eraseOp(getOperation());
- return success();
-}
-
//===----------------------------------------------------------------------===//
// MaterializeInDestinationOp
//===----------------------------------------------------------------------===//
-bool MaterializeInDestinationOp::bufferizesToMemoryRead(
- OpOperand &opOperand, const AnalysisState &state) {
- return opOperand == getSourceMutable();
-}
-
-bool MaterializeInDestinationOp::bufferizesToMemoryWrite(
- OpOperand &opOperand, const AnalysisState &state) {
- if (opOperand == getDestMutable()) {
- assert(isa<TensorType>(getDest().getType()) && "expected tensor type");
- return true;
- }
- return false;
-}
-
-bool MaterializeInDestinationOp::mustBufferizeInPlace(
- OpOperand &opOperand, const AnalysisState &state) {
- // The source is only read and not written, so it always bufferizes in-place
- // by default. The destination is written and is forced to bufferize in-place
- // (if it is a tensor).
- return true;
-}
-
-AliasingValueList
-MaterializeInDestinationOp::getAliasingValues(OpOperand &opOperand,
- const AnalysisState &state) {
- if (opOperand == getDestMutable()) {
- assert(isa<TensorType>(getDest().getType()) && "expected tensor type");
- return {{getOperation()->getResult(0), BufferRelation::Equivalent}};
- }
- return {};
-}
-
-LogicalResult
-MaterializeInDestinationOp::bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state) {
- bool tensorDest = isa<TensorType>(getDest().getType());
- Value buffer;
- if (tensorDest) {
- FailureOr<Value> maybeBuffer =
- getBuffer(rewriter, getDest(), options, state);
- if (failed(maybeBuffer))
- return failure();
- buffer = *maybeBuffer;
- } else {
- assert(isa<BaseMemRefType>(getDest().getType()) && "expected memref type");
- buffer = getDest();
- }
- auto srcBuffer = getBuffer(rewriter, getSource(), options, state);
- if (failed(srcBuffer))
- return failure();
- if (failed(options.memCpyFn(rewriter, getLoc(), *srcBuffer, buffer)))
- return failure();
- replaceOpWithBufferizedValues(rewriter, getOperation(),
- tensorDest ? ValueRange(buffer) : ValueRange());
- return success();
-}
-
-bool MaterializeInDestinationOp::bufferizesToElementwiseAccess(
- const AnalysisState &state, ArrayRef<OpOperand *> opOperands) {
- // As elements are copied from the "source" buffer to the "dest" buffer,
- // already copied elements are not read a second time.
- return true;
-}
-
LogicalResult MaterializeInDestinationOp::reifyResultShapes(
OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
if (getOperation()->getNumResults() == 1) {
@@ -743,11 +560,6 @@ void MaterializeInDestinationOp::build(OpBuilder &builder,
source, dest);
}
-bool MaterializeInDestinationOp::isWritable(Value value,
- const AnalysisState &state) {
- return isa<TensorType>(getDest().getType()) ? true : getWritable();
-}
-
MutableOperandRange MaterializeInDestinationOp::getDpsInitsMutable() {
return getDestMutable();
}
@@ -764,10 +576,6 @@ void MaterializeInDestinationOp::getEffects(
// ToTensorOp
//===----------------------------------------------------------------------===//
-bool ToTensorOp::isWritable(Value value, const AnalysisState &state) {
- return getWritable();
-}
-
OpFoldResult ToTensorOp::fold(FoldAdaptor) {
if (auto toBuffer = getBuffer().getDefiningOp<ToBufferOp>())
// Approximate alias analysis by conservatively folding only when no there
@@ -897,16 +705,6 @@ void ToBufferOp::getCanonicalizationPatterns(RewritePatternSet &results,
ToBufferToTensorFolding>(context);
}
-LogicalResult ToBufferOp::bufferize(RewriterBase &rewriter,
- const BufferizationOptions &options,
- BufferizationState &state) {
- // Fold to_buffer(to_tensor(x)) to x. Insert a cast if necessary.
- (void)foldToBufferToTensorPair(rewriter, *this, options);
- // Note: The return value of `bufferize` indicates whether there was an error
- // or not. (And not whether the pattern matched or not.)
- return success();
-}
-
std::optional<Operation *> CloneOp::buildDealloc(OpBuilder &builder,
Value alloc) {
return memref::DeallocOp::create(builder, alloc.getLoc(), alloc)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
new file mode 100644
index 0000000000000..27e7ddb5d040e
--- /dev/null
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -0,0 +1,324 @@
+//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
+//
+// 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/Bufferization/Transforms/BufferizableOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.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;
+
+namespace mlir {
+namespace bufferization {
+namespace {
+
+struct AllocTensorOpInterface
+ : public BufferizableOpInterface::ExternalModel<AllocTensorOpInterface,
+ AllocTensorOp> {
+ bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
+
+ bool resultBufferizesToMemoryWrite(Operation *op, OpResult opResult,
+ const AnalysisState &state) const {
+ // AllocTensorOps do not write unless they have a `copy` value.
+ return static_cast<bool>(cast<AllocTensorOp>(op).getCopy());
+ }
+
+ bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ assert(opOperand.getOperandNumber() == op->getNumOperands() - 1 &&
+ "expected copy operand");
+ return true;
+ }
+
+ bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ assert(opOperand.getOperandNumber() == op->getNumOperands() - 1 &&
+ "expected copy operand");
+ return false;
+ }
+
+ AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ // This is a new allocation. It does not alias with any other buffer.
+ return {};
+ }
+
+ FailureOr<BufferLikeType>
+ getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+ const BufferizationState &state,
+ SmallVector<Value> &invocationStack) const {
+ auto allocTensorOp = cast<AllocTensorOp>(op);
+ assert(value == allocTensorOp.getResult() && "invalid value");
+
+ // Compute memory space of this allocation.
+ Attribute memorySpace;
+ if (allocTensorOp.getMemorySpace().has_value()) {
+ memorySpace = *allocTensorOp.getMemorySpace();
+ } else if (allocTensorOp.getCopy()) {
+ auto copyBufferType =
+ bufferization::detail::asMemRefType(bufferization::getBufferType(
+ allocTensorOp.getCopy(), options, state, invocationStack));
+ if (failed(copyBufferType))
+ return failure();
+ memorySpace = copyBufferType->getMemorySpace();
+ } else if (auto ms = options.defaultMemorySpaceFn(
+ cast<TensorLikeType>(allocTensorOp.getType()))) {
+ memorySpace = *ms;
+ } else {
+ return op->emitError("could not infer memory space");
+ }
+
+ return cast<BufferLikeType>(getMemRefTypeWithStaticIdentityLayout(
+ allocTensorOp.getType(), memorySpace));
+ }
+
+ LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+ const BufferizationOptions &options,
+ BufferizationState &state) const {
+ auto allocTensorOp = cast<AllocTensorOp>(op);
+ OpBuilder::InsertionGuard g(rewriter);
+ Location loc = allocTensorOp.getLoc();
+
+ // Nothing to do for dead AllocTensorOps.
+ if (op->getUses().empty()) {
+ rewriter.eraseOp(op);
+ return success();
+ }
+
+ // Get "copy" buffer.
+ Value copyBuffer;
+ if (allocTensorOp.getCopy()) {
+ FailureOr<Value> maybeCopyBuffer = bufferization::getBuffer(
+ rewriter, allocTensorOp.getCopy(), options, state);
+ if (failed(maybeCopyBuffer))
+ return failure();
+ copyBuffer = *maybeCopyBuffer;
+ }
+
+ // Create memory allocation.
+ auto allocType =
+ bufferization::getBufferType(allocTensorOp.getResult(), options, state);
+ if (failed(allocType))
+ return failure();
+ SmallVector<Value> dynamicDims = allocTensorOp.getDynamicSizes();
+ if (allocTensorOp.getCopy()) {
+ assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
+ populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
+ }
+ FailureOr<Value> alloc =
+ options.allocationFn(rewriter, loc, llvm::cast<MemRefType>(*allocType),
+ dynamicDims, options.bufferAlignment);
+ if (failed(alloc))
+ return failure();
+
+ // Create memory copy (if any).
+ if (allocTensorOp.getCopy()) {
+ if (failed(options.memCpyFn(rewriter, loc, copyBuffer, *alloc)))
+ return failure();
+ }
+
+ // Replace op.
+ replaceOpWithBufferizedValues(rewriter, op, *alloc);
+
+ return success();
+ }
+};
+
+struct DeallocTensorOpInterface
+ : public BufferizableOpInterface::ExternalModel<DeallocTensorOpInterface,
+ DeallocTensorOp> {
+ bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ return false;
+ }
+
+ bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ return false;
+ }
+
+ AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ return {};
+ }
+
+ LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+ const BufferizationOptions &options,
+ BufferizationState &state) const {
+ auto deallocTensorOp = cast<DeallocTensorOp>(op);
+ FailureOr<Value> buffer = bufferization::getBuffer(
+ rewriter, deallocTensorOp.getTensor(), options, state);
+ if (failed(buffer))
+ return failure();
+ memref::DeallocOp::create(rewriter, deallocTensorOp.getLoc(), *buffer);
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
+struct MaterializeInDestinationOpInterface
+ : public BufferizableOpInterface::ExternalModel<
+ MaterializeInDestinationOpInterface, MaterializeInDestinationOp> {
+ bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ return opOperand == cast<MaterializeInDestinationOp>(op).getSourceMutable();
+ }
+
+ bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ auto materializeOp = cast<MaterializeInDestinationOp>(op);
+ if (opOperand == materializeOp.getDestMutable()) {
+ assert(isa<TensorType>(materializeOp.getDest().getType()) &&
+ "expected tensor type");
+ return true;
+ }
+ return false;
+ }
+
+ bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ // The source is only read and not written, so it always bufferizes in-place
+ // by default. The destination is written and is forced to bufferize
+ // in-place (if it is a tensor).
+ return true;
+ }
+
+ AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ auto materializeOp = cast<MaterializeInDestinationOp>(op);
+ if (opOperand == materializeOp.getDestMutable()) {
+ assert(isa<TensorType>(materializeOp.getDest().getType()) &&
+ "expected tensor type");
+ return {{op->getResult(0), BufferRelation::Equivalent}};
+ }
+ return {};
+ }
+
+ LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+ const BufferizationOptions &options,
+ BufferizationState &state) const {
+ auto materializeOp = cast<MaterializeInDestinationOp>(op);
+ bool tensorDest = isa<TensorType>(materializeOp.getDest().getType());
+ Value buffer;
+ if (tensorDest) {
+ FailureOr<Value> maybeBuffer = bufferization::getBuffer(
+ rewriter, materializeOp.getDest(), options, state);
+ if (failed(maybeBuffer))
+ return failure();
+ buffer = *maybeBuffer;
+ } else {
+ assert(isa<BaseMemRefType>(materializeOp.getDest().getType()) &&
+ "expected memref type");
+ buffer = materializeOp.getDest();
+ }
+ auto srcBuffer = bufferization::getBuffer(
+ rewriter, materializeOp.getSource(), options, state);
+ if (failed(srcBuffer))
+ return failure();
+ if (failed(options.memCpyFn(rewriter, materializeOp.getLoc(), *srcBuffer,
+ buffer)))
+ return failure();
+ replaceOpWithBufferizedValues(
+ rewriter, op, tensorDest ? ValueRange(buffer) : ValueRange());
+ return success();
+ }
+
+ bool bufferizesToElementwiseAccess(Operation *op, const AnalysisState &state,
+ ArrayRef<OpOperand *> opOperands) const {
+ // As elements are copied from the "source" buffer to the "dest" buffer,
+ // already copied elements are not read a second time.
+ return true;
+ }
+
+ bool isWritable(Operation *op, Value value,
+ const AnalysisState &state) const {
+ auto materializeOp = cast<MaterializeInDestinationOp>(op);
+ return isa<TensorType>(materializeOp.getDest().getType())
+ ? true
+ : materializeOp.getWritable();
+ }
+};
+
+// Note: ToBufferOp / ToTensorOp are temporary ops that are inserted at the
+// bufferization boundary. When One-Shot bufferization is complete, there should
+// be no such ops left over. If `allowUnknownOps` (or after running a partial
+// bufferization pass), such ops may be part of the resulting IR, but such IR
+// may no longer be analyzable by One-Shot analysis.
+
+struct ToTensorOpInterface
+ : public BufferizableOpInterface::ExternalModel<ToTensorOpInterface,
+ ToTensorOp> {
+ bool isWritable(Operation *op, Value value,
+ const AnalysisState &state) const {
+ return cast<ToTensorOp>(op).getWritable();
+ }
+
+ FailureOr<BufferLikeType>
+ getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+ const BufferizationState &state,
+ SmallVector<Value> &invocationStack) const {
+ return cast<ToTensorOp>(op).getBuffer().getType();
+ }
+
+ LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+ const BufferizationOptions &options,
+ BufferizationState &state) const {
+ // to_tensor/to_buffer pairs fold away after bufferization.
+ return success();
+ }
+};
+
+struct ToBufferOpInterface
+ : public BufferizableOpInterface::ExternalModel<ToBufferOpInterface,
+ ToBufferOp> {
+ bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ // It is unknown whether the resulting memref will be read or not.
+ return true;
+ }
+
+ bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ return !cast<ToBufferOp>(op).getReadOnly();
+ }
+
+ AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
+ const AnalysisState &state) const {
+ return {};
+ }
+
+ LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+ const BufferizationOptions &options,
+ BufferizationState &state) const {
+ // Fold to_buffer(to_tensor(x)) to x. Insert a cast if necessary.
+ (void)foldToBufferToTensorPair(rewriter, cast<ToBufferOp>(op), options);
+ // Note: The return value of `bufferize` indicates whether there was an
+ // error or not. (And not whether the pattern matched or not.)
+ return success();
+ }
+};
+
+} // namespace
+} // namespace bufferization
+} // namespace mlir
+
+void mlir::bufferization::registerBufferizableOpInterfaceExternalModels(
+ DialectRegistry ®istry) {
+ registry.addExtension(+[](MLIRContext *ctx, BufferizationDialect *dialect) {
+ AllocTensorOp::attachInterface<AllocTensorOpInterface>(*ctx);
+ DeallocTensorOp::attachInterface<DeallocTensorOpInterface>(*ctx);
+ MaterializeInDestinationOp::attachInterface<
+ MaterializeInDestinationOpInterface>(*ctx);
+ ToBufferOp::attachInterface<ToBufferOpInterface>(*ctx);
+ ToTensorOp::attachInterface<ToTensorOpInterface>(*ctx);
+ });
+}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
index 006fcd1ce0ec7..2e6bc4d9503b9 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
@@ -1,4 +1,5 @@
add_mlir_dialect_library(MLIRBufferizationTransforms
+ BufferizableOpInterfaceImpl.cpp
Bufferize.cpp
BufferDeallocationSimplification.cpp
BufferOptimizations.cpp
diff --git a/mlir/lib/RegisterAllDialects.cpp b/mlir/lib/RegisterAllDialects.cpp
index 948abc2ae4f4b..aa502e5ee08a6 100644
--- a/mlir/lib/RegisterAllDialects.cpp
+++ b/mlir/lib/RegisterAllDialects.cpp
@@ -29,6 +29,7 @@
#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
@@ -168,6 +169,7 @@ void mlir::registerAllDialects(DialectRegistry ®istry) {
arith::registerBufferViewFlowOpInterfaceExternalModels(registry);
arith::registerShardingInterfaceExternalModels(registry);
arith::registerValueBoundsOpInterfaceExternalModels(registry);
+ bufferization::registerBufferizableOpInterfaceExternalModels(registry);
bufferization::registerValueBoundsOpInterfaceExternalModels(registry);
bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
registry);
diff --git a/mlir/test/Dialect/Bufferization/bufferize.mlir b/mlir/test/Dialect/Bufferization/bufferize.mlir
new file mode 100644
index 0000000000000..4719689114285
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/bufferize.mlir
@@ -0,0 +1,162 @@
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries" -split-input-file | FileCheck %s
+
+// CHECK-LABEL: @alloc_tensor_static
+func.func @alloc_tensor_static() -> tensor<8x16xf32> {
+ // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 64 : i64} : memref<8x16xf32>
+ // CHECK: return %[[ALLOC]]
+ %0 = bufferization.alloc_tensor() : tensor<8x16xf32>
+ return %0 : tensor<8x16xf32>
+}
+
+// -----
+
+// The dynamic extents become the memref.alloc operands.
+
+// CHECK-LABEL: @alloc_tensor_dynamic
+// CHECK-SAME: %[[D0:.*]]: index
+func.func @alloc_tensor_dynamic(%d0: index) -> tensor<?x16xf32> {
+ // CHECK: %[[ALLOC:.*]] = memref.alloc(%[[D0]]) {alignment = 64 : i64} : memref<?x16xf32>
+ // CHECK: return %[[ALLOC]]
+ %0 = bufferization.alloc_tensor(%d0) : tensor<?x16xf32>
+ return %0 : tensor<?x16xf32>
+}
+
+// -----
+
+// The `memory_space` attribute wins over `defaultMemorySpaceFn`. The result
+// stays inside the function so the function-boundary type conversion does not
+// drive the allocated type.
+
+// CHECK-LABEL: @alloc_tensor_memory_space
+func.func @alloc_tensor_memory_space(%i: index) -> f32 {
+ // CHECK: memref.alloc() {alignment = 64 : i64} : memref<8x16xf32, 1>
+ %0 = bufferization.alloc_tensor() {memory_space = 1 : i64} : tensor<8x16xf32>
+ %1 = tensor.extract %0[%i, %i] : tensor<8x16xf32>
+ return %1 : f32
+}
+
+// -----
+
+// A `copy` operand allocates a fresh buffer and copies into it.
+
+// CHECK-LABEL: @alloc_tensor_copy
+// CHECK-SAME: %[[ARG:.*]]: memref<8x16xf32
+func.func @alloc_tensor_copy(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
+ // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 64 : i64} : memref<8x16xf32>
+ // CHECK: memref.copy %[[ARG]], %[[ALLOC]]
+ // CHECK: return %[[ALLOC]]
+ %0 = bufferization.alloc_tensor() copy(%arg0) : tensor<8x16xf32>
+ return %0 : tensor<8x16xf32>
+}
+
+// -----
+
+// An unused alloc_tensor is erased rather than allocated.
+
+// CHECK-LABEL: @alloc_tensor_dead
+// CHECK-NOT: memref.alloc
+func.func @alloc_tensor_dead() {
+ %0 = bufferization.alloc_tensor() : tensor<8x16xf32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @dealloc_tensor
+func.func @dealloc_tensor() {
+ // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 64 : i64} : memref<8x16xf32>
+ // CHECK: memref.dealloc %[[ALLOC]]
+ %0 = bufferization.alloc_tensor() : tensor<8x16xf32>
+ bufferization.dealloc_tensor %0 : tensor<8x16xf32>
+ return
+}
+
+// -----
+
+// A tensor destination is written in place and returned.
+
+// CHECK-LABEL: @materialize_in_destination_tensor
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9_]*]]: memref<5xf32,
+// CHECK-SAME: %[[DST:[a-zA-Z0-9_]*]]: memref<5xf32,
+func.func @materialize_in_destination_tensor(%src: tensor<5xf32>, %dst: tensor<5xf32>) -> tensor<5xf32> {
+ // CHECK: memref.copy %[[SRC]], %[[DST]]
+ // CHECK: return %[[DST]]
+ %0 = bufferization.materialize_in_destination %src in %dst : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32>
+ return %0 : tensor<5xf32>
+}
+
+// -----
+
+// A memref destination is copied into directly and the op has no result.
+
+// CHECK-LABEL: @materialize_in_destination_memref
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9_]*]]: memref<5xf32,
+// CHECK-SAME: %[[DST:[a-zA-Z0-9_]*]]: memref<5xf32>
+func.func @materialize_in_destination_memref(%src: tensor<5xf32>, %dst: memref<5xf32>) {
+ // CHECK: memref.copy %[[SRC]], %[[DST]]
+ bufferization.materialize_in_destination %src in restrict writable %dst
+ : (tensor<5xf32>, memref<5xf32>) -> ()
+ return
+}
+
+// -----
+
+// Without `writable`, the buffer of a to_tensor must not be written, so the
+// insert bufferizes out of place.
+
+// CHECK-LABEL: @to_tensor_not_writable
+// CHECK-SAME: %[[M:[a-zA-Z0-9_]*]]: memref<5xf32>
+func.func @to_tensor_not_writable(%m: memref<5xf32>, %f: f32, %idx: index) -> tensor<5xf32> {
+ // CHECK: %[[ALLOC:.*]] = memref.alloc()
+ // CHECK: memref.copy %[[M]], %[[ALLOC]]
+ // CHECK: memref.store %{{.*}}, %[[ALLOC]]
+ %t = bufferization.to_tensor %m restrict : memref<5xf32> to tensor<5xf32>
+ %r = tensor.insert %f into %t[%idx] : tensor<5xf32>
+ return %r : tensor<5xf32>
+}
+
+// -----
+
+// With `writable`, the insert bufferizes in place.
+
+// CHECK-LABEL: @to_tensor_writable
+// CHECK-SAME: %[[M:[a-zA-Z0-9_]*]]: memref<5xf32>
+// CHECK-NOT: memref.alloc
+func.func @to_tensor_writable(%m: memref<5xf32>, %f: f32, %idx: index) -> tensor<5xf32> {
+ // CHECK: memref.store %{{.*}}, %[[M]]
+ %t = bufferization.to_tensor %m restrict writable : memref<5xf32> to tensor<5xf32>
+ %r = tensor.insert %f into %t[%idx] : tensor<5xf32>
+ return %r : tensor<5xf32>
+}
+
+// -----
+
+// to_buffer/to_tensor pairs fold away.
+
+// CHECK-LABEL: @to_buffer_of_to_tensor
+// CHECK-SAME: %[[M:[a-zA-Z0-9_]*]]: memref<5xf32>
+// CHECK-NOT: bufferization.to_tensor
+// CHECK-NOT: bufferization.to_buffer
+func.func @to_buffer_of_to_tensor(%m: memref<5xf32>, %f: f32, %idx: index) {
+ // CHECK: memref.store %{{.*}}, %[[M]]
+ %t = bufferization.to_tensor %m restrict writable : memref<5xf32> to tensor<5xf32>
+ %r = bufferization.to_buffer %t : tensor<5xf32> to memref<5xf32>
+ memref.store %f, %r[%idx] : memref<5xf32>
+ return
+}
+
+// -----
+
+// A `read_only` to_buffer does not write, so no copy of the source buffer is
+// needed.
+
+// CHECK-LABEL: @to_buffer_read_only
+// CHECK-SAME: %[[M:[a-zA-Z0-9_]*]]: memref<5xf32>
+// CHECK-NOT: memref.alloc
+func.func @to_buffer_read_only(%m: memref<5xf32>, %idx: index) -> f32 {
+ // CHECK: memref.load %[[M]]
+ %t = bufferization.to_tensor %m restrict : memref<5xf32> to tensor<5xf32>
+ %r = bufferization.to_buffer %t read_only : tensor<5xf32> to memref<5xf32>
+ %v = memref.load %r[%idx] : memref<5xf32>
+ return %v : f32
+}
More information about the Mlir-commits
mailing list