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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tensor

@llvm/pr-subscribers-mlir

Author: Stefan Weigl-Bosker (sweiglbosker)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/173692.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+10) 
- (added) mlir/test/Dialect/Tensor/one-shot-bufferize-invalid.mlir (+8) 


``````````diff
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>>
+}
+

``````````

</details>


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