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

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/184773

>From 16434111d891540353911308d9cf611b2eb0b654 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 15:19:10 -0800
Subject: [PATCH] [mlir][bytecode] Fix crash when reading
 DenseIntOrFPElementsAttr with unsupported element type

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
---
 mlir/lib/IR/BuiltinDialectBytecode.cpp            | 11 +++++++++++
 .../invalid-dense-elem-type-interface.mlir        | 15 +++++++++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 mlir/test/Bytecode/invalid/invalid-dense-elem-type-interface.mlir

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/invalid-dense-elem-type-interface.mlir b/mlir/test/Bytecode/invalid/invalid-dense-elem-type-interface.mlir
new file mode 100644
index 0000000000000..2a03dd8ef4b23
--- /dev/null
+++ b/mlir/test/Bytecode/invalid/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>
+  }
+}