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

## 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.

Differential Revision: [D53957853](https://our.internmc.facebook.com/intern/diff/D53957853/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/serialization/vulkan_graph_serialize.py

@@ -44,23 +44,19 @@ class VulkanDelegateHeader:
     HEADER_SIZE_IX: ClassVar[slice] = slice(8, 10)
     FLATBUFFER_OFFSET_IX: ClassVar[slice] = slice(10, 14)
     FLATBUFFER_SIZE_IX: ClassVar[slice] = slice(14, 18)
-    CONSTANTS_OFFSET_IX: ClassVar[slice] = slice(18, 22)
-    CONSTANTS_SIZE_IX: ClassVar[slice] = slice(22, 30)
-    SHADERS_OFFSET_IX: ClassVar[slice] = slice(30, 34)
-    SHADERS_SIZE_IX: ClassVar[slice] = slice(34, 42)
+    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"VKDG"
+    EXPECTED_MAGIC: ClassVar[bytes] = b"VH00"
     # The length of the header in bytes
-    EXPECTED_LENGTH: ClassVar[int] = 42
+    EXPECTED_LENGTH: ClassVar[int] = 30
 
     # Instance attributes, @dataclass will turn these into constructor args
     flatbuffer_offset: int
     flatbuffer_size: int
-    constants_offset: int
-    constants_size: int
-    shaders_offset: int
-    shaders_size: int
+    bytes_offset: int
+    bytes_size: int
 
     @staticmethod
     def from_bytes(data: bytes) -> "VulkanDelegateHeader":
@@ -90,37 +86,27 @@ def from_bytes(data: bytes) -> "VulkanDelegateHeader":
 
         flatbuffer_offset_b: bytes = data[VulkanDelegateHeader.FLATBUFFER_OFFSET_IX]
         flatbuffer_size_b: bytes = data[VulkanDelegateHeader.FLATBUFFER_SIZE_IX]
-        constants_offset_b: bytes = data[VulkanDelegateHeader.CONSTANTS_OFFSET_IX]
-        constants_size_b: bytes = data[VulkanDelegateHeader.CONSTANTS_SIZE_IX]
-        shaders_offset_b: bytes = data[VulkanDelegateHeader.SHADERS_OFFSET_IX]
-        shaders_size_b: bytes = data[VulkanDelegateHeader.SHADERS_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"),
-            constants_offset=int.from_bytes(constants_offset_b, byteorder="little"),
-            constants_size=int.from_bytes(constants_size_b, byteorder="little"),
-            shaders_offset=int.from_bytes(shaders_offset_b, byteorder="little"),
-            shaders_size=int.from_bytes(shaders_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.constants_offset < expected_offset:
+        if self.bytes_offset < expected_offset:
             return False
 
-        if self.constants_size <= 0:
+        if self.bytes_size < 0:
             return False
 
-        expected_offset = self.constants_offset + self.constants_size
-        if self.shaders_offset < expected_offset:
-            return False
-
-        # shaders_size can be 0
-
         return True
 
     def to_bytes(self) -> bytes:
@@ -135,10 +121,8 @@ def to_bytes(self) -> bytes:
             + 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.constants_offset.to_bytes(4, byteorder="little")
-            + self.constants_size.to_bytes(8, byteorder="little")
-            + self.shaders_offset.to_bytes(4, byteorder="little")
-            + self.shaders_size.to_bytes(8, byteorder="little")
+            + self.bytes_offset.to_bytes(4, byteorder="little")
+            + self.bytes_size.to_bytes(8, byteorder="little")
         )
 
         assert len(data) == VulkanDelegateHeader.EXPECTED_LENGTH

backends/vulkan/test/test_vulkan_delegate_header.py

@@ -12,33 +12,27 @@
 
 EXAMPLE_FLATBUFFER_OFFSET: int = 0x11223344
 EXAMPLE_FLATBUFFER_SIZE: int = 0x55667788
-EXAMPLE_CONSTANTS_OFFSET: int = EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE
-EXAMPLE_CONSTANTS_SIZE: int = 0x99AABB
-EXAMPLE_SHADERS_OFFSET: int = EXAMPLE_CONSTANTS_OFFSET + EXAMPLE_CONSTANTS_SIZE
-EXAMPLE_SHADERS_SIZE: int = 0x11223344
+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"VKDG"
+    + b"VH00"
     # All Values below are littl Endian
     # header length
-    + b"\x2A\x00"
+    + b"\x1E\x00"
     # Flatbuffer Offset
     + b"\x44\x33\x22\x11"
     # Flatbuffer Size
     + b"\x88\x77\x66\x55"
-    # Constant Data Offset
+    # Bytes Data Offset
     + b"\xCC\xAA\x88\x66"
-    # Constant Data Size
-    + b"\xBB\xAA\x99\x00\x00\x00\x00\x00"
-    # Shader Offset
-    + b"\x87\x55\x22\x67"
-    # Shader Size
-    + b"\x44\x33\x22\x11\x00\x00\x00\x00"
+    # Bytes Data Size
+    + b"\xCC\xBB\xAA\x99\xCC\xBB\xAA\x99"
 )
 
 
@@ -47,10 +41,8 @@ def test_to_bytes(self) -> None:
         header = VulkanDelegateHeader(
             EXAMPLE_FLATBUFFER_OFFSET,
             EXAMPLE_FLATBUFFER_SIZE,
-            EXAMPLE_CONSTANTS_OFFSET,
-            EXAMPLE_CONSTANTS_SIZE,
-            EXAMPLE_SHADERS_OFFSET,
-            EXAMPLE_SHADERS_SIZE,
+            EXAMPLE_BYTES_OFFSET,
+            EXAMPLE_BYTES_SIZE,
         )
         self.assertEqual(header.to_bytes(), EXAMPLE_HEADER_DATA)
         self.assertTrue(header.is_valid())
@@ -59,40 +51,36 @@ 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.constants_offset, EXAMPLE_CONSTANTS_OFFSET)
-        self.assertEqual(header.constants_size, EXAMPLE_CONSTANTS_SIZE)
-        self.assertEqual(header.shaders_offset, EXAMPLE_SHADERS_OFFSET)
-        self.assertEqual(header.shaders_size, EXAMPLE_SHADERS_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'VKDG', but got b'YT01'",
+            "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 42 bytes, but got 29 bytes."
+            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 42 bytes, but got 43 bytes."
+            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_CONSTANTS_OFFSET,
-            EXAMPLE_CONSTANTS_SIZE,
-            EXAMPLE_SHADERS_OFFSET,
-            EXAMPLE_SHADERS_SIZE,
+            EXAMPLE_BYTES_OFFSET,
+            EXAMPLE_BYTES_SIZE,
         )
 
         with self.assertRaises(ValueError):
@@ -103,9 +91,7 @@ def test_invalid_constants_offset(self) -> None:
             EXAMPLE_FLATBUFFER_OFFSET,
             EXAMPLE_FLATBUFFER_SIZE,
             EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE - 1,
-            EXAMPLE_CONSTANTS_SIZE,
-            EXAMPLE_SHADERS_OFFSET,
-            EXAMPLE_SHADERS_SIZE,
+            EXAMPLE_BYTES_SIZE,
         )
 
         with self.assertRaises(ValueError):