[ET-VK] Introduce VulkanDelegateHeader to manage constant data and shader data

SS-JIA · SS-JIA · commit d6b1dcb80076 · 2024-02-20T15:05:19.000-08:00
Pull Request resolved: #2013 ## Context This changeset is essentially a mirror of #1523 in the XNNPACK delegate which is the first step to enabling constant weight data to be serialized outside the flatbuffer blob to speed up deserialization time for large models. `VulkanDelegateHeader` is introduced which will be used later on to help interpret what different sections of the serialized binary blob corresponds to. The primary difference compared to `XNNHeader` which was added for the XNNPACK delegate is that fields are added to support custom compute shaders that will be serialized with the model. ghstack-source-id: 215799821 @exported-using-ghexport Differential Revision: [D53957853](https://our.internmc.facebook.com/intern/diff/D53957853/)
diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py
@@ -8,6 +8,9 @@
 import os
 import tempfile
 
+from dataclasses import dataclass
+from typing import ClassVar
+
 # pyre-ignore[21]: Could not find module `executorch.exir._serialize._bindings`.
 import executorch.exir._serialize._bindings as bindings  # @manual=//executorch/exir/_serialize:_bindings
 
@@ -32,3 +35,96 @@ def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes:
         output_path = os.path.join(d, "schema.bin")
         with open(output_path, "rb") as output_file:
             return output_file.read()
+
+
+@dataclass
+class VulkanDelegateHeader:
+    # Defines the byte region that each component of the header corresponds to
+    MAGIC_IX: ClassVar[slice] = slice(4, 8)
+    HEADER_SIZE_IX: ClassVar[slice] = slice(8, 10)
+    FLATBUFFER_OFFSET_IX: ClassVar[slice] = slice(10, 14)
+    FLATBUFFER_SIZE_IX: ClassVar[slice] = slice(14, 18)
+    BYTES_OFFSET_IX: ClassVar[slice] = slice(18, 22)
+    BYTES_SIZE_IX: ClassVar[slice] = slice(22, 30)
+
+    # magic bytes that should be at the beginning of the header
+    EXPECTED_MAGIC: ClassVar[bytes] = b"VH00"
+    # The length of the header in bytes
+    EXPECTED_LENGTH: ClassVar[int] = 30
+
+    # Instance attributes, @dataclass will turn these into constructor args
+    flatbuffer_offset: int
+    flatbuffer_size: int
+    bytes_offset: int
+    bytes_size: int
+
+    @staticmethod
+    def from_bytes(data: bytes) -> "VulkanDelegateHeader":
+        if len(data) > VulkanDelegateHeader.EXPECTED_LENGTH:
+            raise ValueError(
+                f"Expected header to be {VulkanDelegateHeader.EXPECTED_LENGTH} bytes, "
+                f"but got {len(data)} bytes."
+            )
+
+        magic_b: bytes = data[VulkanDelegateHeader.MAGIC_IX]
+
+        if magic_b != VulkanDelegateHeader.EXPECTED_MAGIC:
+            raise ValueError(
+                f"Expected magic bytes to be {VulkanDelegateHeader.EXPECTED_MAGIC}, "
+                f"but got {magic_b}."
+            )
+
+        length: int = int.from_bytes(
+            data[VulkanDelegateHeader.HEADER_SIZE_IX], byteorder="little"
+        )
+
+        if length != VulkanDelegateHeader.EXPECTED_LENGTH:
+            raise ValueError(
+                f"Expected header to be {VulkanDelegateHeader.EXPECTED_LENGTH} bytes, "
+                f"but got {length} bytes."
+            )
+
+        flatbuffer_offset_b: bytes = data[VulkanDelegateHeader.FLATBUFFER_OFFSET_IX]
+        flatbuffer_size_b: bytes = data[VulkanDelegateHeader.FLATBUFFER_SIZE_IX]
+        bytes_offset_b: bytes = data[VulkanDelegateHeader.BYTES_OFFSET_IX]
+        bytes_size_b: bytes = data[VulkanDelegateHeader.BYTES_SIZE_IX]
+
+        return VulkanDelegateHeader(
+            flatbuffer_offset=int.from_bytes(flatbuffer_offset_b, byteorder="little"),
+            flatbuffer_size=int.from_bytes(flatbuffer_size_b, byteorder="little"),
+            bytes_offset=int.from_bytes(bytes_offset_b, byteorder="little"),
+            bytes_size=int.from_bytes(bytes_size_b, byteorder="little"),
+        )
+
+    def is_valid(self) -> bool:
+        if self.flatbuffer_size <= 0:
+            return False
+
+        expected_offset = self.flatbuffer_offset + self.flatbuffer_size
+        if self.bytes_offset < expected_offset:
+            return False
+
+        if self.bytes_size < 0:
+            return False
+
+        return True
+
+    def to_bytes(self) -> bytes:
+        if not self.is_valid():
+            raise ValueError("VulkanDelegateHeader instance contains invalid values")
+
+        data: bytes = (
+            # 4 bytes of padding for magic bytes, this is so that the header magic
+            # bytes is in the same position as the flatbuffer header magic bytes
+            b"\x00\x00\x00\x00"
+            + self.EXPECTED_MAGIC
+            + self.EXPECTED_LENGTH.to_bytes(2, byteorder="little")
+            + self.flatbuffer_offset.to_bytes(4, byteorder="little")
+            + self.flatbuffer_size.to_bytes(4, byteorder="little")
+            + self.bytes_offset.to_bytes(4, byteorder="little")
+            + self.bytes_size.to_bytes(8, byteorder="little")
+        )
+
+        assert len(data) == VulkanDelegateHeader.EXPECTED_LENGTH
+
+        return data
diff --git a/backends/vulkan/test/TARGETS b/backends/vulkan/test/TARGETS
@@ -20,3 +20,13 @@ python_unittest(
         "//executorch/kernels/portable:custom_ops_generated_lib",
     ],
 )
+
+python_unittest(
+    name = "test_vulkan_delegate_header",
+    srcs = [
+        "test_vulkan_delegate_header.py",
+    ],
+    deps = [
+        "//executorch/backends/vulkan:vulkan_preprocess",
+    ],
+)
diff --git a/backends/vulkan/test/test_vulkan_delegate_header.py b/backends/vulkan/test/test_vulkan_delegate_header.py
@@ -0,0 +1,104 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
+    VulkanDelegateHeader,
+)
+
+EXAMPLE_FLATBUFFER_OFFSET: int = 0x11223344
+EXAMPLE_FLATBUFFER_SIZE: int = 0x55667788
+EXAMPLE_BYTES_OFFSET: int = EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE
+EXAMPLE_BYTES_SIZE: int = 0x99AABBCC99AABBCC
+
+# If header layout or magic changes, this test must change too.
+# The layout of the header is a contract, not an implementation detail
+EXAMPLE_HEADER_DATA: bytes = (
+    # zeros
+    b"\x00\x00\x00\x00"
+    # magic
+    + b"VH00"
+    # All Values below are littl Endian
+    # header length
+    + b"\x1E\x00"
+    # Flatbuffer Offset
+    + b"\x44\x33\x22\x11"
+    # Flatbuffer Size
+    + b"\x88\x77\x66\x55"
+    # Bytes Data Offset
+    + b"\xCC\xAA\x88\x66"
+    # Bytes Data Size
+    + b"\xCC\xBB\xAA\x99\xCC\xBB\xAA\x99"
+)
+
+
+class TestVulkanDelegateHeader(unittest.TestCase):
+    def test_to_bytes(self) -> None:
+        header = VulkanDelegateHeader(
+            EXAMPLE_FLATBUFFER_OFFSET,
+            EXAMPLE_FLATBUFFER_SIZE,
+            EXAMPLE_BYTES_OFFSET,
+            EXAMPLE_BYTES_SIZE,
+        )
+        self.assertEqual(header.to_bytes(), EXAMPLE_HEADER_DATA)
+        self.assertTrue(header.is_valid())
+
+    def test_from_bytes(self) -> None:
+        header = VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA)
+        self.assertEqual(header.flatbuffer_offset, EXAMPLE_FLATBUFFER_OFFSET)
+        self.assertEqual(header.flatbuffer_size, EXAMPLE_FLATBUFFER_SIZE)
+        self.assertEqual(header.bytes_offset, EXAMPLE_BYTES_OFFSET)
+        self.assertEqual(header.bytes_size, EXAMPLE_BYTES_SIZE)
+
+    def test_invalid_metadata(self) -> None:
+        WRONG_MAGIC_DATA = EXAMPLE_HEADER_DATA[0:4] + b"YT01" + EXAMPLE_HEADER_DATA[8:]
+        with self.assertRaisesRegex(
+            ValueError,
+            "Expected magic bytes to be b'VH00', but got b'YT01'",
+        ):
+            VulkanDelegateHeader.from_bytes(WRONG_MAGIC_DATA)
+
+        WRONG_LENGTH_DATA = (
+            EXAMPLE_HEADER_DATA[0:8] + b"\x1D\x00" + EXAMPLE_HEADER_DATA[10:]
+        )
+        with self.assertRaisesRegex(
+            ValueError, "Expected header to be 30 bytes, but got 29 bytes."
+        ):
+            VulkanDelegateHeader.from_bytes(WRONG_LENGTH_DATA)
+
+        with self.assertRaisesRegex(
+            ValueError, "Expected header to be 30 bytes, but got 31 bytes."
+        ):
+            VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA + b"\x00")
+
+    def test_invalid_flatbuffer_size(self) -> None:
+        header = VulkanDelegateHeader(
+            EXAMPLE_FLATBUFFER_OFFSET,
+            0,
+            EXAMPLE_BYTES_OFFSET,
+            EXAMPLE_BYTES_SIZE,
+        )
+
+        with self.assertRaises(ValueError):
+            header.to_bytes()
+
+    def test_invalid_constants_offset(self) -> None:
+        header = VulkanDelegateHeader(
+            EXAMPLE_FLATBUFFER_OFFSET,
+            EXAMPLE_FLATBUFFER_SIZE,
+            EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE - 1,
+            EXAMPLE_BYTES_SIZE,
+        )
+
+        with self.assertRaises(ValueError):
+            header.to_bytes()
+
+    def test_to_bytes_same_as_from_bytes(self) -> None:
+        header = VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA)
+
+        to_bytes = header.to_bytes()
+        self.assertEquals(EXAMPLE_HEADER_DATA, to_bytes)
