[Mlir-commits] [mlir] [mlir][bufferization] Alternate op bufferization approach (PR #199898)
Andrei Golubev
llvmlistbot at llvm.org
Wed May 27 00:53:08 PDT 2026
https://github.com/andrey-golubev created https://github.com/llvm/llvm-project/pull/199898
Make op bufferization go through type bufferization, by differentiating
between "known" and "unknown" cases better: when the type is known, i.e.
it is upstream-MLIR-aligned, TensorLikeType::getBufferType() would
dispatch to an op-provided bufferization callback; when the type is
unknown, unknown type converter would be called instead.
In return, this allows operation bufferization to be extended with
custom semantics without sacrificing the upstream's default behaviour.
>From 24f007b33eb6090f09b6b7a8d8c4c333730b560c Mon Sep 17 00:00:00 2001
From: "Golubev, Andrey" <andrey.golubev at intel.com>
Date: Thu, 21 May 2026 14:00:17 +0000
Subject: [PATCH 1/2] [NFC][mlir][bufferization] Simplify
TensorType::getBufferType()
The implementation for a builtin TensorType's TensorLikeType interface
is a bit overly complex for no reason: it calls `getMemRefType()`
whereas in reality `options.unknownTypeConverterFn()` would directly be
called due to the way the current code is written. Remove what is
effectively a wrapper function so that the current semantics is a bit
clearer.
---
mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
index bd177ba1afccd..b36aac8c68d36 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
@@ -48,7 +48,7 @@ struct BuiltinTensorExternalModel
return emitError() << "could not infer memory space";
return cast<BufferLikeType>(
- getMemRefType(tensorType, options, /*layout=*/{}, *memSpace));
+ options.unknownTypeConverterFn(tensorType, *memSpace, options));
}
mlir::LogicalResult verifyCompatibleBufferType(
>From 7a87ac64bd11f31d89f0eaccebfdaefb87ad76e8 Mon Sep 17 00:00:00 2001
From: "Golubev, Andrey" <andrey.golubev at intel.com>
Date: Tue, 26 May 2026 15:46:52 +0000
Subject: [PATCH 2/2] [mlir][bufferization] Alternate op bufferization approach
Make op bufferization go through type bufferization, by differentiating
between "known" and "unknown" cases better: when the type is known, i.e.
it is upstream-MLIR-aligned, TensorLikeType::getBufferType() would
dispatch to an op-provided bufferization callback; when the type is
unknown, unknown type converter would be called instead.
In return, this allows operation bufferization to be extended with
custom semantics without sacrificing the upstream's default behaviour.
---
.../IR/BufferizableOpInterface.h | 10 +++
.../IR/BufferizationTypeInterfaces.td | 3 +-
.../IR/BufferizableOpInterface.cpp | 66 ++++++++------
.../Bufferization/IR/BufferizationDialect.cpp | 13 ++-
.../Bufferization/IR/BufferizationOps.cpp | 43 +++++----
.../Bufferization/Transforms/Bufferize.cpp | 9 ++
.../SparsificationAndBufferizationPass.cpp | 9 ++
.../one-shot-non-module-bufferize.mlir | 38 --------
.../test-one-shot-module-bufferize.mlir | 90 +++++++++++++++++++
.../TestOneShotModuleBufferize.cpp | 11 ++-
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 54 +++--------
mlir/test/lib/Dialect/Test/TestTypeDefs.td | 13 +--
mlir/test/lib/Dialect/Test/TestTypes.cpp | 7 +-
13 files changed, 224 insertions(+), 142 deletions(-)
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 3f8392e3b8970..e8eeb3e96d4f1 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -272,6 +272,9 @@ struct BufferizationOptions {
// Produce a MemorySpace attribute from a tensor type
using DefaultMemorySpaceFn =
std::function<std::optional<Attribute>(TensorType t)>;
+ /// Returns whether the tensor type's encoding is bufferizable using upstream
+ /// bufferization.
+ using HasUpstreamBufferizableEncodingFn = std::function<bool(TensorType)>;
BufferizationOptions();
@@ -364,6 +367,13 @@ struct BufferizationOptions {
DefaultMemorySpaceFn defaultMemorySpaceFn =
[](TensorType t) -> std::optional<Attribute> { return Attribute(); };
+ /// Specifies whether a given tensor type has upstream compatible encoding to
+ /// be successfully bufferized using default upstream semantics. If the
+ /// encoding is not compatible, the unknown type converter is expected to
+ /// bufferize the type. By default, only tensor types without encoding are
+ /// considered compatible.
+ HasUpstreamBufferizableEncodingFn hasUpstreamBufferizableEncodingFn = nullptr;
+
/// If set to `true`, the analysis is skipped. A buffer is copied before every
/// write. This flag cannot be used together with `testAnalysisOnly = true`.
bool copyBeforeWrite = false;
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td
index fb6fc4f5ad964..92b72c06476c2 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td
@@ -31,7 +31,8 @@ def Bufferization_TensorLikeTypeInterface
/*methodName=*/"getBufferType",
/*args=*/(ins
"const ::mlir::bufferization::BufferizationOptions &":$options,
- "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError
+ "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError,
+ "::llvm::function_ref<::mlir::FailureOr<::mlir::bufferization::BufferLikeType>(::mlir::bufferization::TensorLikeType)>":$localGetBufferType
)
>,
InterfaceMethod<[{
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index f77edf23d4bc4..618dfbbdb2bc5 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -351,14 +351,14 @@ BufferLikeType
defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
func::FuncOp funcOp,
const BufferizationOptions &options) {
- if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
- return cast<BufferLikeType>(
- getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace));
- }
-
- // If not builtin, fallback to TensorLikeType::getBufferType()
- auto bufferType =
- type.getBufferType(options, [&]() { return funcOp->emitError(); });
+ auto bufferType = type.getBufferType(
+ options, [&]() { return funcOp->emitError(); },
+ [&](mlir::bufferization::TensorLikeType tensorLikeType)
+ -> mlir::FailureOr<mlir::bufferization::BufferLikeType> {
+ auto tensorType = cast<mlir::TensorType>(tensorLikeType);
+ return cast<BufferLikeType>(
+ getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace));
+ });
assert(succeeded(bufferType) &&
"a valid buffer is always expected at function boundary");
return *bufferType;
@@ -370,12 +370,23 @@ defaultUnknownTypeConverter(TensorType tensorType, Attribute memorySpace,
return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
}
+bool defaultHasUpstreamBufferizableEncoding(TensorType tensorType) {
+ if (isa<UnrankedTensorType>(tensorType)) {
+ // consider unranked tensor with no encoding bufferizable
+ return true;
+ }
+ const auto rankedTensorType = cast<RankedTensorType>(tensorType);
+ return rankedTensorType.getEncoding() == nullptr;
+}
+
} // namespace
// Default constructor for BufferizationOptions.
BufferizationOptions::BufferizationOptions()
: functionArgTypeConverterFn(defaultFunctionArgTypeConverter),
- unknownTypeConverterFn(defaultUnknownTypeConverter) {}
+ unknownTypeConverterFn(defaultUnknownTypeConverter),
+ hasUpstreamBufferizableEncodingFn(
+ defaultHasUpstreamBufferizableEncoding) {}
bool BufferizationOptions::isOpAllowed(Operation *op) const {
// Special case: If function boundary bufferization is deactivated, do not
@@ -407,19 +418,23 @@ void BufferizationOptions::setFunctionBoundaryTypeConversion(
functionArgTypeConverterFn = [=](TensorLikeType type, Attribute memorySpace,
func::FuncOp funcOp,
const BufferizationOptions &options) {
- if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
- if (layoutMapOption == LayoutMapOption::IdentityLayoutMap)
- return cast<BufferLikeType>(
- bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType,
- memorySpace));
- return cast<BufferLikeType>(
- bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
- memorySpace));
- }
-
- // If not builtin, fallback to TensorLikeType::getBufferType()
- auto bufferType =
- type.getBufferType(options, [&]() { return funcOp->emitError(); });
+ // allow generic bufferization to intercept the upstream behaviour
+ auto bufferType = type.getBufferType(
+ options, [&]() { return funcOp->emitError(); },
+ [&](mlir::bufferization::TensorLikeType tensorLikeType)
+ -> mlir::FailureOr<mlir::bufferization::BufferLikeType> {
+ if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
+ if (layoutMapOption == LayoutMapOption::IdentityLayoutMap)
+ return cast<BufferLikeType>(
+ bufferization::getMemRefTypeWithStaticIdentityLayout(
+ tensorType, memorySpace));
+ return cast<BufferLikeType>(
+ bufferization::getMemRefTypeWithFullyDynamicLayout(
+ tensorType, memorySpace));
+ }
+ return funcOp->emitError(
+ "unknown tensor-like type in function argument type converter");
+ });
assert(succeeded(bufferType) &&
"a valid buffer is always expected at function boundary");
return *bufferType;
@@ -739,9 +754,10 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
return bufferizableOp.getBufferType(value, options, state, invocationStack);
// Op is not bufferizable.
- return cast<TensorLikeType>(value.getType()).getBufferType(options, [&]() {
- return op->emitError();
- });
+ return cast<TensorLikeType>(value.getType())
+ .getBufferType(
+ options, [&]() { return op->emitError(); },
+ /*localGetBufferType=*/nullptr);
}
bool bufferization::hasTensorSemantics(Operation *op) {
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
index b36aac8c68d36..5ccf70f468ebe 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
@@ -41,8 +41,19 @@ struct BuiltinTensorExternalModel
Tensor> {
llvm::FailureOr<BufferLikeType> getBufferType(
mlir::Type tensor, const BufferizationOptions &options,
- llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
+ llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
+ llvm::function_ref<mlir::FailureOr<mlir::bufferization::BufferLikeType>(
+ mlir::bufferization::TensorLikeType)>
+ localGetBufferType) const {
auto tensorType = cast<TensorType>(tensor);
+ if (localGetBufferType &&
+ options.hasUpstreamBufferizableEncodingFn(tensorType)) {
+ return localGetBufferType(cast<TensorLikeType>(tensorType));
+ }
+
+ // if there's a non-bufferizable encoding, let unknown type converter handle
+ // the bufferization. assume that there's a non-default behaviour associated
+ // with it.
auto memSpace = options.defaultMemorySpaceFn(tensorType);
if (!memSpace.has_value())
return emitError() << "could not infer memory space";
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index c525ec116f699..5226219c6b2c7 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -229,25 +229,32 @@ AllocTensorOp::getBufferType(Value value, const BufferizationOptions &options,
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(getType())) {
- memorySpace = *ms;
- } else {
- return getOperation()->emitError("could not infer memory space");
- }
+ const auto defaultGetBufferType =
+ [&](TensorLikeType tensorLikeType) -> FailureOr<BufferLikeType> {
+ auto tensorType = cast<TensorType>(tensorLikeType);
+ // 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(tensorType)) {
+ memorySpace = *ms;
+ } else {
+ return getOperation()->emitError("could not infer memory space");
+ }
+
+ return cast<BufferLikeType>(
+ getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace));
+ };
- return cast<BufferLikeType>(
- getMemRefTypeWithStaticIdentityLayout(getType(), memorySpace));
+ return cast<TensorLikeType>(getType()).getBufferType(
+ options, [&]() { return emitError(); }, defaultGetBufferType);
}
LogicalResult AllocTensorOp::verify() {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 701ab52a491a8..d7b1423d8a728 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -91,6 +91,15 @@ struct OneShotBufferizePass
return rtt.getEncoding();
return std::nullopt;
};
+ opt.hasUpstreamBufferizableEncodingFn = [](TensorType t) {
+ if (isa<UnrankedTensorType>(t)) {
+ // consider unranked tensor with no encoding bufferizable
+ return true;
+ }
+ const auto rankedTensorType = cast<RankedTensorType>(t);
+ const auto encoding = rankedTensorType.getEncoding();
+ return !encoding || isa<IntegerAttr>(encoding);
+ };
}
opt.printConflicts = printConflicts;
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 7e8d3600293f8..14689422d3666 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -225,6 +225,15 @@ mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
const BufferizationOptions &options) {
return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
};
+ options.hasUpstreamBufferizableEncodingFn = [](TensorType tensorType) {
+ if (isa<UnrankedTensorType>(tensorType)) {
+ return true;
+ }
+ const auto rankedTensorType = cast<RankedTensorType>(tensorType);
+ const auto encoding = rankedTensorType.getEncoding();
+ return !encoding || isa<sparse_tensor::SparseTensorEncodingAttr>(encoding);
+ };
+
if (analysisOnly) {
options.testAnalysisOnly = true;
options.printConflicts = true;
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-non-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-non-module-bufferize.mlir
index b52612d0d1f10..09fdf8231b7cc 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-non-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-non-module-bufferize.mlir
@@ -29,41 +29,3 @@
}
"test.finish" () : () -> ()
}) : () -> ()
-
-// -----
-
-#enc1 = #test.tensor_encoding<"hello">
-#enc2 = #test.tensor_encoding<"not hello">
-
-"test.symbol_scope_isolated"() ({
- // CHECK: func @inner_func(
- // CHECK-SAME: %[[arg0:.*]]: memref<?xf32, #test.memref_layout<"hello">>)
- // CHECK-SAME: -> memref<?xf32, #test.memref_layout<"hello">>
- func.func @inner_func(%t: tensor<?xf32, #enc1>)
- -> tensor<?xf32, #enc1> {
- // CHECK: return %[[arg0]]
- return %t : tensor<?xf32, #enc1>
- }
-
- // CHECK: func @outer_func(
- // CHECK-SAME: %[[arg0:.*]]: memref<?xf32, #test.memref_layout<"hello">>)
- // CHECK-SAME: -> (memref<?xf32, #test.memref_layout<"hello">>,
- // CHECK-SAME: memref<?xf32, #test.memref_layout<"not hello">>)
- func.func @outer_func(%t0: tensor<?xf32, #enc1>)
- -> (tensor<?xf32, #enc1>, tensor<?xf32, #enc2>) {
- // CHECK: %[[call:.*]] = call @inner_func(%[[arg0]])
- %0 = call @inner_func(%t0)
- : (tensor<?xf32, #enc1>) -> (tensor<?xf32, #enc1>)
-
- // CHECK: %[[local:.*]] = "test.create_memref_op"() : ()
- // CHECK-SAME: -> memref<?xf32, #test.memref_layout<"not hello">>
- %local = "test.create_tensor_op"() : () -> tensor<?xf32, #enc2>
- // CHECK: %[[dummy:.*]] = "test.dummy_memref_op"(%[[local]])
- %1 = "test.dummy_tensor_op"(%local) : (tensor<?xf32, #enc2>)
- -> tensor<?xf32, #enc2>
-
- // CHECK: return %[[call]], %[[dummy]]
- return %0, %1 : tensor<?xf32, #enc1>, tensor<?xf32, #enc2>
- }
- "test.finish" () : () -> ()
-}) : () -> ()
diff --git a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
new file mode 100644
index 0000000000000..452a7746af46c
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
@@ -0,0 +1,90 @@
+// RUN: mlir-opt %s -test-one-shot-module-bufferize -split-input-file | FileCheck %s
+
+#enc1 = #test.tensor_encoding<"hello">
+#enc2 = #test.tensor_encoding<"not hello">
+
+module @BufferizeEncodingThroughFunctionBoundaryAndCustomOps {
+ // CHECK: func @inner_func(
+ // CHECK-SAME: %[[arg0:.*]]: memref<?xf32, #test.memref_layout<"hello">>)
+ // CHECK-SAME: -> memref<?xf32, #test.memref_layout<"hello">>
+ func.func @inner_func(%t: tensor<?xf32, #enc1>)
+ -> tensor<?xf32, #enc1> {
+ // CHECK: return %[[arg0]]
+ return %t : tensor<?xf32, #enc1>
+ }
+
+ // CHECK: func @outer_func(
+ // CHECK-SAME: %[[arg0:.*]]: memref<?xf32, #test.memref_layout<"hello">>)
+ // CHECK-SAME: -> (memref<?xf32, #test.memref_layout<"hello">>,
+ // CHECK-SAME: memref<?xf32, #test.memref_layout<"not hello">>)
+ func.func @outer_func(%t0: tensor<?xf32, #enc1>)
+ -> (tensor<?xf32, #enc1>, tensor<?xf32, #enc2>) {
+ // CHECK: %[[call:.*]] = call @inner_func(%[[arg0]])
+ %0 = call @inner_func(%t0)
+ : (tensor<?xf32, #enc1>) -> (tensor<?xf32, #enc1>)
+
+ // CHECK: %[[local:.*]] = "test.create_memref_op"() : ()
+ // CHECK-SAME: -> memref<?xf32, #test.memref_layout<"not hello">>
+ %local = "test.create_tensor_op"() : () -> tensor<?xf32, #enc2>
+ // CHECK: %[[dummy:.*]] = "test.dummy_memref_op"(%[[local]])
+ %1 = "test.dummy_tensor_op"(%local) : (tensor<?xf32, #enc2>)
+ -> tensor<?xf32, #enc2>
+
+ // CHECK: return %[[call]], %[[dummy]]
+ return %0, %1 : tensor<?xf32, #enc1>, tensor<?xf32, #enc2>
+ }
+}
+
+// -----
+
+#enc1 = #test.tensor_encoding<"hello">
+#enc2 = #test.tensor_encoding<"not hello">
+
+// The memref's layout must come from the encoding, not from the default
+// static-identity layout.
+module @BufferizeEncodingForAlloc {
+ // CHECK: func @some_func(
+ // CHECK-SAME: %[[arg0:.*]]: memref<42xf32, #test.memref_layout<"hello">>)
+ // CHECK-SAME: -> (memref<42xf32, #test.memref_layout<"hello">>,
+ // CHECK-SAME: memref<42xf32, #test.memref_layout<"not hello">>)
+ func.func @some_func(%t0: tensor<42xf32, #enc1>)
+ -> (tensor<42xf32, #enc1>, tensor<42xf32, #enc2>) {
+ // CHECK: %[[T0:.+]] = memref.alloc() {{.*}} : memref<42xf32, #test.memref_layout<"hello">>
+ %0 = bufferization.alloc_tensor() : tensor<42xf32, #enc1>
+
+ // CHECK: %[[T1:.+]] = memref.alloc() {{.*}} : memref<42xf32, #test.memref_layout<"not hello">>
+ %1 = bufferization.alloc_tensor() : tensor<42xf32, #enc2>
+
+ // CHECK: return %[[T0]], %[[T1]]
+ return %0, %1 : tensor<42xf32, #enc1>, tensor<42xf32, #enc2>
+ }
+}
+
+// -----
+
+#enc1 = #test.tensor_encoding<"custom">
+
+module @BufferizeEncodingForCustomOpsInsideScf {
+ // CHECK: func.func @custom_encoding_inside_scf(
+ // CHECK-SAME: %[[arg:.*]]: memref<42xf64, #test.memref_layout<"custom">>,
+ // CHECK-SAME: %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
+ // CHECK-SAME: -> memref<42xf64, #test.memref_layout<"custom">>
+ func.func @custom_encoding_inside_scf(
+ %arg: tensor<42xf64, #enc1>,
+ %lb: index, %ub: index, %step: index)
+ -> tensor<42xf64, #enc1> {
+ // CHECK: %[[loop:.+]] = scf.for %{{.*}} = %[[lb]] to %[[ub]] step %[[step]]
+ // CHECK-SAME: iter_args(%[[iter:.+]] = %[[arg]]) -> (memref<42xf64, #test.memref_layout<"custom">>) {
+ // CHECK: %[[call:.+]] = "test.dummy_memref_op"(%[[iter]])
+ // CHECK: scf.yield %[[call]] : memref<42xf64, #test.memref_layout<"custom">>
+ %loop = scf.for %i = %lb to %ub step %step
+ iter_args(%iter = %arg) -> (tensor<42xf64, #enc1>) {
+ %call = "test.dummy_tensor_op"(%iter) : (tensor<42xf64, #enc1>)
+ -> tensor<42xf64, #enc1>
+ scf.yield %call : tensor<42xf64, #enc1>
+ }
+
+ // CHECK: return %[[loop]]
+ return %loop : tensor<42xf64, #enc1>
+ }
+}
diff --git a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
index dead1a4b7e047..bbd4dc7421890 100644
--- a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
+++ b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
@@ -56,15 +56,14 @@ struct TestOneShotModuleBufferizePass
bufferization::OneShotBufferizationOptions opt;
opt.bufferizeFunctionBoundaries = true;
- opt.functionArgTypeConverterFn =
- [&](bufferization::TensorLikeType tensor, Attribute memSpace,
- func::FuncOp, const bufferization::BufferizationOptions &) {
+ opt.unknownTypeConverterFn =
+ [&](TensorType tensor, Attribute memSpace,
+ const bufferization::BufferizationOptions &) {
assert(isa<RankedTensorType>(tensor) && "tests only builtin tensors");
auto tensorType = cast<RankedTensorType>(tensor);
auto layout = getMemRefLayoutForTensorEncoding(tensorType);
- return cast<bufferization::BufferLikeType>(
- MemRefType::get(tensorType.getShape(),
- tensorType.getElementType(), layout, memSpace));
+ return MemRefType::get(tensorType.getShape(),
+ tensorType.getElementType(), layout, memSpace);
};
bufferization::BufferizationState bufferizationState;
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index a3ff397ac26db..1bbf32eb8a42d 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1769,39 +1769,6 @@ TestMultiSlotAlloca::handleDestructuringComplete(
return createNewMultiAllocaWithoutSlot(slot, builder, *this);
}
-namespace {
-/// Returns test dialect's memref layout for test dialect's tensor encoding when
-/// applicable.
-MemRefLayoutAttrInterface
-getMemRefLayoutForTensorEncoding(RankedTensorType tensorType) {
- if (auto encoding =
- dyn_cast<test::TestTensorEncodingAttr>(tensorType.getEncoding())) {
- return cast<MemRefLayoutAttrInterface>(test::TestMemRefLayoutAttr::get(
- tensorType.getContext(), encoding.getDummy()));
- }
- return {};
-}
-
-/// Auxiliary bufferization function for test and builtin tensors.
-bufferization::BufferLikeType
-convertTensorToBuffer(mlir::Operation *op,
- const bufferization::BufferizationOptions &options,
- bufferization::TensorLikeType tensorLike) {
- auto buffer =
- *tensorLike.getBufferType(options, [&]() { return op->emitError(); });
- if (auto memref = dyn_cast<MemRefType>(buffer)) {
- // Note: For the sake of testing, we want to ensure that encoding -> layout
- // bufferization happens. This is currently achieved manually.
- auto layout =
- getMemRefLayoutForTensorEncoding(cast<RankedTensorType>(tensorLike));
- return cast<bufferization::BufferLikeType>(
- MemRefType::get(memref.getShape(), memref.getElementType(), layout,
- memref.getMemorySpace()));
- }
- return buffer;
-}
-} // namespace
-
::mlir::LogicalResult test::TestDummyTensorOp::bufferize(
::mlir::RewriterBase &rewriter,
const ::mlir::bufferization::BufferizationOptions &options,
@@ -1811,12 +1778,16 @@ ::mlir::LogicalResult test::TestDummyTensorOp::bufferize(
if (mlir::failed(buffer))
return failure();
- const auto outType = getOutput().getType();
+ // Note: mlir::bufferization::getBufferType() would internally call
+ // TestDummyTensorOp::getBufferType()
const auto bufferizedOutType =
- convertTensorToBuffer(getOperation(), options, outType);
+ mlir::bufferization::getBufferType(getOutput(), options, state);
+ if (mlir::failed(bufferizedOutType))
+ return failure();
+
// replace op with memref analogy
auto dummyMemrefOp = test::TestDummyMemrefOp::create(
- rewriter, getLoc(), bufferizedOutType, *buffer);
+ rewriter, getLoc(), *bufferizedOutType, *buffer);
mlir::bufferization::replaceOpWithBufferizedValues(rewriter, getOperation(),
dummyMemrefOp.getResult());
@@ -1826,15 +1797,15 @@ ::mlir::LogicalResult test::TestDummyTensorOp::bufferize(
mlir::FailureOr<mlir::bufferization::BufferLikeType>
test::TestDummyTensorOp::getBufferType(
- mlir::Value value, const mlir::bufferization::BufferizationOptions &,
+ mlir::Value value, const mlir::bufferization::BufferizationOptions &options,
const mlir::bufferization::BufferizationState &,
llvm::SmallVector<::mlir::Value> &) {
- const auto type = dyn_cast<test::TestTensorType>(value.getType());
+ const auto type = dyn_cast<bufferization::TensorLikeType>(value.getType());
if (type == nullptr)
return failure();
- return cast<mlir::bufferization::BufferLikeType>(test::TestMemrefType::get(
- getContext(), type.getShape(), type.getElementType(), nullptr));
+ return type.getBufferType(
+ options, [&]() { return emitError(); }, /*localGetBufferType=*/nullptr);
}
::mlir::LogicalResult test::TestCreateTensorOp::bufferize(
@@ -1867,7 +1838,8 @@ test::TestCreateTensorOp::getBufferType(
if (type == nullptr)
return failure();
- return convertTensorToBuffer(getOperation(), options, type);
+ return type.getBufferType(
+ options, [&]() { return emitError(); }, /*localGetBufferType=*/nullptr);
}
// Define a custom builder for ManyRegionsOp declared in TestOps.td.
diff --git a/mlir/test/lib/Dialect/Test/TestTypeDefs.td b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
index 08600ce713a17..4121608cefa65 100644
--- a/mlir/test/lib/Dialect/Test/TestTypeDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
@@ -448,7 +448,9 @@ def TestTypeOpAsmTypeInterfaceTablegenDefault : Test_Type<"TestTypeOpAsmTypeInte
}
def TestTensorType : Test_Type<"TestTensor",
- [Bufferization_TensorLikeTypeInterface, ShapedTypeInterface]> {
+ [DeclareTypeInterfaceMethods<Bufferization_TensorLikeTypeInterface,
+ ["getBufferType", "verifyCompatibleBufferType"]>,
+ ShapedTypeInterface]> {
let mnemonic = "test_tensor";
let parameters = (ins
ArrayRefParameter<"int64_t">:$shape,
@@ -466,15 +468,6 @@ def TestTensorType : Test_Type<"TestTensor",
return test::TestTensorType::get(
getContext(), shape.value_or(getShape()), elementType);
}
-
- // TensorLikeTypeInterface:
- ::mlir::FailureOr<::mlir::bufferization::BufferLikeType>
- getBufferType(const ::mlir::bufferization::BufferizationOptions& options,
- ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError);
-
- ::mlir::LogicalResult verifyCompatibleBufferType(
- ::mlir::bufferization::BufferLikeType bufferType,
- ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError);
}];
}
diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp
index ef3396fc4f610..570a81121e7bc 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp
@@ -563,14 +563,17 @@ TestTypeOpAsmTypeInterfaceType::getAlias(::llvm::raw_ostream &os) const {
::mlir::FailureOr<::mlir::bufferization::BufferLikeType>
TestTensorType::getBufferType(
const ::mlir::bufferization::BufferizationOptions &,
- ::llvm::function_ref<::mlir::InFlightDiagnostic()>) {
+ ::llvm::function_ref<::mlir::InFlightDiagnostic()>,
+ ::llvm::function_ref<
+ ::mlir::FailureOr<::mlir::bufferization::BufferLikeType>(
+ ::mlir::bufferization::TensorLikeType)>) const {
return cast<bufferization::BufferLikeType>(
TestMemrefType::get(getContext(), getShape(), getElementType(), nullptr));
}
::mlir::LogicalResult TestTensorType::verifyCompatibleBufferType(
::mlir::bufferization::BufferLikeType bufferType,
- ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError) {
+ ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError) const {
if (auto testMemref = dyn_cast<TestMemrefType>(bufferType)) {
const bool valid = getShape() == testMemref.getShape() &&
getElementType() == testMemref.getElementType();
More information about the Mlir-commits
mailing list