[Mlir-commits] [mlir] [Bufferization]: Handle invalid memref element types (PR #173692)

[Stefan Weigl-Bosker]

https://github.com/sweiglbosker created https://github.com/llvm/llvm-project/pull/173692

Fixes https://github.com/llvm/llvm-project/issues/128329, https://github.com/llvm/llvm-project/issues/128330, https://github.com/llvm/llvm-project/issues/173565, https://github.com/llvm/llvm-project/issues/114730

There is an assertion failure in `-one-shot-bufferize` when tensors that have an element type that can't be a memref element type are encountered. https://github.com/llvm/llvm-project/blob/f8d3f47e1fd09392aa30df83849b25acd8c59a25/mlir/include/mlir/IR/BuiltinTypes.h#L440

We can't emit a to_tensor for ops that do implement `BufferizableOpInterface`, and i don't think quantizing is the right move either, so erroring seemed like the best fit.

After some trial and error, `defaultGetBufferType` seems like the most functional and least invasive place to put this check.

>From f9ff5daf2f5bcf964dff8d7714ceb6cd555af52e Mon Sep 17 00:00:00 2001
From: Stefan Weigl-Bosker <stefan at s00.xyz>
Date: Fri, 26 Dec 2025 20:06:30 -0500
Subject: [PATCH] [Bufferization]: Handle invalid memref element types

---
 .../Bufferization/IR/BufferizableOpInterface.cpp       | 10 ++++++++++
 .../Dialect/Tensor/one-shot-bufferize-invalid.mlir     |  8 ++++++++
 2 files changed, 18 insertions(+)
 create mode 100644 mlir/test/Dialect/Tensor/one-shot-bufferize-invalid.mlir

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 9b11270e7bbe2..9746f012261a1 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -961,6 +961,16 @@ FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
   assert(llvm::isa<TensorType>(value.getType()) && "expected tensor type");
   auto tensorType = cast<TensorType>(value.getType());
 
+  auto elementType = tensorType.getElementType();
+
+  if (!BaseMemRefType::isValidElementType(elementType)) {
+    if (Operation *def = value.getDefiningOp())
+      return def->emitError() << "cannot bufferize value of type " << tensorType
+                              << ": element type " << elementType
+                              << " is not a valid memref element type";
+    return failure();
+  }
+
   // No further analysis is possible for a block argument.
   if (llvm::isa<BlockArgument>(value)) {
     return cast<BufferLikeType>(
diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize-invalid.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize-invalid.mlir
new file mode 100644
index 0000000000000..c84ce03054ed6
--- /dev/null
+++ b/mlir/test/Dialect/Tensor/one-shot-bufferize-invalid.mlir
@@ -0,0 +1,8 @@
+// RUN: not mlir-opt %s -one-shot-bufferize -verify-diagnostics
+
+// expected-error cannot bufferize value of type 'tensor<1x!quant.uniform<i8:f32, 1.000000e+00>>': element type '!quant.uniform<i8:f32, 1.000000e+00>' is not a valid memref element type
+func.func @internal_value() -> tensor<1x!quant.uniform<i8:f32, 1.0>> {
+  %t = tensor.empty() : tensor<1x!quant.uniform<i8:f32, 1.0>>
+  return %t : tensor<1x!quant.uniform<i8:f32, 1.0>>
+}
+