From 531698c5ef124c207ebc6434cef78e5297a9e84b Mon Sep 17 00:00:00 2001 From: Justin Yip Date: Mon, 25 Nov 2024 14:26:51 -0800 Subject: [PATCH 1/2] [ET-VK] Deserialize VkGraph in ET-VK Add logic to deserialize a VkGraph blob back python object. This allows us to get a implement debugging / visualization directly on the vulkan-exported program. Still extra works need to be done: From the entire bundle, need to extract the specific vulkan delegate first. Differential Revision: [D66443780](https://our.internmc.facebook.com/intern/diff/D66443780/) [ghstack-poisoned] --- .../serialization/vulkan_graph_serialize.py | 21 +++++++++- backends/vulkan/test/test_serialization.py | 40 ++++++++++++++++++- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index 37785f47521..0cc2a9540cc 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -19,9 +19,9 @@ VkBytes, VkGraph, ) -from executorch.exir._serialize._dataclass import _DataclassEncoder +from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass -from executorch.exir._serialize._flatbuffer import _flatc_compile +from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes: @@ -40,6 +40,23 @@ def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes: return output_file.read() +def flatbuffer_to_vk_graph(flatbuffers: bytes) -> VkGraph: + with tempfile.TemporaryDirectory() as d: + schema_path = os.path.join(d, "schema.fbs") + with open(schema_path, "wb") as schema_file: + schema_file.write(pkg_resources.resource_string(__name__, "schema.fbs")) + + bin_path = os.path.join(d, "schema.bin") + with open(bin_path, "wb") as bin_file: + bin_file.write(flatbuffers) + + _flatc_decompile(d, schema_path, bin_path, ["--raw-binary"]) + + 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) + + @dataclass class VulkanDelegateHeader: # Defines the byte region that each component of the header corresponds to diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index eb112d7b12b..a0802efcea8 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -11,9 +11,17 @@ import torch -from executorch.backends.vulkan.serialization.vulkan_graph_schema import VkGraph +from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( + IntList, + OperatorCall, + String, + VkGraph, + VkValue, +) from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( + convert_to_flatbuffer, + flatbuffer_to_vk_graph, serialize_vulkan_graph, VulkanDelegateHeader, ) @@ -93,3 +101,33 @@ def test_serialize_vulkan_binary(self): tensor_bytes = bytes(array) self.assertEqual(constant_data_bytes, tensor_bytes) + + def test_serialize_deserialize_vkgraph(self): + in_vk_graph = VkGraph( + version="1", + chain=[ + OperatorCall(node_id=1, name="foo", args=[1, 2, 3]), + OperatorCall(node_id=2, name="bar", args=[]), + ], + values=[ + VkValue( + value=String( + string_val="abc", + ), + ), + VkValue( + value=IntList( + items=[-1, -4, 2], + ), + ), + ], + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + + bs = convert_to_flatbuffer(in_vk_graph) + out_vk_graph = flatbuffer_to_vk_graph(bs) + + self.assertEqual(in_vk_graph, out_vk_graph) From 85bc8c5baf49d7effc74db0669f85e02e9a33f42 Mon Sep 17 00:00:00 2001 From: Justin Yip Date: Mon, 25 Nov 2024 14:50:10 -0800 Subject: [PATCH 2/2] Update on "[ET-VK] Deserialize VkGraph in ET-VK" Add logic to deserialize a VkGraph blob back python object. This allows us to get a implement debugging / visualization directly on the vulkan-exported program. Still extra works need to be done: From the entire bundle, need to extract the specific vulkan delegate first. Differential Revision: [D66443780](https://our.internmc.facebook.com/intern/diff/D66443780/) [ghstack-poisoned] --- backends/vulkan/serialization/vulkan_graph_serialize.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index 0cc2a9540cc..ade7530c34e 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -41,6 +41,8 @@ def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes: def flatbuffer_to_vk_graph(flatbuffers: bytes) -> VkGraph: + # Following similar (de)serialization logic on other backends: + # https://github.com/pytorch/executorch/blob/main/backends/qualcomm/serialization/qc_schema_serialize.py#L33 with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") with open(schema_path, "wb") as schema_file: