[Mlir-commits] [mlir] [mlir] Reject invalid boundary types (PR #208787)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 21 22:53:01 PDT 2026
https://github.com/mygitljf updated https://github.com/llvm/llvm-project/pull/208787
>From a66e3ae961cdcc17f6645dbf575f08964e7782a2 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Sat, 11 Jul 2026 01:29:55 +0000
Subject: [PATCH 1/2] [mlir][bufferization] Reject invalid boundary types
---
.../FuncBufferizableOpInterfaceImpl.cpp | 29 +++++++++++++++++--
...est-one-shot-module-bufferize-invalid.mlir | 8 +++++
2 files changed, 34 insertions(+), 3 deletions(-)
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 8ca968367b026..f38de559353a8 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -60,16 +60,32 @@ getDefaultMemorySpace(const BufferizationOptions &options,
return nullptr;
}
+static LogicalResult verifyMemRefElementType(FuncOp funcOp,
+ TensorLikeType type) {
+ auto tensorType = dyn_cast<TensorType>(type);
+ if (!tensorType)
+ return success();
+ Type elementType = tensorType.getElementType();
+ if (BaseMemRefType::isValidElementType(elementType))
+ return success();
+ return funcOp.emitError() << "cannot bufferize function boundary type "
+ << tensorType << ": element type " << elementType
+ << " is not a valid memref element type";
+}
+
/// Return the index-th bufferized function argument type. This assumes that the
/// specified argument is a tensor. If the tensor is ranked, a layout map may be
/// specified by the user (as per `options.functionArgTypeConverterFn`).
-static BufferLikeType
+static FailureOr<BufferLikeType>
getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
const BufferizationOptions &options) {
auto type =
dyn_cast<TensorLikeType>(funcOp.getFunctionType().getInput(index));
assert(type && "expected TensorLikeType");
+ if (failed(verifyMemRefElementType(funcOp, type)))
+ return failure();
+
// Note: For builtin tensors there is additional logic related to layout.
if (auto tensorType = dyn_cast<TensorType>(type)) {
BufferLikeType memrefType = options.functionArgTypeConverterFn(
@@ -245,6 +261,8 @@ struct CallOpInterface
// Otherwise, call the type converter to compute the bufferized type.
auto tensorType = cast<TensorLikeType>(resultType);
+ if (failed(verifyMemRefElementType(funcOp, tensorType)))
+ return failure();
return cast<BufferLikeType>(options.functionArgTypeConverterFn(
tensorType, getDefaultMemorySpace(options, tensorType), funcOp,
options));
@@ -444,8 +462,11 @@ struct FuncOpInterface
for (const auto &it : llvm::enumerate(funcType.getInputs())) {
Type argType = it.value();
if (isa<TensorLikeType>(argType)) {
- argTypes.push_back(
- getBufferizedFunctionArgType(funcOp, it.index(), options));
+ FailureOr<BufferLikeType> bufferType =
+ getBufferizedFunctionArgType(funcOp, it.index(), options);
+ if (failed(bufferType))
+ return failure();
+ argTypes.push_back(*bufferType);
continue;
}
argTypes.push_back(argType);
@@ -455,6 +476,8 @@ struct FuncOpInterface
SmallVector<Type> retTypes;
for (Type resultType : funcType.getResults()) {
if (auto tensorType = dyn_cast<TensorLikeType>(resultType)) {
+ if (failed(verifyMemRefElementType(funcOp, tensorType)))
+ return failure();
BufferLikeType resultType = options.functionArgTypeConverterFn(
tensorType, getDefaultMemorySpace(options, tensorType), funcOp,
options);
diff --git a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir
new file mode 100644
index 0000000000000..4a2de50db1de9
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir
@@ -0,0 +1,8 @@
+// RUN: mlir-opt %s -test-one-shot-module-bufferize -verify-diagnostics
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1" -verify-diagnostics
+
+module {
+ // expected-error @below {{cannot bufferize function boundary type 'tensor<!llvm.array<1 x i32>>': element type '!llvm.array<1 x i32>' is not a valid memref element type}}
+ // expected-error @below {{failed to bufferize op}}
+ func.func private @sparse_csr(tensor<f64>) -> tensor<!llvm.array<1 x i32>>
+}
>From cad752a93f49ac53d2a421032428de11b1aff565 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Wed, 22 Jul 2026 13:52:12 +0800
Subject: [PATCH 2/2] [mlir][bufferization] Validate function boundary
conversions
---
.../IR/BufferizableOpInterface.h | 1 +
.../IR/BufferizableOpInterface.cpp | 14 +++++
.../FuncBufferizableOpInterfaceImpl.cpp | 61 +++++++++----------
...-shot-module-bufferize-invalid-layout.mlir | 29 +++++++++
...est-one-shot-module-bufferize-invalid.mlir | 12 +++-
.../test-one-shot-module-bufferize.mlir | 34 +++++++++++
.../TestOneShotModuleBufferize.cpp | 22 ++++++-
7 files changed, 138 insertions(+), 35 deletions(-)
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid-layout.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 9f7a6a638f307..6dcadee7f0fee 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -262,6 +262,7 @@ struct BufferizationOptions {
using AnalysisStateInitFn = std::function<void(AnalysisState &)>;
/// Tensor-like -> Buffer-like type conversion.
/// Parameters: tensor-like type, memory space, func op, bufferization options
+ /// Returning a null type signals conversion failure.
using FunctionArgTypeConverterFn =
std::function<BufferLikeType(TensorLikeType, Attribute memorySpace,
func::FuncOp, const BufferizationOptions &)>;
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index e0a35e1932f3d..ff407bfa35251 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -358,12 +358,24 @@ BaseMemRefType getMemRefTypeWithFullyDynamicLayout(ArrayRef<int64_t> shape,
return MemRefType::get(shape, elementType, stridedLayout, memorySpace);
}
+LogicalResult verifyMemRefElementType(func::FuncOp funcOp,
+ TensorType tensorType) {
+ Type elementType = tensorType.getElementType();
+ if (BaseMemRefType::isValidElementType(elementType))
+ return success();
+ return funcOp.emitError() << "cannot bufferize function boundary type "
+ << tensorType << ": element type " << elementType
+ << " is not a valid memref element type";
+}
+
/// Default function arg type converter: Use a fully dynamic layout map.
BufferLikeType
defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
func::FuncOp funcOp,
const BufferizationOptions &options) {
if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
+ if (failed(verifyMemRefElementType(funcOp, tensorType)))
+ return {};
return cast<BufferLikeType>(
bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
memorySpace));
@@ -440,6 +452,8 @@ void BufferizationOptions::setFunctionBoundaryTypeConversion(
func::FuncOp funcOp,
const BufferizationOptions &options) {
if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
+ if (failed(verifyMemRefElementType(funcOp, tensorType)))
+ return BufferLikeType{};
if (layoutMapOption == LayoutMapOption::IdentityLayoutMap)
return cast<BufferLikeType>(
bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType,
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index f38de559353a8..90023bebda1bd 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -60,19 +60,6 @@ getDefaultMemorySpace(const BufferizationOptions &options,
return nullptr;
}
-static LogicalResult verifyMemRefElementType(FuncOp funcOp,
- TensorLikeType type) {
- auto tensorType = dyn_cast<TensorType>(type);
- if (!tensorType)
- return success();
- Type elementType = tensorType.getElementType();
- if (BaseMemRefType::isValidElementType(elementType))
- return success();
- return funcOp.emitError() << "cannot bufferize function boundary type "
- << tensorType << ": element type " << elementType
- << " is not a valid memref element type";
-}
-
/// Return the index-th bufferized function argument type. This assumes that the
/// specified argument is a tensor. If the tensor is ranked, a layout map may be
/// specified by the user (as per `options.functionArgTypeConverterFn`).
@@ -83,29 +70,34 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
dyn_cast<TensorLikeType>(funcOp.getFunctionType().getInput(index));
assert(type && "expected TensorLikeType");
- if (failed(verifyMemRefElementType(funcOp, type)))
- return failure();
-
// Note: For builtin tensors there is additional logic related to layout.
if (auto tensorType = dyn_cast<TensorType>(type)) {
- BufferLikeType memrefType = options.functionArgTypeConverterFn(
+ BufferLikeType bufferType = options.functionArgTypeConverterFn(
type, *options.defaultMemorySpaceFn(type), funcOp, options);
+ if (!bufferType)
+ return failure();
auto layoutAttr = funcOp.getArgAttrOfType<MemRefLayoutAttrInterface>(
index, BufferizationDialect::kBufferLayoutAttrName);
if (!layoutAttr)
- return memrefType;
+ return bufferType;
- auto rankedMemrefType = dyn_cast<MemRefType>(memrefType);
- assert(rankedMemrefType &&
- "buffer layout not supported on unranked tensors");
+ auto rankedMemrefType = dyn_cast<MemRefType>(bufferType);
+ if (!rankedMemrefType) {
+ funcOp.emitError() << "cannot apply buffer layout to buffer type "
+ << bufferType;
+ return failure();
+ }
return cast<BufferLikeType>(MemRefType::get(
rankedMemrefType.getShape(), rankedMemrefType.getElementType(),
layoutAttr, rankedMemrefType.getMemorySpace()));
}
- return options.functionArgTypeConverterFn(type, /*memSpace=*/nullptr, funcOp,
- options);
+ BufferLikeType bufferType = options.functionArgTypeConverterFn(
+ type, /*memSpace=*/nullptr, funcOp, options);
+ if (!bufferType)
+ return failure();
+ return bufferType;
}
/// Return the FuncOp called by `callOp`.
@@ -261,11 +253,12 @@ struct CallOpInterface
// Otherwise, call the type converter to compute the bufferized type.
auto tensorType = cast<TensorLikeType>(resultType);
- if (failed(verifyMemRefElementType(funcOp, tensorType)))
- return failure();
- return cast<BufferLikeType>(options.functionArgTypeConverterFn(
+ BufferLikeType bufferType = options.functionArgTypeConverterFn(
tensorType, getDefaultMemorySpace(options, tensorType), funcOp,
- options));
+ options);
+ if (!bufferType)
+ return failure();
+ return bufferType;
}
/// All function arguments are writable. It is the responsibility of the
@@ -337,9 +330,13 @@ struct CallOpInterface
// something better. Insert a reallocation + copy if it cannot be
// statically guaranteed that a direct cast would be valid.
if (buffer.getType() != bufferType) {
- auto memrefDstType = dyn_cast<MemRefType>(bufferType);
- assert(memrefDstType &&
- "buffer layout not supported on unranked tensors");
+ if (!isa<MemRefType>(buffer.getType()) ||
+ !isa<MemRefType>(bufferType)) {
+ callOp.emitError() << "cannot reconcile buffer types "
+ << buffer.getType() << " and " << bufferType;
+ return failure();
+ }
+ auto memrefDstType = cast<MemRefType>(bufferType);
FailureOr<Value> replacement = bufferization::castOrReallocMemRefValue(
rewriter, buffer, memrefDstType, options);
if (failed(replacement))
@@ -476,11 +473,11 @@ struct FuncOpInterface
SmallVector<Type> retTypes;
for (Type resultType : funcType.getResults()) {
if (auto tensorType = dyn_cast<TensorLikeType>(resultType)) {
- if (failed(verifyMemRefElementType(funcOp, tensorType)))
- return failure();
BufferLikeType resultType = options.functionArgTypeConverterFn(
tensorType, getDefaultMemorySpace(options, tensorType), funcOp,
options);
+ if (!resultType)
+ return failure();
retTypes.push_back(resultType);
continue;
}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid-layout.mlir b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid-layout.mlir
new file mode 100644
index 0000000000000..b1fdb44c4193f
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid-layout.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s -test-one-shot-module-bufferize -verify-diagnostics -split-input-file
+
+#custom_buffer = #test.tensor_encoding<"custom_buffer">
+
+module {
+ // expected-error @below {{cannot apply buffer layout to buffer type '!test.test_memref<[1], !llvm.array<1 x i32>>'}}
+ // expected-error @below {{failed to bufferize op}}
+ func.func private @custom_buffer_with_layout(
+ tensor<1x!llvm.array<1 x i32>, #custom_buffer>
+ {bufferization.buffer_layout = affine_map<(d0) -> (d0)>})
+}
+
+// -----
+
+#custom_buffer = #test.tensor_encoding<"custom_buffer">
+
+module {
+ func.func private @custom_buffer_callee(
+ tensor<1xf32, #custom_buffer>)
+
+ func.func @custom_buffer_caller() {
+ %tensor = "test.create_tensor_op"() : () -> tensor<1xf32, #custom_buffer>
+ // expected-error @below {{cannot reconcile buffer types 'memref<1xf32, #test.memref_layout<"custom_buffer">>' and '!test.test_memref<[1], f32>'}}
+ // expected-error @below {{failed to bufferize op}}
+ func.call @custom_buffer_callee(%tensor)
+ : (tensor<1xf32, #custom_buffer>) -> ()
+ return
+ }
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir
index 4a2de50db1de9..410d74c2a44eb 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize-invalid.mlir
@@ -1,8 +1,16 @@
-// RUN: mlir-opt %s -test-one-shot-module-bufferize -verify-diagnostics
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1" -verify-diagnostics
+// RUN: mlir-opt %s -test-one-shot-module-bufferize -verify-diagnostics -split-input-file
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1" -verify-diagnostics -split-input-file
module {
// expected-error @below {{cannot bufferize function boundary type 'tensor<!llvm.array<1 x i32>>': element type '!llvm.array<1 x i32>' is not a valid memref element type}}
// expected-error @below {{failed to bufferize op}}
func.func private @sparse_csr(tensor<f64>) -> tensor<!llvm.array<1 x i32>>
}
+
+// -----
+
+module {
+ // expected-error @below {{cannot bufferize function boundary type 'tensor<!llvm.array<1 x i32>>': element type '!llvm.array<1 x i32>' is not a valid memref element type}}
+ // expected-error @below {{failed to bufferize op}}
+ func.func private @invalid_arg(tensor<!llvm.array<1 x i32>>)
+}
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
index b2674e9497927..a3701e45fa14c 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
@@ -92,6 +92,40 @@ func.func @custom_types_bar(%arg: !test.test_tensor<[4, 4], f64>)
return %out : !test.test_tensor<[4, 8], f64>
}
+// -----
+
+#custom_buffer = #test.tensor_encoding<"custom_buffer">
+
+// CHECK-LABEL: func.func @custom_buffer_like_callee(
+// CHECK-SAME: %[[ARG:.*]]: !test.test_memref<[1], !llvm.array<1 x i32>>
+// CHECK-SAME: ) -> !test.test_memref<[1], !llvm.array<1 x i32>>
+func.func @custom_buffer_like_callee(
+ %arg: tensor<1x!llvm.array<1 x i32>, #custom_buffer>)
+ -> tensor<1x!llvm.array<1 x i32>, #custom_buffer> {
+ // CHECK: %[[CALL:.*]] = call @custom_buffer_like_caller(%[[ARG]])
+ // CHECK-SAME: (!test.test_memref<[1], !llvm.array<1 x i32>>)
+ // CHECK-SAME: -> !test.test_memref<[1], !llvm.array<1 x i32>>
+ %call = func.call @custom_buffer_like_caller(%arg)
+ : (tensor<1x!llvm.array<1 x i32>, #custom_buffer>)
+ -> tensor<1x!llvm.array<1 x i32>, #custom_buffer>
+ return %call : tensor<1x!llvm.array<1 x i32>, #custom_buffer>
+}
+
+// CHECK-LABEL: func.func @custom_buffer_like_caller(
+// CHECK-SAME: %[[ARG:.*]]: !test.test_memref<[1], !llvm.array<1 x i32>>
+// CHECK-SAME: ) -> !test.test_memref<[1], !llvm.array<1 x i32>>
+func.func @custom_buffer_like_caller(
+ %arg: tensor<1x!llvm.array<1 x i32>, #custom_buffer>)
+ -> tensor<1x!llvm.array<1 x i32>, #custom_buffer> {
+ // CHECK: %[[CALL:.*]] = call @custom_buffer_like_callee(%[[ARG]])
+ // CHECK-SAME: (!test.test_memref<[1], !llvm.array<1 x i32>>)
+ // CHECK-SAME: -> !test.test_memref<[1], !llvm.array<1 x i32>>
+ %call = func.call @custom_buffer_like_callee(%arg)
+ : (tensor<1x!llvm.array<1 x i32>, #custom_buffer>)
+ -> tensor<1x!llvm.array<1 x i32>, #custom_buffer>
+ return %call : tensor<1x!llvm.array<1 x i32>, #custom_buffer>
+}
+
// -----
diff --git a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
index 00eb2dc92e14f..4360352c6ea06 100644
--- a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
+++ b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
@@ -58,7 +58,27 @@ struct TestOneShotModuleBufferizePass
opt.bufferizeFunctionBoundaries = true;
opt.functionArgTypeConverterFn =
[&](bufferization::TensorLikeType tensor, Attribute memSpace,
- func::FuncOp, const bufferization::BufferizationOptions &options) {
+ func::FuncOp funcOp,
+ const bufferization::BufferizationOptions &options) {
+ if (auto rankedTensor = dyn_cast<RankedTensorType>(tensor)) {
+ auto encoding = dyn_cast_if_present<test::TestTensorEncodingAttr>(
+ rankedTensor.getEncoding());
+ if (encoding && encoding.getDummy().getValue() == "custom_buffer")
+ return cast<bufferization::BufferLikeType>(
+ test::TestMemrefType::get(
+ rankedTensor.getContext(), rankedTensor.getShape(),
+ rankedTensor.getElementType(), memSpace));
+ }
+ if (auto tensorType = dyn_cast<TensorType>(tensor)) {
+ Type elementType = tensorType.getElementType();
+ if (!BaseMemRefType::isValidElementType(elementType)) {
+ funcOp.emitError()
+ << "cannot bufferize function boundary type " << tensorType
+ << ": element type " << elementType
+ << " is not a valid memref element type";
+ return bufferization::BufferLikeType{};
+ }
+ }
return options.unknownTypeConverterFn(tensor, memSpace, options);
};
opt.unknownTypeConverterFn =
More information about the Mlir-commits
mailing list