[Mlir-commits] [mlir] [mlir][bytecode] Fix crashes when reading bytecode with unsupported types (PR #186354)
Mehdi Amini
llvmlistbot at llvm.org
Mon Mar 16 06:08:18 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/186354
>From 22cbb49d97da69b7c656b30e75e6f90fc63a0a73 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 1/3] [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 53a859e32d64b..64cc8a8ff5e20 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 f7430784dd222..b3124d1d9a4bd 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
+ }
+}
>From 40d95318ec55f8df5a04a175fb4f87039fa45ba2 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Mon, 16 Mar 2026 04:30:11 -0700
Subject: [PATCH 2/3] Address review comments: fix stale comments and improve
tests
- Fix stale comments in getIntegerBitWidth/readAPIntWithKnownWidth
(the function returns std::nullopt, not 0)
- Merge the two separate test files into a single split-input-file
test (mlir/test/Bytecode/invalid/invalid-type-remapping.mlir)
- Simplify the IntegerAttr test: replace the TOSA ops with a plain
arith.constant to avoid the TOSA dialect dependency
- Add a scalable VectorType test case to cover the
VectorTypeWithScalableDims code path that was also fixed
Assisted-by: Claude Code
---
mlir/lib/IR/BuiltinDialectBytecode.cpp | 7 ++--
.../invalid/invalid-type-remapping.mlir | 42 +++++++++++++++++++
.../test_integer_attr_unsupported_type.mlir | 15 -------
.../test_vector_elem_type_interface.mlir | 14 -------
4 files changed, 46 insertions(+), 32 deletions(-)
create mode 100644 mlir/test/Bytecode/invalid/invalid-type-remapping.mlir
delete mode 100644 mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir
delete mode 100644 mlir/test/Bytecode/invalid/test_vector_elem_type_interface.mlir
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index b3124d1d9a4bd..14dc665184099 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -33,7 +33,7 @@ namespace {
// TODO: Move these to separate file.
-// Returns the bitwidth if known, else return 0.
+// Returns the bitwidth if known, else return std::nullopt.
static std::optional<unsigned> getIntegerBitWidth(DialectBytecodeReader &reader,
Type type) {
if (auto intType = dyn_cast<IntegerType>(type))
@@ -48,8 +48,9 @@ static std::optional<unsigned> getIntegerBitWidth(DialectBytecodeReader &reader,
static LogicalResult readAPIntWithKnownWidth(DialectBytecodeReader &reader,
Type type, FailureOr<APInt> &val) {
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.
+ // getIntegerBitWidth returns std::nullopt 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);
diff --git a/mlir/test/Bytecode/invalid/invalid-type-remapping.mlir b/mlir/test/Bytecode/invalid/invalid-type-remapping.mlir
new file mode 100644
index 0000000000000..09dbe409d8237
--- /dev/null
+++ b/mlir/test/Bytecode/invalid/invalid-type-remapping.mlir
@@ -0,0 +1,42 @@
+// RUN: not mlir-opt %s -split-input-file --test-bytecode-roundtrip="test-kind=2" 2>&1 | FileCheck %s
+
+// Tests that proper errors are emitted (rather than crashes) when the type
+// callback replaces types with ones that are incompatible with built-in types
+// and attributes (test-kind=2 replaces i32 with !test.i32).
+
+// CHECK: expected integer or index type for IntegerAttr, but got: '!test.i32'
+// CHECK: failed to read bytecode
+// IntegerAttr whose type is replaced by one that is neither IntegerType nor
+// IndexType — previously crashed with an APInt assertion.
+module {
+ func.func @integer_attr_unsupported_type() {
+ %c = arith.constant 1 : i32
+ return
+ }
+}
+
+// -----
+
+// CHECK: failed to verify 'elementType': VectorElementTypeInterface instance
+// CHECK: failed to read bytecode
+// Fixed-size VectorType whose element type is replaced by one that does not
+// implement VectorElementTypeInterface — previously crashed in VectorType::get.
+module {
+ func.func @vector_unsupported_elem_type() {
+ %cst = arith.constant dense<42> : vector<3xi32>
+ return
+ }
+}
+
+// -----
+
+// CHECK: failed to verify 'elementType': VectorElementTypeInterface instance
+// CHECK: failed to read bytecode
+// Scalable VectorType whose element type is replaced by one that does not
+// implement VectorElementTypeInterface — exercises the VectorTypeWithScalableDims
+// bytecode path.
+module {
+ func.func @scalable_vector_unsupported_elem_type(%v : vector<[3]xi32>) {
+ return
+ }
+}
diff --git a/mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir b/mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir
deleted file mode 100644
index 8f7d85c0a469f..0000000000000
--- a/mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir
+++ /dev/null
@@ -1,15 +0,0 @@
-// 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
deleted file mode 100644
index 31180c71d7138..0000000000000
--- a/mlir/test/Bytecode/invalid/test_vector_elem_type_interface.mlir
+++ /dev/null
@@ -1,14 +0,0 @@
-// 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
- }
-}
>From af17ee4d8fc78de35b5fa32a169369373d84e126 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Mon, 16 Mar 2026 06:05:58 -0700
Subject: [PATCH 3/3] Address review comment: fold dense-elem test into
invalid-type-remapping
Merge mlir/test/Bytecode/invalid/invalid-dense-elem-type-interface.mlir
into the consolidated invalid-type-remapping.mlir file, as both exercise
the same --test-bytecode-roundtrip="test-kind=2" type-remapping path.
Assisted-by: Claude Code
---
.../invalid-dense-elem-type-interface.mlir | 15 ---------------
.../Bytecode/invalid/invalid-type-remapping.mlir | 13 +++++++++++++
2 files changed, 13 insertions(+), 15 deletions(-)
delete mode 100644 mlir/test/Bytecode/invalid/invalid-dense-elem-type-interface.mlir
diff --git a/mlir/test/Bytecode/invalid/invalid-dense-elem-type-interface.mlir b/mlir/test/Bytecode/invalid/invalid-dense-elem-type-interface.mlir
deleted file mode 100644
index f076dcb9b2f1f..0000000000000
--- a/mlir/test/Bytecode/invalid/invalid-dense-elem-type-interface.mlir
+++ /dev/null
@@ -1,15 +0,0 @@
-// 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 DenseTypedElementsAttr.
-
-// CHECK: DenseTypedElementsAttr 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>
- }
-}
diff --git a/mlir/test/Bytecode/invalid/invalid-type-remapping.mlir b/mlir/test/Bytecode/invalid/invalid-type-remapping.mlir
index 09dbe409d8237..44d0a4eb8bb4a 100644
--- a/mlir/test/Bytecode/invalid/invalid-type-remapping.mlir
+++ b/mlir/test/Bytecode/invalid/invalid-type-remapping.mlir
@@ -40,3 +40,16 @@ module {
return
}
}
+
+// -----
+
+// CHECK: DenseTypedElementsAttr element type must implement DenseElementTypeInterface, but got: '!test.i32'
+// CHECK: failed to read bytecode
+// DenseTypedElementsAttr whose element type is replaced by one that does not
+// implement DenseElementTypeInterface — previously crashed with an assertion.
+module {
+ func.func @dense_elem_unsupported_type() -> tensor<10xi32> {
+ %0 = arith.constant dense<42> : tensor<10xi32>
+ return %0 : tensor<10xi32>
+ }
+}
More information about the Mlir-commits
mailing list