[Mlir-commits] [mlir] [mlir][bufferization] Refine tensor-buffer compatibility checks (PR #167705)
Andrei Golubev
llvmlistbot at llvm.org
Tue Nov 18 00:56:07 PST 2025
https://github.com/andrey-golubev updated https://github.com/llvm/llvm-project/pull/167705
>From 3744279782a97d828ef4c5501f704b8b3474086d Mon Sep 17 00:00:00 2001
From: "Golubev, Andrey" <andrey.golubev at intel.com>
Date: Wed, 12 Nov 2025 15:25:11 +0000
Subject: [PATCH 1/2] [mlir][bufferization] Refine tensor-buffer compatibility
checks
Generally, to_tensor and to_buffer already perform sufficient
verification. However, there are some unnecessary strict constraints:
* builtin tensor requires its buffer counterpart to always be memref
* to_buffer on ranked tensor requires to always return memref
These checks are assertions (i.e. preconditions), however, they actually
prevent an apparently useful bufferization where builtin tensors could
become custom buffers. Lift these assertions, maintaining the verification
procedure unchanged, to allow builtin -> custom bufferizations at
operation boundary level.
---
.../IR/BufferizableOpInterface.cpp | 12 +-------
.../Bufferization/IR/BufferizationDialect.cpp | 3 --
.../Transforms/tensorlike-bufferlike.mlir | 28 ++++++++++++++++++-
mlir/test/lib/Dialect/Test/TestTypes.cpp | 18 ++++++++----
4 files changed, 40 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index e0cf353da207f..9b11270e7bbe2 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -680,16 +680,6 @@ bool AnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
return false;
}
-// bufferization.to_buffer is not allowed to change the rank.
-static void ensureToBufferOpIsValid(Value tensor, Type memrefType) {
-#ifndef NDEBUG
- auto rankedTensorType = llvm::dyn_cast<RankedTensorType>(tensor.getType());
- assert((!rankedTensorType || llvm::cast<MemRefType>(memrefType).getRank() ==
- rankedTensorType.getRank()) &&
- "to_buffer would be invalid: mismatching ranks");
-#endif
-}
-
FailureOr<Value> bufferization::getBuffer(RewriterBase &rewriter, Value value,
const BufferizationOptions &options,
const BufferizationState &state) {
@@ -708,7 +698,7 @@ FailureOr<Value> bufferization::getBuffer(RewriterBase &rewriter, Value value,
FailureOr<BufferLikeType> bufferType = getBufferType(value, options, state);
if (failed(bufferType))
return failure();
- ensureToBufferOpIsValid(value, *bufferType);
+
return bufferization::ToBufferOp::create(rewriter, value.getLoc(),
*bufferType, value)
.getResult();
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
index d6c3cd62ee742..bd177ba1afccd 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp
@@ -54,9 +54,6 @@ struct BuiltinTensorExternalModel
mlir::LogicalResult verifyCompatibleBufferType(
mlir::Type tensor, BufferLikeType bufferType,
llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
- assert(isa<TensorType>(tensor) && "expected tensor type");
- assert(isa<BaseMemRefType>(bufferType) && "expected memref type");
-
auto tensorType = cast<ShapedType>(tensor);
auto memrefType = cast<ShapedType>(bufferType);
diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir
index d8b1a00522ab6..39474d6e15c92 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -test-tensorlike-bufferlike -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-tensorlike-bufferlike -verify-diagnostics -split-input-file | FileCheck %s
// CHECK: func.func @builtin_unranked
// CHECK-SAME: {found = {operand_0 = "is_tensor_like", result_0 = "is_buffer_like"}}
@@ -35,3 +35,29 @@ func.func @custom_memref(%t: !test.test_memref<[42], f32>) -> ()
{
return
}
+
+// -----
+
+// CHECK: func.func @builtin_custom_builtin_roundtrip
+// CHECK-SAME: {found = {operand_0 = "is_tensor_like", result_0 = "is_tensor_like"}}
+func.func @builtin_custom_builtin_roundtrip(%t: tensor<42xf32>)
+ -> tensor<42xf32> {
+ %buffer = bufferization.to_buffer %t
+ : tensor<42xf32> to !test.test_memref<[42], f32>
+ %tensor = bufferization.to_tensor %buffer
+ : !test.test_memref<[42], f32> to tensor<42xf32>
+ return %tensor : tensor<42xf32>
+}
+
+// -----
+
+// CHECK: func.func @custom_builtin_custom_roundtrip
+// CHECK-SAME: {found = {operand_0 = "is_tensor_like", result_0 = "is_tensor_like"}}
+func.func @custom_builtin_custom_roundtrip(%t: !test.test_tensor<[42], f32>)
+ -> !test.test_tensor<[42], f32> {
+ %buffer = bufferization.to_buffer %t
+ : !test.test_tensor<[42], f32> to memref<42xf32>
+ %tensor = bufferization.to_tensor %buffer
+ : memref<42xf32> to !test.test_tensor<[42], f32>
+ return %tensor : !test.test_tensor<[42], f32>
+}
diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp
index 614121f1d43dd..9cf64a896d28a 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp
@@ -569,11 +569,17 @@ TestTensorType::getBufferType(
::mlir::LogicalResult TestTensorType::verifyCompatibleBufferType(
::mlir::bufferization::BufferLikeType bufferType,
::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError) {
- auto testMemref = dyn_cast<TestMemrefType>(bufferType);
- if (!testMemref)
- return emitError() << "expected TestMemrefType";
+ if (auto testMemref = dyn_cast<TestMemrefType>(bufferType)) {
+ const bool valid = getShape() == testMemref.getShape() &&
+ getElementType() == testMemref.getElementType();
+ return mlir::success(valid);
+ }
+
+ if (auto builtinMemref = dyn_cast<MemRefType>(bufferType)) {
+ const bool valid = getShape() == builtinMemref.getShape() &&
+ getElementType() == builtinMemref.getElementType();
+ return mlir::success(valid);
+ }
- const bool valid = getShape() == testMemref.getShape() &&
- getElementType() == testMemref.getElementType();
- return mlir::success(valid);
+ return emitError() << "expected MemRefType or TestMemrefType";
}
>From f390e90c4c60f5e0650a2a04a6cca1935bc1e59f Mon Sep 17 00:00:00 2001
From: "Golubev, Andrey" <andrey.golubev at intel.com>
Date: Tue, 18 Nov 2025 08:57:56 +0000
Subject: [PATCH 2/2] Drop unused -verify-diagnostics
---
.../Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir
index 39474d6e15c92..0d516ea6ee401 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -test-tensorlike-bufferlike -verify-diagnostics -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-tensorlike-bufferlike -split-input-file | FileCheck %s
// CHECK: func.func @builtin_unranked
// CHECK-SAME: {found = {operand_0 = "is_tensor_like", result_0 = "is_buffer_like"}}
More information about the Mlir-commits
mailing list