[mlir][bufferization] Refine tensor-buffer compatibility checks #167705
+110
−20
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Member
|
@llvm/pr-subscribers-mlir Author: Andrei Golubev (andrey-golubev) ChangesGenerally, to_tensor and to_buffer already perform sufficient verification. However, there are some unnecessary strict constraints:
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. Full diff: https://github.com/llvm/llvm-project/pull/167705.diff 4 Files Affected:
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..5117d7563a915 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 TestMemrefType";
}
|
Member
|
@llvm/pr-subscribers-mlir-bufferization Author: Andrei Golubev (andrey-golubev) ChangesGenerally, to_tensor and to_buffer already perform sufficient verification. However, there are some unnecessary strict constraints:
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. Full diff: https://github.com/llvm/llvm-project/pull/167705.diff 4 Files Affected:
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..5117d7563a915 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 TestMemrefType";
}
|
andrey-golubev
commented
Nov 12, 2025
| rankedTensorType.getRank()) && | ||
| "to_buffer would be invalid: mismatching ranks"); | ||
| #endif | ||
| } |
Contributor
Author
There was a problem hiding this comment.
note: this actually blocked our downstream to switch to mlir::bufferization::getBuffer() implementation. from what I can tell, this check is superseded a long time ago by ToBufferOp's verifier (that checks, for builtins, both shape - and thus rank - and element type).
unfortunately, testing this specifically is a pain: i need a fairly large setup to showcase how this fails. I hope this is fine that I just delete it and provide some general tests that exercise op's verifiers instead.
andrey-golubev
force-pushed
the
refine_buffer_verification
branch
from
c1067fd to
3744279
Compare
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.
Contributor
Author
|
Gentle ping. |
| @@ -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 | |||
Member
There was a problem hiding this comment.
You added -verify-diagnostics, but I don't see a test case utilizing it.
Contributor
Author
There was a problem hiding this comment.
perhaps this is my lack of knowledge: i wanted to make sure verifiers are run and so added -verify-diagnostics. I think I've seen in the past (at least in our downstream) that verifiers are not guaranteed if e.g. IR parsing uses ::get() instead of ::getChecked().
Contributor
Author
There was a problem hiding this comment.
removed the -verify-diagnostics part. thanks for noticing!
Member
There was a problem hiding this comment.
-verify-diagnostics just tries match up errors, warning, remarks etc. with // expected-error etc. in the input file.
| const bool valid = getShape() == testMemref.getShape() && | ||
| getElementType() == testMemref.getElementType(); | ||
| return mlir::success(valid); | ||
| return emitError() << "expected MemRefType or TestMemrefType"; |
Member
There was a problem hiding this comment.
Did you mean to trigger this in the test above?
Contributor
Author
There was a problem hiding this comment.
this specifically - not really.
the bulk of the change is for allowing a mix of builtin memref and test memref (this is what the test above exercises).
the emit-error message is changed just to better align with the new logic (so that the error, if it happens, is less ambiguous). i don't really intend to test code that's written for the sake of testing :)
🐧 Linux x64 Test Results
|
andrey-golubev
commented
Nov 18, 2025
| @@ -127,3 +127,63 @@ func.func @invalid_manual_deallocation() { | |||
| // expected-error @below{{op attribute 'bufferization.manual_deallocation' can be used only on ops that have an allocation and/or free side effect}} | |||
| arith.constant {bufferization.manual_deallocation} 0 : index | |||
| } | |||
|
|
|||
| // ----- | |||
Contributor
Author
There was a problem hiding this comment.
@matthias-springer I have extended these tests with rank / shape / element type checks for builtin types. now, this does "confirm" that ensureToBufferOpIsValid() removed in ::getBuffer() is indeed unnecessary.
andrey-golubev
commented
Nov 18, 2025
| @@ -83,3 +83,40 @@ func.func @test_dealloc_op(%arg0: memref<2xf32>, %arg1: memref<4xi32>, | |||
| bufferization.dealloc | |||
| return %0#0, %0#1 : i1, i1 | |||
| } | |||
|
|
|||
Contributor
Author
There was a problem hiding this comment.
note: moved the tests that were originally in tensorlike-bufferlike.mlir to ops.mlir. this seems to be a better place since the tests essentially check what operations can do.
matthias-springer
approved these changes
Nov 18, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Generally, to_tensor and to_buffer already perform sufficient verification. However, there are some unnecessary strict constraints:
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.