[Mlir-commits] [mlir] [mlir][bufferization] Drop TensorLikeType::getBufferType() (PR #201350)
Andrei Golubev
llvmlistbot at llvm.org
Tue Jun 9 03:10:09 PDT 2026
https://github.com/andrey-golubev updated https://github.com/llvm/llvm-project/pull/201350
>From efcabaae21820598b6e2599699aab11225b3394e Mon Sep 17 00:00:00 2001
From: "Golubev, Andrey" <andrey.golubev at intel.com>
Date: Tue, 2 Jun 2026 14:38:09 +0000
Subject: [PATCH] [mlir][bufferization] Drop TensorLikeType::getBufferType()
Replace TensorLikeType::getBufferType() with
options.unknownTypeConverterFn() hook. Make the hook work with
tensor-like and buffer-like types (instead of builtins) to maintain the
same behaviour at the API boundary level and still allow custom types
to be properly supported.
Historically, an attempt to support custom types within the one-shot
bufferization framework was made. As part of it,
TensorLikeType::getBufferType() was introduced to allow custom types to
provide custom bufferization. However, the whole affair still proved to
be challenging: there is a completely custom API implementation for
custom types, and options-based (not sufficient) implementation for
builtin tensors. On top of this, there was always a function-specific
hook to customize function-level behaviour further. As a result of this,
users would need to implement two different mechanisms on their end:
interface implementation + custom options in order to achieve correct
bufferization.
It seems more reasonable at this stage to thus drop the interface part
and just allow unknown type conversion hook to be used directly. That
way, users would only need to supply that hook themselves (with the
caveat of having to perform a type dispatch inside of the hook
manually).
On the positive side, this unites builtin and custom tensors
implementation as they both feed into unknown type conversion.
On the negative side, default one-shot bufferization no longer supports
custom types directly - instead, one has to implement their own "pass",
overwriting the default bufferization options / hooks to customize the
behaviour.
As a drive by, multiple tests that exercise the usage of custom tensor
type (e.g. !test.test_tensor) are moved to work on top of a test pass.
Default memory space inference hook is also updated to work on top of
tensor-like type to align the signature to the unknown type conversion.
---
.../IR/BufferizableOpInterface.h | 10 +-
.../IR/BufferizationTypeInterfaces.td | 12 +-
.../BufferizableOpInterfaceImpl.cpp | 3 +-
.../IR/BufferizableOpInterface.cpp | 46 ++-
.../Bufferization/IR/BufferizationDialect.cpp | 12 -
.../Bufferization/IR/BufferizationOps.cpp | 3 +-
.../Bufferization/Transforms/Bufferize.cpp | 17 +-
.../FuncBufferizableOpInterfaceImpl.cpp | 6 +-
.../SparsificationAndBufferizationPass.cpp | 6 +-
.../BufferizableOpInterfaceImpl.cpp | 8 +-
.../Transforms/one-shot-bufferize.mlir | 38 ---
.../Transforms/one-shot-module-bufferize.mlir | 232 --------------
.../one-shot-non-module-bufferize.mlir | 38 ---
.../test-one-shot-module-bufferize.mlir | 298 ++++++++++++++++++
.../TestOneShotModuleBufferize.cpp | 39 ++-
.../Bufferization/TestTensorCopyInsertion.cpp | 4 +-
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 64 +---
mlir/test/lib/Dialect/Test/TestTypeDefs.td | 13 +-
mlir/test/lib/Dialect/Test/TestTypes.cpp | 10 +-
19 files changed, 407 insertions(+), 452 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 1d959dae56180..bf94f56497160 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -267,11 +267,11 @@ struct BufferizationOptions {
func::FuncOp, const BufferizationOptions &)>;
/// Tensor -> MemRef type conversion.
/// Parameters: tensor type, memory space, bufferization options
- using UnknownTypeConverterFn = std::function<BaseMemRefType(
- TensorType, Attribute memorySpace, const BufferizationOptions &)>;
+ using UnknownTypeConverterFn = std::function<BufferLikeType(
+ TensorLikeType, Attribute memorySpace, const BufferizationOptions &)>;
// Produce a MemorySpace attribute from a tensor type
using DefaultMemorySpaceFn =
- std::function<std::optional<Attribute>(TensorType t)>;
+ std::function<std::optional<Attribute>(TensorLikeType t)>;
BufferizationOptions();
@@ -340,7 +340,7 @@ struct BufferizationOptions {
///
/// By default, if tensor is a (builtin) tensor type, it is converted to a
/// memref type with a fully dynamic layout map; if tensor is a (generic)
- /// tensor-like type, it is converted using TensorLikeType::getBufferType().
+ /// tensor-like type, it is converted using unknownTypeConverterFn.
///
/// If `bufferizeFunctionBoundaries` is not set, this function isn't used.
FunctionArgTypeConverterFn functionArgTypeConverterFn = nullptr;
@@ -362,7 +362,7 @@ struct BufferizationOptions {
// Returning std::nullopt will cause bufferization to fail (useful to indicate
// failure to determine memory space for a tensor type).
DefaultMemorySpaceFn defaultMemorySpaceFn =
- [](TensorType t) -> std::optional<Attribute> { return Attribute(); };
+ [](TensorLikeType t) -> std::optional<Attribute> { return Attribute(); };
/// If set to `true`, the analysis is skipped. A buffer is copied before every
/// write. This flag cannot be used together with `testAnalysisOnly = true`.
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td
index fb6fc4f5ad964..6bbc6cb517766 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td
@@ -24,19 +24,9 @@ def Bufferization_TensorLikeTypeInterface
}];
let methods = [
- InterfaceMethod<[{
- Returns a BufferLike type for this TensorLike type.
- }],
- /*retTy=*/"::mlir::FailureOr<::mlir::bufferization::BufferLikeType>",
- /*methodName=*/"getBufferType",
- /*args=*/(ins
- "const ::mlir::bufferization::BufferizationOptions &":$options,
- "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError
- )
- >,
InterfaceMethod<[{
Returns whether a BufferLike type is compatible to this TensorLike type.
- The BufferLike type is assumed to be created by getBufferType().
+ The BufferLike type is assumed to be created by unknown type converter.
}],
/*retTy=*/"::mlir::LogicalResult",
/*methodName=*/"verifyCompatibleBufferType",
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index b073a31850678..91d03a0356de1 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -34,7 +34,8 @@ struct ConstantOpInterface
return failure();
Attribute memorySpace;
- if (auto memSpace = options.defaultMemorySpaceFn(type))
+ if (auto memSpace =
+ options.defaultMemorySpaceFn(cast<TensorLikeType>(type)))
memorySpace = *memSpace;
else
return constantOp->emitError("could not infer memory space");
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 55ec23aaf40c8..e60736a20dc3b 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -225,7 +225,8 @@ FailureOr<Value> bufferization::allocateTensorForShapedValue(
return failure();
std::optional<Attribute> memorySpace = copyBufferType->getMemorySpace();
if (!memorySpace)
- memorySpace = options.defaultMemorySpaceFn(tensorType);
+ memorySpace =
+ options.defaultMemorySpaceFn(cast<TensorLikeType>(tensorType));
if (memorySpace.has_value())
allocTensorOp.setMemorySpaceAttr(memorySpace.value());
return allocTensorOp.getResult();
@@ -356,18 +357,15 @@ defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace));
}
- // If not builtin, fallback to TensorLikeType::getBufferType()
- auto bufferType =
- type.getBufferType(options, [&]() { return funcOp->emitError(); });
- assert(succeeded(bufferType) &&
- "a valid buffer is always expected at function boundary");
- return *bufferType;
+ // If not builtin, fallback to unknown type conversion.
+ return options.unknownTypeConverterFn(type, memorySpace, options);
}
/// Default unknown type converter: Use a fully dynamic layout map.
-BaseMemRefType
-defaultUnknownTypeConverter(TensorType tensorType, Attribute memorySpace,
+BufferLikeType
+defaultUnknownTypeConverter(TensorLikeType tensorType, Attribute memorySpace,
const BufferizationOptions &options) {
- return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
+ return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(
+ cast<TensorType>(tensorType), memorySpace));
}
} // namespace
@@ -417,12 +415,8 @@ void BufferizationOptions::setFunctionBoundaryTypeConversion(
memorySpace));
}
- // If not builtin, fallback to TensorLikeType::getBufferType()
- auto bufferType =
- type.getBufferType(options, [&]() { return funcOp->emitError(); });
- assert(succeeded(bufferType) &&
- "a valid buffer is always expected at function boundary");
- return *bufferType;
+ // If not builtin, fallback to unknown type conversion.
+ return options.unknownTypeConverterFn(type, memorySpace, options);
};
inferFunctionResultLayout =
layoutMapOption == LayoutMapOption::InferLayoutMap;
@@ -739,9 +733,13 @@ 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();
- });
+ auto memSpace =
+ options.defaultMemorySpaceFn(cast<TensorLikeType>(value.getType()));
+ if (!memSpace.has_value())
+ return op->emitError("could not infer memory space");
+
+ return options.unknownTypeConverterFn(cast<TensorLikeType>(value.getType()),
+ *memSpace, options);
}
bool bufferization::hasTensorSemantics(Operation *op) {
@@ -957,8 +955,8 @@ FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
// No further analysis is possible for a block argument.
if (llvm::isa<BlockArgument>(value)) {
- return cast<BufferLikeType>(options.unknownTypeConverterFn(
- tensorType, /*memorySpace=*/nullptr, options));
+ return options.unknownTypeConverterFn(cast<TensorLikeType>(tensorType),
+ /*memorySpace=*/nullptr, options);
}
// Value is an OpResult.
@@ -978,12 +976,12 @@ FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
// If we do not know the memory space and there is no default memory space,
// report a failure.
auto memSpace =
- options.defaultMemorySpaceFn(cast<TensorType>(value.getType()));
+ options.defaultMemorySpaceFn(cast<TensorLikeType>(tensorType));
if (!memSpace.has_value())
return op->emitError("could not infer memory space");
- return cast<BufferLikeType>(
- options.unknownTypeConverterFn(tensorType, *memSpace, options));
+ return options.unknownTypeConverterFn(cast<TensorLikeType>(tensorType),
+ *memSpace, options);
}
bool bufferization::detail::defaultIsRepetitiveRegion(
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
index b36aac8c68d36..3bf2df7a3eae9 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
@@ -39,18 +39,6 @@ template <typename Tensor>
struct BuiltinTensorExternalModel
: TensorLikeType::ExternalModel<BuiltinTensorExternalModel<Tensor>,
Tensor> {
- llvm::FailureOr<BufferLikeType> getBufferType(
- mlir::Type tensor, const BufferizationOptions &options,
- llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
- auto tensorType = cast<TensorType>(tensor);
- auto memSpace = options.defaultMemorySpaceFn(tensorType);
- if (!memSpace.has_value())
- return emitError() << "could not infer memory space";
-
- return cast<BufferLikeType>(
- options.unknownTypeConverterFn(tensorType, *memSpace, options));
- }
-
mlir::LogicalResult verifyCompatibleBufferType(
mlir::Type tensor, BufferLikeType bufferType,
llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index c525ec116f699..2a1dcc89a83ed 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -240,7 +240,8 @@ AllocTensorOp::getBufferType(Value value, const BufferizationOptions &options,
if (failed(copyBufferType))
return failure();
memorySpace = copyBufferType->getMemorySpace();
- } else if (auto ms = options.defaultMemorySpaceFn(getType())) {
+ } else if (auto ms = options.defaultMemorySpaceFn(
+ cast<TensorLikeType>(getType()))) {
memorySpace = *ms;
} else {
return getOperation()->emitError("could not infer memory space");
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 84d5bb2d713c1..8bdee09d0d6e7 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -79,14 +79,14 @@ struct OneShotBufferizePass
if (mustInferMemorySpace) {
opt.defaultMemorySpaceFn =
- [](TensorType t) -> std::optional<Attribute> {
+ [](TensorLikeType t) -> std::optional<Attribute> {
return std::nullopt;
};
}
if (useEncodingForMemorySpace) {
opt.defaultMemorySpaceFn =
- [](TensorType t) -> std::optional<Attribute> {
+ [](TensorLikeType t) -> std::optional<Attribute> {
if (auto rtt = dyn_cast<RankedTensorType>(t))
return rtt.getEncoding();
return std::nullopt;
@@ -108,17 +108,20 @@ struct OneShotBufferizePass
"'unknown-type-conversion'");
return signalPassFailure();
}
- opt.unknownTypeConverterFn = [=](TensorType tensorType,
+ opt.unknownTypeConverterFn = [=](TensorLikeType type,
Attribute memorySpace,
const BufferizationOptions &options) {
+ const auto tensorType = cast<TensorType>(type);
if (unknownTypeConversionOption == LayoutMapOption::IdentityLayoutMap)
- return bufferization::getMemRefTypeWithStaticIdentityLayout(
- tensorType, memorySpace);
+ return cast<bufferization::BufferLikeType>(
+ bufferization::getMemRefTypeWithStaticIdentityLayout(
+ tensorType, memorySpace));
assert(unknownTypeConversionOption ==
LayoutMapOption::FullyDynamicLayoutMap &&
"invalid layout map option");
- return bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
- memorySpace);
+ return cast<bufferization::BufferLikeType>(
+ bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
+ memorySpace));
};
// Configure op filter.
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 3aaa38272935d..8ca968367b026 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -54,8 +54,8 @@ void FuncAnalysisState::startFunctionAnalysis(FuncOp funcOp) {
static mlir::Attribute
getDefaultMemorySpace(const BufferizationOptions &options,
TensorLikeType type) {
- if (auto tensorType = dyn_cast<TensorType>(type)) {
- return *options.defaultMemorySpaceFn(tensorType);
+ if (isa<TensorType>(type)) {
+ return *options.defaultMemorySpaceFn(type);
}
return nullptr;
}
@@ -73,7 +73,7 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
// Note: For builtin tensors there is additional logic related to layout.
if (auto tensorType = dyn_cast<TensorType>(type)) {
BufferLikeType memrefType = options.functionArgTypeConverterFn(
- type, *options.defaultMemorySpaceFn(tensorType), funcOp, options);
+ type, *options.defaultMemorySpaceFn(type), funcOp, options);
auto layoutAttr = funcOp.getArgAttrOfType<MemRefLayoutAttrInterface>(
index, BufferizationDialect::kBufferLayoutAttrName);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 7e8d3600293f8..b38d43476efb7 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -220,10 +220,12 @@ mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
OneShotBufferizationOptions options;
options.bufferizeFunctionBoundaries = true;
options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
- options.unknownTypeConverterFn = [](TensorType tensorType,
+ options.unknownTypeConverterFn = [](TensorLikeType tensorType,
Attribute memorySpace,
const BufferizationOptions &options) {
- return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
+ return cast<bufferization::BufferLikeType>(
+ getMemRefTypeWithStaticIdentityLayout(cast<TensorType>(tensorType),
+ memorySpace));
};
if (analysisOnly) {
options.testAnalysisOnly = true;
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 310e72587eb81..41a3d96b832d1 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -610,7 +610,7 @@ struct GenerateOpInterface
auto type = generateOp.getResult().getType();
// TODO: Implement memory space for this op.
- if (options.defaultMemorySpaceFn(type) != Attribute())
+ if (options.defaultMemorySpaceFn(cast<TensorLikeType>(type)) != Attribute())
return op->emitError("memory space not implemented yet");
// Allocate memory.
@@ -1062,7 +1062,8 @@ struct SplatOpInterface
auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
// TODO: Implement memory space for this op.
- if (options.defaultMemorySpaceFn(tensorType) != Attribute())
+ if (options.defaultMemorySpaceFn(cast<TensorLikeType>(tensorType)) !=
+ Attribute())
return op->emitError("memory space not implemented yet");
auto linalgOp = linalg::MapOp::create(rewriter, loc, tensorType,
@@ -1120,7 +1121,8 @@ struct ConcatOpInterface
auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
// TODO: Implement memory space for this op.
- if (options.defaultMemorySpaceFn(tensorType) != Attribute())
+ if (options.defaultMemorySpaceFn(cast<TensorLikeType>(tensorType)) !=
+ Attribute())
return op->emitError("memory space not implemented yet");
MemRefLayoutAttrInterface layout;
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
index 8031732011839..4cb08b68fffa0 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
@@ -269,41 +269,3 @@ func.func @materialize_in_dest_raw(%f: f32, %f2: f32, %idx: index) -> (tensor<5x
return %0, %r : tensor<5xf32>, f32
}
-
-// -----
-
-// CHECK: func.func @custom_op(
-// CHECK-SAME: %[[ARG:.*]]: !test.test_tensor<[32, 64], f64>
-// CHECK-SAME: ) -> !test.test_tensor<[32, 128], f64> {
-func.func @custom_op(%arg: !test.test_tensor<[32, 64], f64>)
- -> !test.test_tensor<[32, 128], f64> {
- // CHECK: %[[MEMREF:.*]] = bufferization.to_buffer %[[ARG]]
- // CHECK: %[[DUMMY:.*]] = "test.dummy_memref_op"(%[[MEMREF]])
- // CHECK-SAME: : (!test.test_memref<[32, 64], f64>)
- // CHECK-SAME: -> !test.test_memref<[32, 128], f64>
- // CHECK: %[[OUT:.*]] = bufferization.to_tensor %[[DUMMY]]
- %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[32, 64], f64>)
- -> !test.test_tensor<[32, 128], f64>
-
- // CHECK: return %[[OUT]]
- return %out : !test.test_tensor<[32, 128], f64>
-}
-
-// -----
-
-// CHECK: func.func @custom_origin_op()
-// CHECK-SAME: -> !test.test_tensor<[42], f64> {
-func.func @custom_origin_op() -> !test.test_tensor<[42], f64> {
- // CHECK: %[[MEMREF:.*]] = "test.create_memref_op"() : ()
- // CHECK-SAME: -> !test.test_memref<[21], f64>
- // CHECK: %[[DUMMY:.*]] = "test.dummy_memref_op"(%[[MEMREF]])
- // CHECK-SAME: : (!test.test_memref<[21], f64>)
- // CHECK-SAME: -> !test.test_memref<[42], f64>
- %in = "test.create_tensor_op"() : () -> !test.test_tensor<[21], f64>
- %out = "test.dummy_tensor_op"(%in) : (!test.test_tensor<[21], f64>)
- -> !test.test_tensor<[42], f64>
-
- // CHECK: %[[OUT:.*]] = bufferization.to_tensor %[[DUMMY]]
- // CHECK: return %[[OUT]]
- return %out : !test.test_tensor<[42], f64>
-}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index f8d5a1310ebdf..fee3c23367d9c 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -831,62 +831,6 @@ module @inner_module {
// -----
-// CHECK: func.func @custom_types(
-// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>
-// CHECK-SAME: ) -> (!test.test_memref<[4, 8], f64>,
-// CHECK-SAME: !test.test_memref<[4, 8], f64>)
-func.func @custom_types(%arg: !test.test_tensor<[4, 4], f64>)
- -> (!test.test_tensor<[4, 8], f64>, !test.test_tensor<[4, 8], f64>) {
- // CHECK: %[[out1:.*]] = "test.dummy_memref_op"(%[[arg]]) :
- // CHECK-SAME: (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 8], f64>
- %out1 = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 8], f64>
-
- // CHECK: %[[alloc:.*]] = "test.create_memref_op"
- // CHECK: %[[out2:.*]] = "test.dummy_memref_op"(%[[alloc]])
- // CHECK-SAME: (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 8], f64>
- %alloc = "test.create_tensor_op"() : () -> !test.test_tensor<[4, 4], f64>
- %out2 = "test.dummy_tensor_op"(%alloc) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 8], f64>
-
- // CHECK: return %[[out1]], %[[out2]]
- return %out1, %out2 :
- !test.test_tensor<[4, 8], f64>, !test.test_tensor<[4, 8], f64>
-}
-
-// -----
-
-// CHECK: func.func @custom_types_foo(
-// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>
-// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
-func.func @custom_types_foo(%arg: !test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64> {
- // CHECK: %[[out:.*]] = "test.dummy_memref_op"(%[[arg]])
- %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
- // CHECK: return %[[out]]
- return %out : !test.test_tensor<[4, 4], f64>
-}
-
-// CHECK: func.func @custom_types_bar(
-// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>
-// CHECK-SAME: ) -> !test.test_memref<[4, 8], f64>
-func.func @custom_types_bar(%arg: !test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 8], f64> {
- // CHECK: %[[call:.*]] = call @custom_types_foo(%[[arg]])
- %call = func.call @custom_types_foo(%arg) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
-
- // CHECK: %[[out:.*]] = "test.dummy_memref_op"(%[[call]])
- %out = "test.dummy_tensor_op"(%call) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 8], f64>
-
- // CHECK: return %[[out]]
- return %out : !test.test_tensor<[4, 8], f64>
-}
-
-// -----
-
// Test that foldMemRefCasts does not downgrade a ranked return type to unranked
// when the return value is produced by a memref.cast from unranked to ranked.
// CHECK-LABEL: func.func @ranked_return_via_unranked_call(
@@ -905,179 +849,3 @@ func.func @ranked_return_via_unranked_call(%arg0: tensor<64x20x40xf32>) -> tenso
return %b : tensor<64x20x40xf32>
}
func.func private @relu_unranked(tensor<*xf32>) -> tensor<*xf32>
-
-// -----
-
-// CHECK: func.func @custom_types_scf_for_inplace(
-// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>,
-// CHECK-SAME: %[[lb:.+]]: index, %[[ub:.+]]: index, %[[step:.+]]: index
-// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
-func.func @custom_types_scf_for_inplace(
- %arg: !test.test_tensor<[4, 4], f64>,
- %lb: index, %ub: index, %step: index)
- -> !test.test_tensor<[4, 4], f64> {
- // CHECK: %[[loop:.+]] = scf.for %{{.*}} = %[[lb]] to %[[ub]] step %[[step]]
- // CHECK-SAME: iter_args(%[[iter:.+]] = %[[arg]]) -> (!test.test_memref<[4, 4], f64>) {
- // CHECK: %[[call:.+]] = "test.dummy_memref_op"(%[[iter]])
- // CHECK: scf.yield %[[call]] : !test.test_memref<[4, 4], f64>
- %loop = scf.for %i = %lb to %ub step %step
- iter_args(%iter = %arg) -> (!test.test_tensor<[4, 4], f64>) {
- // Inside loop: use iter_args directly (this is inplace modifiable op)
- %call = "test.dummy_tensor_op"(%iter) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
- // Yield: return the same iter_args value (or result of inplace op on it)
- scf.yield %call : !test.test_tensor<[4, 4], f64>
- }
-
- // CHECK: return %[[loop]] : !test.test_memref<[4, 4], f64>
- return %loop : !test.test_tensor<[4, 4], f64>
-}
-
-// -----
-
-func.func private @custom_types_identity_2d(%arg: !test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64> {
- %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
- return %out : !test.test_tensor<[4, 4], f64>
-}
-
-// Same as @custom_types_scf_for_inplace, but with an inner call to test alias analysis
-// through function boundaries.
-// CHECK-LABEL: func.func @custom_types_scf_for_inplace_with_call(
-// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>
-// CHECK-SAME: %[[lb:.+]]: index, %[[ub:.+]]: index, %[[step:.+]]: index
-// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
-// CHECK: %[[loop:.+]] = scf.for %{{.*}} = %[[lb]] to %[[ub]] step %[[step]] iter_args(%[[iter:.+]] = %[[arg]]) -> (!test.test_memref<[4, 4], f64>) {
-// CHECK: %[[call:.+]] = func.call @custom_types_identity_2d(%[[iter]]) : (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 4], f64>
-// CHECK: scf.yield %[[call]] : !test.test_memref<[4, 4], f64>
-// CHECK: return %[[loop]] : !test.test_memref<[4, 4], f64>
-func.func @custom_types_scf_for_inplace_with_call(
- %arg: !test.test_tensor<[4, 4], f64>,
- %lb: index, %ub: index, %step: index)
- -> !test.test_tensor<[4, 4], f64> {
- %loop = scf.for %i = %lb to %ub step %step
- iter_args(%iter = %arg) -> (!test.test_tensor<[4, 4], f64>) {
- %call = func.call @custom_types_identity_2d(%iter)
- : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64>
- scf.yield %call : !test.test_tensor<[4, 4], f64>
- }
-
- return %loop : !test.test_tensor<[4, 4], f64>
-}
-
-// -----
-
-// CHECK-LABEL: func.func @custom_types_scf_if_inplace(
-// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>
-// CHECK-SAME: %[[cond:.+]]: i1
-// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
-// CHECK: %[[res:.+]] = scf.if %[[cond]] -> (!test.test_memref<[4, 4], f64>) {
-// CHECK: %[[dummy:.+]] = "test.dummy_memref_op"(%[[arg]]) : (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 4], f64>
-// CHECK: scf.yield %[[dummy]] : !test.test_memref<[4, 4], f64>
-// CHECK: } else {
-// CHECK: scf.yield %[[arg]] : !test.test_memref<[4, 4], f64>
-// CHECK: }
-// CHECK: return %[[res]] : !test.test_memref<[4, 4], f64>
-func.func @custom_types_scf_if_inplace(
- %arg: !test.test_tensor<[4, 4], f64>,
- %cond: i1)
- -> !test.test_tensor<[4, 4], f64> {
- %res = scf.if %cond -> (!test.test_tensor<[4, 4], f64>) {
- %dummy = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
- scf.yield %dummy : !test.test_tensor<[4, 4], f64>
- } else {
- scf.yield %arg : !test.test_tensor<[4, 4], f64>
- }
- return %res : !test.test_tensor<[4, 4], f64>
-}
-
-// -----
-
-func.func private @custom_types_identity_2d(%arg: !test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64> {
- %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
- return %out : !test.test_tensor<[4, 4], f64>
-}
-
-// CHECK-LABEL: func.func @custom_types_scf_if_inplace_with_call(
-// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>
-// CHECK-SAME: %[[cond:.+]]: i1
-// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
-// CHECK: %[[res:.+]] = scf.if %[[cond]] -> (!test.test_memref<[4, 4], f64>) {
-// CHECK: %[[call:.+]] = func.call @custom_types_identity_2d(%[[arg]]) : (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 4], f64>
-// CHECK: scf.yield %[[call]] : !test.test_memref<[4, 4], f64>
-// CHECK: } else {
-// CHECK: scf.yield %[[arg]] : !test.test_memref<[4, 4], f64>
-// CHECK: }
-// CHECK: return %[[res]] : !test.test_memref<[4, 4], f64>
-func.func @custom_types_scf_if_inplace_with_call(
- %arg: !test.test_tensor<[4, 4], f64>,
- %cond: i1)
- -> !test.test_tensor<[4, 4], f64> {
- %res = scf.if %cond -> (!test.test_tensor<[4, 4], f64>) {
- %call = func.call @custom_types_identity_2d(%arg)
- : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64>
- scf.yield %call : !test.test_tensor<[4, 4], f64>
- } else {
- scf.yield %arg : !test.test_tensor<[4, 4], f64>
- }
- return %res : !test.test_tensor<[4, 4], f64>
-}
-
-// -----
-
-// CHECK-LABEL: func.func @scf_while_inplace(
-// CHECK-SAME: !test.test_memref<[4, 4], f64>
-// CHECK: scf.while
-// CHECK: scf.condition
-// CHECK: scf.yield
-// CHECK: return
-func.func @scf_while_inplace(
- %arg: !test.test_tensor<[4, 4], f64>,
- %cond: i1)
- -> !test.test_tensor<[4, 4], f64> {
- %loop = scf.while (%iter = %arg)
- : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64> {
- scf.condition(%cond) %iter : !test.test_tensor<[4, 4], f64>
- } do {
- ^bb0(%current: !test.test_tensor<[4, 4], f64>):
- %dummy = "test.dummy_tensor_op"(%current) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
- scf.yield %dummy : !test.test_tensor<[4, 4], f64>
- }
- return %loop : !test.test_tensor<[4, 4], f64>
-}
-
-// -----
-
-func.func private @custom_types_identity_2d(%arg: !test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64> {
- %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
- -> !test.test_tensor<[4, 4], f64>
- return %out : !test.test_tensor<[4, 4], f64>
-}
-
-// CHECK-LABEL: func.func @scf_while_inplace(
-// CHECK-SAME: !test.test_memref<[4, 4], f64>
-// CHECK: scf.while
-// CHECK: scf.condition
-// CHECK: scf.yield
-// CHECK: return
-func.func @scf_while_inplace(
- %arg: !test.test_tensor<[4, 4], f64>,
- %cond: i1)
- -> !test.test_tensor<[4, 4], f64> {
- %loop = scf.while (%iter = %arg)
- : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64> {
- scf.condition(%cond) %iter : !test.test_tensor<[4, 4], f64>
- } do {
- ^bb0(%current: !test.test_tensor<[4, 4], f64>):
- %call = func.call @custom_types_identity_2d(%current)
- : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64>
- scf.yield %call : !test.test_tensor<[4, 4], f64>
- }
- return %loop : !test.test_tensor<[4, 4], f64>
-}
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..6ea68bec24b53
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
@@ -0,0 +1,298 @@
+// 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>
+ }
+}
+
+// -----
+
+// CHECK: func.func @custom_types(
+// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>
+// CHECK-SAME: ) -> (!test.test_memref<[4, 8], f64>,
+// CHECK-SAME: !test.test_memref<[4, 8], f64>)
+func.func @custom_types(%arg: !test.test_tensor<[4, 4], f64>)
+ -> (!test.test_tensor<[4, 8], f64>, !test.test_tensor<[4, 8], f64>) {
+ // CHECK: %[[out1:.*]] = "test.dummy_memref_op"(%[[arg]]) :
+ // CHECK-SAME: (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 8], f64>
+ %out1 = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 8], f64>
+
+ // CHECK: %[[alloc:.*]] = "test.create_memref_op"
+ // CHECK: %[[out2:.*]] = "test.dummy_memref_op"(%[[alloc]])
+ // CHECK-SAME: (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 8], f64>
+ %alloc = "test.create_tensor_op"() : () -> !test.test_tensor<[4, 4], f64>
+ %out2 = "test.dummy_tensor_op"(%alloc) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 8], f64>
+
+ // CHECK: return %[[out1]], %[[out2]]
+ return %out1, %out2 :
+ !test.test_tensor<[4, 8], f64>, !test.test_tensor<[4, 8], f64>
+}
+
+// -----
+
+// CHECK: func.func @custom_types_foo(
+// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>
+// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
+func.func @custom_types_foo(%arg: !test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64> {
+ // CHECK: %[[out:.*]] = "test.dummy_memref_op"(%[[arg]])
+ %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+ // CHECK: return %[[out]]
+ return %out : !test.test_tensor<[4, 4], f64>
+}
+
+// CHECK: func.func @custom_types_bar(
+// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>
+// CHECK-SAME: ) -> !test.test_memref<[4, 8], f64>
+func.func @custom_types_bar(%arg: !test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 8], f64> {
+ // CHECK: %[[call:.*]] = call @custom_types_foo(%[[arg]])
+ %call = func.call @custom_types_foo(%arg) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+
+ // CHECK: %[[out:.*]] = "test.dummy_memref_op"(%[[call]])
+ %out = "test.dummy_tensor_op"(%call) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 8], f64>
+
+ // CHECK: return %[[out]]
+ return %out : !test.test_tensor<[4, 8], f64>
+}
+
+
+// -----
+
+// CHECK: func.func @custom_types_scf_for_inplace(
+// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>,
+// CHECK-SAME: %[[lb:.+]]: index, %[[ub:.+]]: index, %[[step:.+]]: index
+// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
+func.func @custom_types_scf_for_inplace(
+ %arg: !test.test_tensor<[4, 4], f64>,
+ %lb: index, %ub: index, %step: index)
+ -> !test.test_tensor<[4, 4], f64> {
+ // CHECK: %[[loop:.+]] = scf.for %{{.*}} = %[[lb]] to %[[ub]] step %[[step]]
+ // CHECK-SAME: iter_args(%[[iter:.+]] = %[[arg]]) -> (!test.test_memref<[4, 4], f64>) {
+ // CHECK: %[[call:.+]] = "test.dummy_memref_op"(%[[iter]])
+ // CHECK: scf.yield %[[call]] : !test.test_memref<[4, 4], f64>
+ %loop = scf.for %i = %lb to %ub step %step
+ iter_args(%iter = %arg) -> (!test.test_tensor<[4, 4], f64>) {
+ // Inside loop: use iter_args directly (this is inplace modifiable op)
+ %call = "test.dummy_tensor_op"(%iter) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+ // Yield: return the same iter_args value (or result of inplace op on it)
+ scf.yield %call : !test.test_tensor<[4, 4], f64>
+ }
+
+ // CHECK: return %[[loop]] : !test.test_memref<[4, 4], f64>
+ return %loop : !test.test_tensor<[4, 4], f64>
+}
+
+// -----
+
+func.func private @custom_types_identity_2d(%arg: !test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64> {
+ %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+ return %out : !test.test_tensor<[4, 4], f64>
+}
+
+// Same as @custom_types_scf_for_inplace, but with an inner call to test alias analysis
+// through function boundaries.
+// CHECK-LABEL: func.func @custom_types_scf_for_inplace_with_call(
+// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>
+// CHECK-SAME: %[[lb:.+]]: index, %[[ub:.+]]: index, %[[step:.+]]: index
+// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
+// CHECK: %[[loop:.+]] = scf.for %{{.*}} = %[[lb]] to %[[ub]] step %[[step]] iter_args(%[[iter:.+]] = %[[arg]]) -> (!test.test_memref<[4, 4], f64>) {
+// CHECK: %[[call:.+]] = func.call @custom_types_identity_2d(%[[iter]]) : (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 4], f64>
+// CHECK: scf.yield %[[call]] : !test.test_memref<[4, 4], f64>
+// CHECK: return %[[loop]] : !test.test_memref<[4, 4], f64>
+func.func @custom_types_scf_for_inplace_with_call(
+ %arg: !test.test_tensor<[4, 4], f64>,
+ %lb: index, %ub: index, %step: index)
+ -> !test.test_tensor<[4, 4], f64> {
+ %loop = scf.for %i = %lb to %ub step %step
+ iter_args(%iter = %arg) -> (!test.test_tensor<[4, 4], f64>) {
+ %call = func.call @custom_types_identity_2d(%iter)
+ : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64>
+ scf.yield %call : !test.test_tensor<[4, 4], f64>
+ }
+
+ return %loop : !test.test_tensor<[4, 4], f64>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @custom_types_scf_if_inplace(
+// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>
+// CHECK-SAME: %[[cond:.+]]: i1
+// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
+// CHECK: %[[res:.+]] = scf.if %[[cond]] -> (!test.test_memref<[4, 4], f64>) {
+// CHECK: %[[dummy:.+]] = "test.dummy_memref_op"(%[[arg]]) : (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 4], f64>
+// CHECK: scf.yield %[[dummy]] : !test.test_memref<[4, 4], f64>
+// CHECK: } else {
+// CHECK: scf.yield %[[arg]] : !test.test_memref<[4, 4], f64>
+// CHECK: }
+// CHECK: return %[[res]] : !test.test_memref<[4, 4], f64>
+func.func @custom_types_scf_if_inplace(
+ %arg: !test.test_tensor<[4, 4], f64>,
+ %cond: i1)
+ -> !test.test_tensor<[4, 4], f64> {
+ %res = scf.if %cond -> (!test.test_tensor<[4, 4], f64>) {
+ %dummy = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+ scf.yield %dummy : !test.test_tensor<[4, 4], f64>
+ } else {
+ scf.yield %arg : !test.test_tensor<[4, 4], f64>
+ }
+ return %res : !test.test_tensor<[4, 4], f64>
+}
+
+// -----
+
+func.func private @custom_types_identity_2d(%arg: !test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64> {
+ %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+ return %out : !test.test_tensor<[4, 4], f64>
+}
+
+// CHECK-LABEL: func.func @custom_types_scf_if_inplace_with_call(
+// CHECK-SAME: %[[arg:.+]]: !test.test_memref<[4, 4], f64>
+// CHECK-SAME: %[[cond:.+]]: i1
+// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>
+// CHECK: %[[res:.+]] = scf.if %[[cond]] -> (!test.test_memref<[4, 4], f64>) {
+// CHECK: %[[call:.+]] = func.call @custom_types_identity_2d(%[[arg]]) : (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 4], f64>
+// CHECK: scf.yield %[[call]] : !test.test_memref<[4, 4], f64>
+// CHECK: } else {
+// CHECK: scf.yield %[[arg]] : !test.test_memref<[4, 4], f64>
+// CHECK: }
+// CHECK: return %[[res]] : !test.test_memref<[4, 4], f64>
+func.func @custom_types_scf_if_inplace_with_call(
+ %arg: !test.test_tensor<[4, 4], f64>,
+ %cond: i1)
+ -> !test.test_tensor<[4, 4], f64> {
+ %res = scf.if %cond -> (!test.test_tensor<[4, 4], f64>) {
+ %call = func.call @custom_types_identity_2d(%arg)
+ : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64>
+ scf.yield %call : !test.test_tensor<[4, 4], f64>
+ } else {
+ scf.yield %arg : !test.test_tensor<[4, 4], f64>
+ }
+ return %res : !test.test_tensor<[4, 4], f64>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @scf_while_inplace(
+// CHECK-SAME: !test.test_memref<[4, 4], f64>
+// CHECK: scf.while
+// CHECK: scf.condition
+// CHECK: scf.yield
+// CHECK: return
+func.func @scf_while_inplace(
+ %arg: !test.test_tensor<[4, 4], f64>,
+ %cond: i1)
+ -> !test.test_tensor<[4, 4], f64> {
+ %loop = scf.while (%iter = %arg)
+ : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64> {
+ scf.condition(%cond) %iter : !test.test_tensor<[4, 4], f64>
+ } do {
+ ^bb0(%current: !test.test_tensor<[4, 4], f64>):
+ %dummy = "test.dummy_tensor_op"(%current) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+ scf.yield %dummy : !test.test_tensor<[4, 4], f64>
+ }
+ return %loop : !test.test_tensor<[4, 4], f64>
+}
+
+// -----
+
+func.func private @custom_types_identity_2d(%arg: !test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64> {
+ %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)
+ -> !test.test_tensor<[4, 4], f64>
+ return %out : !test.test_tensor<[4, 4], f64>
+}
+
+// CHECK-LABEL: func.func @scf_while_inplace(
+// CHECK-SAME: !test.test_memref<[4, 4], f64>
+// CHECK: scf.while
+// CHECK: scf.condition
+// CHECK: scf.yield
+// CHECK: return
+func.func @scf_while_inplace(
+ %arg: !test.test_tensor<[4, 4], f64>,
+ %cond: i1)
+ -> !test.test_tensor<[4, 4], f64> {
+ %loop = scf.while (%iter = %arg)
+ : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64> {
+ scf.condition(%cond) %iter : !test.test_tensor<[4, 4], f64>
+ } do {
+ ^bb0(%current: !test.test_tensor<[4, 4], f64>):
+ %call = func.call @custom_types_identity_2d(%current)
+ : (!test.test_tensor<[4, 4], f64>) -> !test.test_tensor<[4, 4], f64>
+ scf.yield %call : !test.test_tensor<[4, 4], f64>
+ }
+ return %loop : !test.test_tensor<[4, 4], f64>
+}
+
+// -----
+
+#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..87f95e36e9183 100644
--- a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
+++ b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
@@ -58,13 +58,38 @@ struct TestOneShotModuleBufferizePass
opt.bufferizeFunctionBoundaries = true;
opt.functionArgTypeConverterFn =
[&](bufferization::TensorLikeType tensor, Attribute memSpace,
- func::FuncOp, 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));
+ func::FuncOp, const bufferization::BufferizationOptions &options) {
+ return options.unknownTypeConverterFn(tensor, memSpace, options);
+ };
+ opt.unknownTypeConverterFn =
+ [&](bufferization::TensorLikeType tensor, Attribute memSpace,
+ const bufferization::BufferizationOptions &) {
+ return llvm::TypeSwitch<bufferization::TensorLikeType,
+ bufferization::BufferLikeType>(tensor)
+ .Case([&](UnrankedTensorType unrankedTensorType) {
+ return cast<bufferization::BufferLikeType>(
+ UnrankedMemRefType::get(unrankedTensorType.getElementType(),
+ memSpace));
+ })
+ .Case([&](RankedTensorType rankedTensorType) {
+ // Note: builtin ranked tensor with custom encoding to layout
+ // conversion.
+ auto layout =
+ getMemRefLayoutForTensorEncoding(rankedTensorType);
+ return cast<bufferization::BufferLikeType>(MemRefType::get(
+ rankedTensorType.getShape(),
+ rankedTensorType.getElementType(), layout, memSpace));
+ })
+ .Case([&](test::TestTensorType testTensorType)
+ -> bufferization::BufferLikeType {
+ return test::TestMemrefType::get(
+ testTensorType.getContext(), testTensorType.getShape(),
+ testTensorType.getElementType(), memSpace);
+ })
+ .Default([&](bufferization::TensorLikeType tensor) {
+ llvm_unreachable("unexpected tensor type");
+ return bufferization::BufferLikeType{};
+ });
};
bufferization::BufferizationState bufferizationState;
diff --git a/mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp b/mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp
index dfaebccde7dcc..dc971ad68b846 100644
--- a/mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp
+++ b/mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp
@@ -46,7 +46,9 @@ struct TestTensorCopyInsertionPass
options.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
if (mustInferMemorySpace) {
options.defaultMemorySpaceFn =
- [](TensorType t) -> std::optional<Attribute> { return std::nullopt; };
+ [](bufferization::TensorLikeType t) -> std::optional<Attribute> {
+ return std::nullopt;
+ };
}
bufferization::BufferizationState bufferizationState;
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 8315bd7cef783..41f5a6f69fa21 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1950,39 +1950,6 @@ Value TestPartialAlias::projectAliasValueToSlotValue(
.getResult(0);
}
-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,
@@ -1992,12 +1959,14 @@ ::mlir::LogicalResult test::TestDummyTensorOp::bufferize(
if (mlir::failed(buffer))
return failure();
- const auto outType = getOutput().getType();
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());
@@ -2007,15 +1976,14 @@ ::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());
- if (type == nullptr)
- return failure();
-
- return cast<mlir::bufferization::BufferLikeType>(test::TestMemrefType::get(
- getContext(), type.getShape(), type.getElementType(), nullptr));
+ // let unknown type converter handle the types - assume test operations work
+ // with test-one-shot-module-bufferize pass.
+ return options.unknownTypeConverterFn(
+ cast<mlir::bufferization::TensorLikeType>(value.getType()), nullptr,
+ options);
}
::mlir::LogicalResult test::TestCreateTensorOp::bufferize(
@@ -2044,11 +2012,11 @@ test::TestCreateTensorOp::getBufferType(
mlir::Value value, const mlir::bufferization::BufferizationOptions &options,
const mlir::bufferization::BufferizationState &,
llvm::SmallVector<::mlir::Value> &) {
- const auto type = dyn_cast<bufferization::TensorLikeType>(value.getType());
- if (type == nullptr)
- return failure();
-
- return convertTensorToBuffer(getOperation(), options, type);
+ // let unknown type converter handle the types - assume test operations work
+ // with test-one-shot-module-bufferize pass.
+ return options.unknownTypeConverterFn(
+ cast<mlir::bufferization::TensorLikeType>(value.getType()), nullptr,
+ options);
}
// 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..f405940358ca0 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,
+ ["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..7d1cfe8c9ba06 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp
@@ -560,17 +560,9 @@ TestTypeOpAsmTypeInterfaceType::getAlias(::llvm::raw_ostream &os) const {
return ::mlir::OpAsmDialectInterface::AliasResult::FinalAlias;
}
-::mlir::FailureOr<::mlir::bufferization::BufferLikeType>
-TestTensorType::getBufferType(
- const ::mlir::bufferization::BufferizationOptions &,
- ::llvm::function_ref<::mlir::InFlightDiagnostic()>) {
- 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