[Mlir-commits] [mlir] [mlir][bytecode] Fix crash when reading DenseIntOrFPElementsAttr with unsupported element type (PR #184773)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When a bytecode type callback substitutes a type that does not implement DenseElementTypeInterface (e.g., \!test.i32 replacing i32), the bytecode reader attempted to reconstruct a DenseIntOrFPElementsAttr with that type. This unconditionally called getDenseElementBitWidth() which hit an llvm_unreachable on unsupported types.

Fix this by validating the element type implements DenseElementTypeInterface in readDenseIntOrFPElementsAttr before proceeding. If the check fails, a proper diagnostic is emitted and reading fails gracefully instead of crashing.

Fixes #<!-- -->128317

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


2 Files Affected:

- (modified) mlir/lib/IR/BuiltinDialectBytecode.cpp (+11) 
- (added) mlir/test/Bytecode/invalid/dense_elem_type_interface.mlir (+15) 


``````````diff
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index 8cf25391ba11a..73ec3eec69615 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -151,6 +151,17 @@ static void writeFileLineColRangeLocs(DialectBytecodeWriter &writer,
 static LogicalResult
 readDenseIntOrFPElementsAttr(DialectBytecodeReader &reader, ShapedType type,
                              SmallVectorImpl<char> &rawData) {
+  // Validate that the element type implements DenseElementTypeInterface.
+  // Without this check, downstream code unconditionally calls
+  // getDenseElementBitWidth() which asserts on unsupported types.
+  if (!llvm::isa<DenseElementType>(type.getElementType())) {
+    reader.emitError()
+        << "DenseIntOrFPElementsAttr element type must implement "
+           "DenseElementTypeInterface, but got: "
+        << type.getElementType();
+    return failure();
+  }
+
   ArrayRef<char> blob;
   if (failed(reader.readBlob(blob)))
     return failure();
diff --git a/mlir/test/Bytecode/invalid/dense_elem_type_interface.mlir b/mlir/test/Bytecode/invalid/dense_elem_type_interface.mlir
new file mode 100644
index 0000000000000..2a03dd8ef4b23
--- /dev/null
+++ b/mlir/test/Bytecode/invalid/dense_elem_type_interface.mlir
@@ -0,0 +1,15 @@
+// RUN: not mlir-opt %s --test-bytecode-roundtrip="test-kind=2" 2>&1 | FileCheck %s
+
+// Regression test: test-kind=2 replaces i32 with !test.i32 (a type that does
+// not implement DenseElementTypeInterface). This should produce a proper error
+// instead of an assertion failure when deserializing DenseIntOrFPElementsAttr.
+
+// CHECK: DenseIntOrFPElementsAttr element type must implement DenseElementTypeInterface, but got: '!test.i32'
+// CHECK: failed to read bytecode
+
+module {
+  func.func @test() -> tensor<10xi32> {
+    %0 = arith.constant dense<42> : tensor<10xi32>
+    return %0 : tensor<10xi32>
+  }
+}

``````````

</details>


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