From 3744279782a97d828ef4c5501f704b8b3474086d Mon Sep 17 00:00:00 2001 From: "Golubev, Andrey" Date: Wed, 12 Nov 2025 15:25:11 +0000 Subject: [PATCH 1/5] [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(tensor.getType()); - assert((!rankedTensorType || llvm::cast(memrefType).getRank() == - rankedTensorType.getRank()) && - "to_buffer would be invalid: mismatching ranks"); -#endif -} - FailureOr bufferization::getBuffer(RewriterBase &rewriter, Value value, const BufferizationOptions &options, const BufferizationState &state) { @@ -708,7 +698,7 @@ FailureOr bufferization::getBuffer(RewriterBase &rewriter, Value value, FailureOr 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 emitError) const { - assert(isa(tensor) && "expected tensor type"); - assert(isa(bufferType) && "expected memref type"); - auto tensorType = cast(tensor); auto memrefType = cast(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(bufferType); - if (!testMemref) - return emitError() << "expected TestMemrefType"; + if (auto testMemref = dyn_cast(bufferType)) { + const bool valid = getShape() == testMemref.getShape() && + getElementType() == testMemref.getElementType(); + return mlir::success(valid); + } + + if (auto builtinMemref = dyn_cast(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" Date: Tue, 18 Nov 2025 08:57:56 +0000 Subject: [PATCH 2/5] 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"}} From ae43accb13bef673faa465c3441c50499a1e4e0c Mon Sep 17 00:00:00 2001 From: "Golubev, Andrey" Date: Tue, 18 Nov 2025 09:25:23 +0000 Subject: [PATCH 3/5] Test to_buffer / to_tensor mismatches --- mlir/test/Dialect/Bufferization/invalid.mlir | 60 ++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/mlir/test/Dialect/Bufferization/invalid.mlir b/mlir/test/Dialect/Bufferization/invalid.mlir index 2c8807b66de74..4a839608726a0 100644 --- a/mlir/test/Dialect/Bufferization/invalid.mlir +++ b/mlir/test/Dialect/Bufferization/invalid.mlir @@ -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 } + +// ----- + +func.func @invalid_rank_to_buffer(%t: tensor<1x2x3x4xf32>) { + // expected-error @below{{shapes do not match}} + // expected-error @below{{'bufferization.to_buffer' op failed to verify that specified tensor and buffer types match}} + %b = bufferization.to_buffer %t + : tensor<1x2x3x4xf32> to memref<1x2x3xf32> + return +} + +// ----- + +func.func @invalid_rank_to_tensor(%b: memref<1x2x3xf32>) { + // expected-error @below{{shapes do not match}} + // expected-error @below{{'bufferization.to_tensor' op failed to verify that specified tensor and buffer types match}} + %t = bufferization.to_tensor %b + : memref<1x2x3xf32> to tensor<1x2x3x4xf32> + return +} + +// ----- + +func.func @invalid_shape_to_buffer(%t: tensor<1x2x3x4xf32>) { + // expected-error @below{{'bufferization.to_buffer' op failed to verify that specified tensor and buffer types match}} + // expected-error @below{{shapes do not match}} + %b = bufferization.to_buffer %t + : tensor<1x2x3x4xf32> to memref<1x2x4x3xf32> + return +} + +// ----- + +func.func @invalid_shape_to_tensor(%b: memref<1x2x4x3xf32>) { + // expected-error @below{{'bufferization.to_tensor' op failed to verify that specified tensor and buffer types match}} + // expected-error @below{{shapes do not match}} + %t = bufferization.to_tensor %b + : memref<1x2x4x3xf32> to tensor<1x2x3x4xf32> + return +} + +// ----- + +func.func @invalid_type_to_buffer(%t: tensor<1x2x3x4xf32>) { + // expected-error @below{{'bufferization.to_buffer' op failed to verify that specified tensor and buffer types match}} + // expected-error @below{{element types do not match}} + %b = bufferization.to_buffer %t + : tensor<1x2x3x4xf32> to memref<1x2x3x4xf16> + return +} + +// ----- + +func.func @invalid_type_to_tensor(%b: memref<1x2x3x4xf16>) { + // expected-error @below{{'bufferization.to_tensor' op failed to verify that specified tensor and buffer types match}} + // expected-error @below{{element types do not match}} + %t2 = bufferization.to_tensor %b + : memref<1x2x3x4xf16> to tensor<1x2x3x4xf32> + return +} From 1eceefeb0fd5f833be5b2d6862496ea4aa430191 Mon Sep 17 00:00:00 2001 From: "Golubev, Andrey" Date: Tue, 18 Nov 2025 09:28:43 +0000 Subject: [PATCH 4/5] Reorder expected-error to align with other tests --- mlir/test/Dialect/Bufferization/invalid.mlir | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/test/Dialect/Bufferization/invalid.mlir b/mlir/test/Dialect/Bufferization/invalid.mlir index 4a839608726a0..9884b040119d0 100644 --- a/mlir/test/Dialect/Bufferization/invalid.mlir +++ b/mlir/test/Dialect/Bufferization/invalid.mlir @@ -131,8 +131,8 @@ func.func @invalid_manual_deallocation() { // ----- func.func @invalid_rank_to_buffer(%t: tensor<1x2x3x4xf32>) { - // expected-error @below{{shapes do not match}} // expected-error @below{{'bufferization.to_buffer' op failed to verify that specified tensor and buffer types match}} + // expected-error @below{{shapes do not match}} %b = bufferization.to_buffer %t : tensor<1x2x3x4xf32> to memref<1x2x3xf32> return @@ -141,8 +141,8 @@ func.func @invalid_rank_to_buffer(%t: tensor<1x2x3x4xf32>) { // ----- func.func @invalid_rank_to_tensor(%b: memref<1x2x3xf32>) { - // expected-error @below{{shapes do not match}} // expected-error @below{{'bufferization.to_tensor' op failed to verify that specified tensor and buffer types match}} + // expected-error @below{{shapes do not match}} %t = bufferization.to_tensor %b : memref<1x2x3xf32> to tensor<1x2x3x4xf32> return From 07e040fa6a8d0421a72bc4cedaa0e04347878fe6 Mon Sep 17 00:00:00 2001 From: "Golubev, Andrey" Date: Tue, 18 Nov 2025 09:57:09 +0000 Subject: [PATCH 5/5] Move new tests for cross-type compatibility to ops.mlir --- .../Transforms/tensorlike-bufferlike.mlir | 26 ------------- mlir/test/Dialect/Bufferization/ops.mlir | 37 +++++++++++++++++++ 2 files changed, 37 insertions(+), 26 deletions(-) diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir index 0d516ea6ee401..d8b1a00522ab6 100644 --- a/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir @@ -35,29 +35,3 @@ 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/Dialect/Bufferization/ops.mlir b/mlir/test/Dialect/Bufferization/ops.mlir index fc6df4a09f706..b0db1bb2d0389 100644 --- a/mlir/test/Dialect/Bufferization/ops.mlir +++ b/mlir/test/Dialect/Bufferization/ops.mlir @@ -83,3 +83,40 @@ func.func @test_dealloc_op(%arg0: memref<2xf32>, %arg1: memref<4xi32>, bufferization.dealloc return %0#0, %0#1 : i1, i1 } + +// CHECK: func.func @test_builtin_custom_builtin_type_conversion +// CHECK-SAME: (%[[t:.*]]: tensor<42xf32>) -> tensor<42xf32> +func.func @test_builtin_custom_builtin_type_conversion(%t: tensor<42xf32>) + -> tensor<42xf32> { + // CHECK: %[[buffer:.*]] = bufferization.to_buffer %[[t]] + // CHECK-SAME: to !test.test_memref<[42], f32> + %buffer = bufferization.to_buffer %t + : tensor<42xf32> to !test.test_memref<[42], f32> + + // CHECK: %[[tensor:.*]] = bufferization.to_tensor %[[buffer]] + // CHECK-SAME: to tensor<42xf32> + %tensor = bufferization.to_tensor %buffer + : !test.test_memref<[42], f32> to tensor<42xf32> + + // CHECK: return %[[tensor]] + return %tensor : tensor<42xf32> +} + +// CHECK: func.func @test_custom_builtin_custom_type_conversion +// CHECK-SAME: (%[[t:.*]]: !test.test_tensor<[42], f32>) +// CHECK-SAME: -> !test.test_tensor<[42], f32> +func.func @test_custom_builtin_custom_type_conversion(%t: !test.test_tensor<[42], f32>) + -> !test.test_tensor<[42], f32> { + // CHECK: %[[buffer:.*]] = bufferization.to_buffer %[[t]] + // CHECK-SAME: to memref<42xf32> + %buffer = bufferization.to_buffer %t + : !test.test_tensor<[42], f32> to memref<42xf32> + + // CHECK: %[[tensor:.*]] = bufferization.to_tensor %[[buffer]] + // CHECK-SAME: to !test.test_tensor<[42], f32> + %tensor = bufferization.to_tensor %buffer + : memref<42xf32> to !test.test_tensor<[42], f32> + + // CHECK: return %[[tensor]] + return %tensor : !test.test_tensor<[42], f32> +}