[Mlir-commits] [mlir] [mlir][bytecode] Fix crashes when reading bytecode with unsupported types (PR #186354)

[Mehdi Amini]

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

>From 689992d5d63b64883cdc02ecc52fdb676b005d0e Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 16:39:43 -0800
Subject: [PATCH] [mlir][bytecode] Fix crashes when reading bytecode with
 unsupported types

When using test-kind=2 in the bytecode roundtrip test, integer types
(i32) are replaced by a custom type (TestI32Type) via a type callback.
This exposed two crash scenarios:

1. Reading IntegerAttr with an unsupported type: `getIntegerBitWidth`
   returns 0 for unsupported types and emits an error, but
   `readAPIntWithKnownWidth` would proceed to call
   `reader.readAPIntWithKnownWidth(0)`, creating a zero-width APInt
   with a potentially non-zero value. Fix: early-return failure when
   `bitWidth == 0`.

2. Reading VectorType with an unsupported element type: `VectorType::get`
   asserts that the element type implements VectorElementTypeInterface.
   When the element type is replaced by a custom type that doesn't
   implement this interface, the program crashes. Fix: use
   `VectorType::getChecked` with a diagnostic emitter lambda instead of
   `get<VectorType>` in the bytecode builder.

Fixes #128312
Fixes #128308
---
 .../include/mlir/IR/BuiltinDialectBytecode.td |  7 ++++++-
 mlir/lib/IR/BuiltinDialectBytecode.cpp        | 19 +++++++++++--------
 .../test_integer_attr_unsupported_type.mlir   | 15 +++++++++++++++
 .../test_vector_elem_type_interface.mlir      | 14 ++++++++++++++
 4 files changed, 46 insertions(+), 9 deletions(-)
 create mode 100644 mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir
 create mode 100644 mlir/test/Bytecode/invalid/test_vector_elem_type_interface.mlir

diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index b5ffa3eeb58c4..a28ce5614aeda 100644
--- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td
+++ b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
@@ -296,6 +296,9 @@ def VectorType : DialectType<(type
   Type:$elementType
 )> {
   let printerPredicate = "!$_val.isScalable()";
+  // Use getChecked to produce a null type (and emit a diagnostic) instead of
+  // asserting when the element type does not implement VectorElementTypeInterface.
+  let cBuilder = "VectorType::getChecked([&]() { return reader.emitError(\"invalid vector type\"); }, shape, elementType)";
 }
 
 def VectorTypeWithScalableDims : DialectType<(type
@@ -305,7 +308,9 @@ def VectorTypeWithScalableDims : DialectType<(type
 )> {
   let printerPredicate = "$_val.isScalable()";
   // Note: order of serialization does not match order of builder.
-  let cBuilder = "get<$_resultType>(context, shape, elementType, scalableDims)";
+  // Use getChecked to produce a null type (and emit a diagnostic) instead of
+  // asserting when the element type does not implement VectorElementTypeInterface.
+  let cBuilder = "VectorType::getChecked([&]() { return reader.emitError(\"invalid vector type\"); }, shape, elementType, scalableDims)";
 }
 }
 
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index fe194991a6b2a..667fb6b01841d 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -34,22 +34,25 @@ namespace {
 // TODO: Move these to separate file.
 
 // Returns the bitwidth if known, else return 0.
-static unsigned getIntegerBitWidth(DialectBytecodeReader &reader, Type type) {
-  if (auto intType = dyn_cast<IntegerType>(type)) {
+static std::optional<unsigned> getIntegerBitWidth(DialectBytecodeReader &reader,
+                                                  Type type) {
+  if (auto intType = dyn_cast<IntegerType>(type))
     return intType.getWidth();
-  }
-  if (llvm::isa<IndexType>(type)) {
+  if (llvm::isa<IndexType>(type))
     return IndexType::kInternalStorageBitWidth;
-  }
   reader.emitError()
       << "expected integer or index type for IntegerAttr, but got: " << type;
-  return 0;
+  return std::nullopt;
 }
 
 static LogicalResult readAPIntWithKnownWidth(DialectBytecodeReader &reader,
                                              Type type, FailureOr<APInt> &val) {
-  unsigned bitWidth = getIntegerBitWidth(reader, type);
-  val = reader.readAPIntWithKnownWidth(bitWidth);
+  std::optional<unsigned> bitWidth = getIntegerBitWidth(reader, type);
+  // getIntegerBitWidth returns 0 and emits an error for unsupported types.
+  // Bail out early to avoid creating a zero-width APInt with a non-zero value.
+  if (!bitWidth)
+    return failure();
+  val = reader.readAPIntWithKnownWidth(*bitWidth);
   return val;
 }
 
diff --git a/mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir b/mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir
new file mode 100644
index 0000000000000..8f7d85c0a469f
--- /dev/null
+++ b/mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir
@@ -0,0 +1,15 @@
+// RUN: not mlir-opt %s --test-bytecode-roundtrip="test-kind=2" 2>&1 | FileCheck %s
+
+// CHECK: expected integer or index type for IntegerAttr, but got: '!test.i32'
+// CHECK: failed to read bytecode
+
+// This test verifies that a proper error is emitted (rather than crashing with
+// an APInt assertion) when the type callback replaces an integer type with one
+// that does not implement IntegerType or IndexType.
+module {
+  func.func @combined_operations() -> () {
+    %cst = arith.constant dense<[[true, true, true], [true, false, true]]> : tensor<2x3xi1>
+    %reduction_output = tosa.reduce_all %cst {axis = 1 : i32} : (tensor<2x3xi1>) -> tensor<2x1xi1>
+    return
+  }
+}
diff --git a/mlir/test/Bytecode/invalid/test_vector_elem_type_interface.mlir b/mlir/test/Bytecode/invalid/test_vector_elem_type_interface.mlir
new file mode 100644
index 0000000000000..31180c71d7138
--- /dev/null
+++ b/mlir/test/Bytecode/invalid/test_vector_elem_type_interface.mlir
@@ -0,0 +1,14 @@
+// RUN: not mlir-opt %s --test-bytecode-roundtrip="test-kind=2" 2>&1 | FileCheck %s
+
+// CHECK: failed to verify 'elementType': VectorElementTypeInterface instance
+// CHECK: failed to read bytecode
+
+// This test verifies that a proper error is emitted (rather than crashing with
+// an assertion) when the type callback replaces a vector element type with one
+// that does not implement VectorElementTypeInterface.
+module {
+  func.func @main() -> () {
+    %cst = arith.constant dense<42> : vector<3xi32>
+    return
+  }
+}