[mlir][bytecode] Use getChecked<T>() in bytecode reading to avoid crashes

mlir/include/mlir/Bytecode/BytecodeImplementation.h

@@ -527,6 +527,36 @@ auto get(MLIRContext *context, Ts &&...params) {
   }
 }
 
+namespace detail {
+template <typename T, typename... Ts>
+using has_get_checked_method = decltype(T::getChecked(std::declval<Ts>()...));
+} // namespace detail
+
+/// Helper method analogous to `get`, but uses `getChecked` when available to
+/// allow graceful failure on invalid parameters instead of asserting.
+///
+/// Only the no-context form of `getChecked` is tried here. Types that expose
+/// `getChecked(emitError, params...)` without a leading `MLIRContext*` (e.g.
+/// MemRefType, VectorType, RankedTensorType) will use it for graceful failure.
+/// Everything else falls back to `get<T>()`.  We intentionally do NOT try
+/// `T::getChecked(emitError, context, params...)`: for types that only inherit
+/// the base `StorageUserBase::getChecked` template (e.g. ArrayAttr), that
+/// template instantiation requires a complete storage type which may not be
+/// available in the bytecode reading TU.
+template <typename T, typename... Ts>
+auto getChecked(function_ref<InFlightDiagnostic()> emitError,
+                MLIRContext *context, Ts &&...params) {
+  if constexpr (llvm::is_detected<detail::has_get_checked_method, T,
+                                  function_ref<InFlightDiagnostic()>,
+                                  Ts...>::value) {
+    (void)context;
+    return T::getChecked(emitError, std::forward<Ts>(params)...);
+  } else {
+    // Fall back to get() for types that don't define a no-context getChecked.
+    return get<T>(context, std::forward<Ts>(params)...);
+  }
+}
+
 } // namespace mlir
 
 #endif // MLIR_BYTECODE_BYTECODEIMPLEMENTATION_H

mlir/include/mlir/IR/BuiltinDialectBytecode.td

@@ -79,15 +79,15 @@ def IntegerAttr: DialectAttribute<(attr
   Type:$type,
   KnownWidthAPInt<"type">:$value
 )> {
-  let cBuilder = "get<$_resultType>(context, type, *value)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, type, *value)";
 }
 
 defvar FloatType = Type;
 def FloatAttr : DialectAttribute<(attr
   FloatType:$type,
   KnownSemanticsAPFloat<"type">:$value
 )> {
-  let cBuilder = "get<$_resultType>(context, type, *value)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, type, *value)";
 }
 
 def CallSiteLoc : DialectAttribute<(attr
@@ -117,7 +117,7 @@ def FileLineColLoc : DialectAttribute<(attr
 }
 
 let cType = "FusedLoc",
-    cBuilder = "cast<FusedLoc>(get<FusedLoc>(context, $_args))" in {
+    cBuilder = "cast<FusedLoc>(getChecked<FusedLoc>([&]() { return reader.emitError(); }, context, $_args))" in {
 def FusedLoc : DialectAttribute<(attr
   Array<Location>:$locations
 )> {
@@ -144,7 +144,7 @@ def DenseResourceElementsAttr : DialectAttribute<(attr
   ResourceHandle<"DenseResourceElementsHandle">:$rawHandle
 )> {
   // Note: order of serialization does not match order of builder.
-  let cBuilder = "get<$_resultType>(context, type, *rawHandle)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, type, *rawHandle)";
 }
 
 let cType = "RankedTensorType" in {
@@ -162,7 +162,7 @@ def RankedTensorTypeWithEncoding : DialectType<(type
 )> {
   let printerPredicate = "$_val.getEncoding()";
   // Note: order of serialization does not match order of builder.
-  let cBuilder = "get<$_resultType>(context, shape, elementType, encoding)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, shape, elementType, encoding)";
 }
 }
 
@@ -258,7 +258,7 @@ def MemRefTypeWithMemSpace : DialectType<(type
 )> {
   let printerPredicate = "!!$_val.getMemorySpace()";
   // Note: order of serialization does not match order of builder.
-  let cBuilder = "get<$_resultType>(context, shape, elementType, layout, memorySpace)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, shape, elementType, layout, memorySpace)";
 }
 }
 
@@ -273,7 +273,7 @@ def UnrankedMemRefType : DialectType<(type
   Type:$elementType
 )> {
   let printerPredicate = "!$_val.getMemorySpace()";
-  let cBuilder = "get<$_resultType>(context, elementType, Attribute())";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, elementType, Attribute())";
 }
 
 def UnrankedMemRefTypeWithMemSpace : DialectType<(type
@@ -282,7 +282,7 @@ def UnrankedMemRefTypeWithMemSpace : DialectType<(type
 )> {
   let printerPredicate = "$_val.getMemorySpace()";
   // Note: order of serialization does not match order of builder.
-  let cBuilder = "get<$_resultType>(context, elementType, memorySpace)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, elementType, memorySpace)";
 }
 }
 
@@ -308,9 +308,7 @@ def VectorTypeWithScalableDims : DialectType<(type
 )> {
   let printerPredicate = "$_val.isScalable()";
   // Note: order of serialization does not match order of builder.
-  // 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)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, shape, elementType, scalableDims)";
 }
 }
 

mlir/include/mlir/IR/BytecodeBase.td

@@ -147,11 +147,11 @@ class DialectAttrOrType<dag d> {
 
 class DialectAttribute<dag d> : DialectAttrOrType<d>, AttributeKind {
   let cParser = "succeeded($_reader.readAttribute<$_resultType>($_var))";
-  let cBuilder = "get<$_resultType>(context, $_args)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, $_args)";
 }
 class DialectType<dag d> : DialectAttrOrType<d>, TypeKind {
   let cParser = "succeeded($_reader.readType<$_resultType>($_var))";
-  let cBuilder = "get<$_resultType>(context, $_args)";
+  let cBuilder = "getChecked<$_resultType>([&]() { return reader.emitError(); }, context, $_args)";
 }
 
 class DialectAttributes<string d> {

mlir/test/Bytecode/bytecode_callback_with_custom_type.mlir

@@ -16,3 +16,17 @@ func.func @base_test(%arg0: i32, %arg1: f32) {
 
 // TEST_2: Overriding parsing of TestI32Type encoding...
 // TEST_2: func.func @base_test([[ARG0:%.+]]: !test.i32, [[ARG1:%.+]]: f32) {
+
+// -----
+
+// Regression test: complex types such as memref must round-trip without
+// crashing when the test-kind=2 type callback calls iface->readType() for
+// every builtin type. Previously this crashed because the bytecode reading
+// path used get<T>() (which asserts) instead of getChecked<T>() (which
+// returns null on invalid input).
+
+func.func @test_memref_types(%arg0: memref<4xf32>, %arg1: memref<4x4xf32>) {
+  return
+}
+
+// TEST_2: func.func @test_memref_types([[A0:%.+]]: memref<4xf32>, [[A1:%.+]]: memref<4x4xf32>) {

mlir/tools/mlir-tblgen/BytecodeDialectGen.cpp

@@ -206,7 +206,10 @@ void Generator::emitParseHelper(StringRef kind, StringRef returnType,
   auto funScope = ios.scope("{\n", "}");
 
   if (args.empty()) {
-    ios << formatv("return get<{0}>(context);\n", returnType);
+    ios << formatv(
+        "return getChecked<{0}>([&]() {{ return reader.emitError(); }, "
+        "context);\n",
+        returnType);
     return;
   }