From a58177135c1e7ecad6ee1c1234ba13fc1a38b38f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sat, 29 Aug 2026 16:24:23 +0200 Subject: [PATCH] Vulkan: serialize non-finite floats using flatc's spelling The Vulkan graph is serialized by dumping it to JSON with Python's json module and handing that to flatc. The two disagree on how to spell the non-finite floats: Python emits Infinity / -Infinity / NaN, and flatc accepts none of them. Any model whose graph carries a non-finite scalar therefore fails to lower, with the -inf fill value of a transformer attention mask being the common source. all-MiniLM-L6-v2 and CLIP ViT-B/32 both partition cleanly and then die here: schema.json:1: 6012: error: cannot parse value starting with: - The error points at a byte offset inside a temporary file, so nothing connects it back to the model. flatc does accept "inf" and "-inf", so emit those. The decompile direction needs the inverse rewrite, since flatc writes bare inf tokens that json.load will not parse; string literals are stepped over so a shader or key name containing "inf" is left alone. FlatBuffers JSON has no spelling for NaN at all, so raise there rather than emitting JSON that flatc cannot read. Adds round-trip coverage for the infinities, the NaN error, and the string-literal case. --- .../serialization/vulkan_graph_serialize.py | 62 ++++++++++++++++++- backends/vulkan/test/test_serialization.py | 57 +++++++++++++++++ 2 files changed, 117 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index 96f944560a8..beb7700bcc9 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -11,6 +11,7 @@ import importlib.resources as _resources import json import os +import re import tempfile from dataclasses import dataclass from typing import ClassVar, List @@ -27,8 +28,64 @@ from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile +# flatc's JSON dialect spells the non-finite floats "inf" / "-inf"; Python's +# json module spells them "Infinity" / "-Infinity" and rejects flatc's spelling. +# Neither side has a spelling the other accepts, so both directions are +# translated below. A graph carries a non-finite scalar whenever the model +# does -- the -inf fill value of a transformer attention mask is the common +# case -- and without this the failure surfaces as a flatc byte offset into a +# temporary file rather than anything pointing at the graph. +_JSON_STRING_RE = re.compile(r'"(?:[^"\\]|\\.)*"') +_FLATC_INF_RE = re.compile(r"(? str: + """Rewrite flatc's bare ``inf`` / ``-inf`` tokens so json.load accepts them. + + String literals are copied through untouched so that a shader or key name + containing "inf" is never rewritten. + """ + + def sub(segment: str) -> str: + return _FLATC_INF_RE.sub(lambda m: m.group(1) + "Infinity", segment) + + out = [] + last = 0 + for m in _JSON_STRING_RE.finditer(text): + out.append(sub(text[last : m.start()])) + out.append(m.group(0)) + last = m.end() + out.append(sub(text[last:])) + return "".join(out) + + def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes: - vk_graph_json = json.dumps(vk_graph, cls=_DataclassEncoder) + vk_graph_json = json.dumps(vk_graph, cls=_VkGraphEncoder) with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") @@ -63,7 +120,8 @@ def flatbuffer_to_vk_graph(flatbuffers: bytes) -> VkGraph: json_path = os.path.join(d, "schema.json") with open(json_path, "rb") as output_file: - return _json_to_dataclass(json.load(output_file), VkGraph) + raw = output_file.read().decode("utf-8") + return _json_to_dataclass(json.loads(_flatc_json_to_python_json(raw)), VkGraph) def extract_vk_flatbuffer(data: bytes) -> bytes: diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index 71a6980635a..7cf776c536b 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -19,6 +19,7 @@ ) from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( + Double, IntList, OperatorCall, String, @@ -269,3 +270,59 @@ def test_serialize_deserialize_vkgraph(self) -> None: out_vk_graph = flatbuffer_to_vk_graph(bs) self.assertEqual(in_vk_graph, out_vk_graph) + + def test_serialize_deserialize_non_finite_floats(self) -> None: + # flatc's JSON dialect spells the infinities "inf" / "-inf" while + # Python's json module spells them "Infinity" / "-Infinity" and rejects + # flatc's spelling, so both directions need translating. A graph picks + # up a non-finite scalar whenever the model has one -- the -inf fill + # value of a transformer attention mask being the usual source. + in_vk_graph = VkGraph( + version="1", + chain=[], + values=[ + VkValue(value=Double(double_val=float("-inf"))), + VkValue(value=Double(double_val=float("inf"))), + VkValue(value=Double(double_val=1.5)), + ], + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + + out_vk_graph = flatbuffer_to_vk_graph(convert_to_flatbuffer(in_vk_graph)) + + self.assertEqual(in_vk_graph, out_vk_graph) + + def test_serialize_nan_float_raises(self) -> None: + # FlatBuffers JSON has no spelling for NaN at all, so report it rather + # than emitting JSON that flatc cannot parse. + vk_graph = VkGraph( + version="1", + chain=[], + values=[VkValue(value=Double(double_val=float("nan")))], + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + + with self.assertRaisesRegex(ValueError, "NaN"): + convert_to_flatbuffer(vk_graph) + + def test_deserialize_leaves_inf_inside_strings_alone(self) -> None: + # The inf-token rewrite must not reach into string literals. + in_vk_graph = VkGraph( + version="1", + chain=[OperatorCall(node_id=1, name="inf_shader", args=[])], + values=[VkValue(value=String(string_val="value: inf, -inf"))], + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + + out_vk_graph = flatbuffer_to_vk_graph(convert_to_flatbuffer(in_vk_graph)) + + self.assertEqual(in_vk_graph, out_vk_graph)