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)