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

[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 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

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


4 Files Affected:

- (modified) mlir/include/mlir/IR/BuiltinDialectBytecode.td (+6-1) 
- (modified) mlir/lib/IR/BuiltinDialectBytecode.cpp (+4) 
- (added) mlir/test/Bytecode/invalid/test_integer_attr_unsupported_type.mlir (+15) 
- (added) mlir/test/Bytecode/invalid/test_vector_elem_type_interface.mlir (+14) 


``````````diff
diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index b5ffa3eeb58c4..3644c27bcb8cd 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(); }, 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(); }, shape, elementType, scalableDims)";
 }
 }
 
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index 73ec3eec69615..50dc44070d905 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -49,6 +49,10 @@ static unsigned getIntegerBitWidth(DialectBytecodeReader &reader, Type type) {
 static LogicalResult readAPIntWithKnownWidth(DialectBytecodeReader &reader,
                                              Type type, FailureOr<APInt> &val) {
   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 == 0)
+    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
+  }
+}

``````````

</details>


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