From 61517aefe9e41f66f1774cd5a991e8f74417945c Mon Sep 17 00:00:00 2001 From: jorgep31415 Date: Tue, 4 Feb 2025 16:13:47 -0800 Subject: [PATCH] [ET-VK] Add util to extract VkGraph from Program Pull Request resolved: https://github.com/pytorch/executorch/pull/8199 Missing piece to facilitate VkGraph debugging. Differential Revision: [D69151587](https://our.internmc.facebook.com/intern/diff/D69151587/) ghstack-source-id: 264726843 --- backends/vulkan/serialization/vulkan_graph_serialize.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index c97ea69a435..ebb13bbb97d 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -61,6 +61,15 @@ def flatbuffer_to_vk_graph(flatbuffers: bytes) -> VkGraph: return _json_to_dataclass(json.load(output_file), VkGraph) +def extract_vk_flatbuffer(data: bytes) -> bytes: + h: VulkanDelegateHeader = VulkanDelegateHeader.from_bytes( + data[: VulkanDelegateHeader.EXPECTED_LENGTH] + ) + start = h.flatbuffer_offset + end = h.flatbuffer_offset + h.flatbuffer_size + return data[start:end] + + @dataclass class VulkanDelegateHeader: # Defines the byte region that each component of the header corresponds to