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

[Stefan Weigl-Bosker]

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

>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 1/2] [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>>
+}
+

>From 69103cb8ef809ac1197a294511e0e46bde51490e Mon Sep 17 00:00:00 2001
From: Stefan Weigl-Bosker <stefan at s00.xyz>
Date: Sat, 27 Dec 2025 11:31:35 -0500
Subject: [PATCH 2/2] fixes

---
 .../Bufferization/IR/BufferizableOpInterface.cpp     | 12 +++++-------
 ...-shot-bufferize-invalid-memref-element-type.mlir} |  0
 2 files changed, 5 insertions(+), 7 deletions(-)
 rename mlir/test/Dialect/{Tensor/one-shot-bufferize-invalid.mlir => Bufferization/Transforms/one-shot-bufferize-invalid-memref-element-type.mlir} (100%)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 9746f012261a1..4fa580a8c78b5 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -963,13 +963,11 @@ FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
 
   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();
-  }
+  if (!BaseMemRefType::isValidElementType(elementType))
+    return getOwnerOfValue(value)->emitError()
+           << "cannot bufferize value of type " << tensorType
+           << ": element type " << elementType
+           << " is not a valid memref element type";
 
   // No further analysis is possible for a block argument.
   if (llvm::isa<BlockArgument>(value)) {
diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-invalid-memref-element-type.mlir
similarity index 100%
rename from mlir/test/Dialect/Tensor/one-shot-bufferize-invalid.mlir
rename to mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-invalid-memref-element-type.mlir