Skip to content

Commit 5a9b190

Browse files
SS-JIASS-JIA
committed
[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-source-id: 215762567 Pull Request resolved: #2013
1 parent d25a219 commit 5a9b190

File tree

3 files changed

+240
-0
lines changed

3 files changed

+240
-0
lines changed

‎backends/vulkan/serialization/vulkan_graph_serialize.py

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@
88
import os
99
import tempfile
1010

11+
from dataclasses import dataclass
12+
from typing import ClassVar
13+
1114
# pyre-ignore[21]: Could not find module `executorch.exir._serialize._bindings`.
1215
import executorch.exir._serialize._bindings as bindings # @manual=//executorch/exir/_serialize:_bindings
1316

@@ -32,3 +35,112 @@ def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes:
3235
output_path = os.path.join(d, "schema.bin")
3336
with open(output_path, "rb") as output_file:
3437
return output_file.read()
38+
39+
40+
@dataclass
41+
class VulkanDelegateHeader:
42+
# Defines the byte region that each component of the header corresponds to
43+
MAGIC_IX: ClassVar[slice] = slice(4, 8)
44+
HEADER_SIZE_IX: ClassVar[slice] = slice(8, 10)
45+
FLATBUFFER_OFFSET_IX: ClassVar[slice] = slice(10, 14)
46+
FLATBUFFER_SIZE_IX: ClassVar[slice] = slice(14, 18)
47+
CONSTANTS_OFFSET_IX: ClassVar[slice] = slice(18, 22)
48+
CONSTANTS_SIZE_IX: ClassVar[slice] = slice(22, 30)
49+
SHADERS_OFFSET_IX: ClassVar[slice] = slice(30, 34)
50+
SHADERS_SIZE_IX: ClassVar[slice] = slice(34, 42)
51+
52+
# magic bytes that should be at the beginning of the header
53+
EXPECTED_MAGIC: ClassVar[bytes] = b"VKDG"
54+
# The length of the header in bytes
55+
EXPECTED_LENGTH: ClassVar[int] = 42
56+
57+
# Instance attributes, @dataclass will turn these into constructor args
58+
flatbuffer_offset: int
59+
flatbuffer_size: int
60+
constants_offset: int
61+
constants_size: int
62+
shaders_offset: int
63+
shaders_size: int
64+
65+
@staticmethod
66+
def from_bytes(data: bytes) -> "VulkanDelegateHeader":
67+
if len(data) > VulkanDelegateHeader.EXPECTED_LENGTH:
68+
raise ValueError(
69+
f"Expected header to be {VulkanDelegateHeader.EXPECTED_LENGTH} bytes, "
70+
f"but got {len(data)} bytes."
71+
)
72+
73+
magic_b: bytes = data[VulkanDelegateHeader.MAGIC_IX]
74+
75+
if magic_b != VulkanDelegateHeader.EXPECTED_MAGIC:
76+
raise ValueError(
77+
f"Expected magic bytes to be {VulkanDelegateHeader.EXPECTED_MAGIC}, "
78+
f"but got {magic_b}."
79+
)
80+
81+
length: int = int.from_bytes(
82+
data[VulkanDelegateHeader.HEADER_SIZE_IX], byteorder="little"
83+
)
84+
85+
if length != VulkanDelegateHeader.EXPECTED_LENGTH:
86+
raise ValueError(
87+
f"Expected header to be {VulkanDelegateHeader.EXPECTED_LENGTH} bytes, "
88+
f"but got {length} bytes."
89+
)
90+
91+
flatbuffer_offset_b: bytes = data[VulkanDelegateHeader.FLATBUFFER_OFFSET_IX]
92+
flatbuffer_size_b: bytes = data[VulkanDelegateHeader.FLATBUFFER_SIZE_IX]
93+
constants_offset_b: bytes = data[VulkanDelegateHeader.CONSTANTS_OFFSET_IX]
94+
constants_size_b: bytes = data[VulkanDelegateHeader.CONSTANTS_SIZE_IX]
95+
shaders_offset_b: bytes = data[VulkanDelegateHeader.SHADERS_OFFSET_IX]
96+
shaders_size_b: bytes = data[VulkanDelegateHeader.SHADERS_SIZE_IX]
97+
98+
return VulkanDelegateHeader(
99+
flatbuffer_offset=int.from_bytes(flatbuffer_offset_b, byteorder="little"),
100+
flatbuffer_size=int.from_bytes(flatbuffer_size_b, byteorder="little"),
101+
constants_offset=int.from_bytes(constants_offset_b, byteorder="little"),
102+
constants_size=int.from_bytes(constants_size_b, byteorder="little"),
103+
shaders_offset=int.from_bytes(shaders_offset_b, byteorder="little"),
104+
shaders_size=int.from_bytes(shaders_size_b, byteorder="little"),
105+
)
106+
107+
def is_valid(self) -> bool:
108+
if self.flatbuffer_size <= 0:
109+
return False
110+
111+
expected_offset = self.flatbuffer_offset + self.flatbuffer_size
112+
if self.constants_offset < expected_offset:
113+
return False
114+
115+
if self.constants_size <= 0:
116+
return False
117+
118+
expected_offset = self.constants_offset + self.constants_size
119+
if self.shaders_offset < expected_offset:
120+
return False
121+
122+
# shaders_size can be 0
123+
124+
return True
125+
126+
def to_bytes(self) -> bytes:
127+
if not self.is_valid():
128+
raise ValueError("VulkanDelegateHeader instance contains invalid values")
129+
130+
data: bytes = (
131+
# 4 bytes of padding for magic bytes, this is so that the header magic
132+
# bytes is in the same position as the flatbuffer header magic bytes
133+
b"\x00\x00\x00\x00"
134+
+ self.EXPECTED_MAGIC
135+
+ self.EXPECTED_LENGTH.to_bytes(2, byteorder="little")
136+
+ self.flatbuffer_offset.to_bytes(4, byteorder="little")
137+
+ self.flatbuffer_size.to_bytes(4, byteorder="little")
138+
+ self.constants_offset.to_bytes(4, byteorder="little")
139+
+ self.constants_size.to_bytes(8, byteorder="little")
140+
+ self.shaders_offset.to_bytes(4, byteorder="little")
141+
+ self.shaders_size.to_bytes(8, byteorder="little")
142+
)
143+
144+
assert len(data) == VulkanDelegateHeader.EXPECTED_LENGTH
145+
146+
return data

‎backends/vulkan/test/TARGETS

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,3 +20,13 @@ python_unittest(
2020
"//executorch/kernels/portable:custom_ops_generated_lib",
2121
],
2222
)
23+
24+
python_unittest(
25+
name = "test_vulkan_delegate_header",
26+
srcs = [
27+
"test_vulkan_delegate_header.py",
28+
],
29+
deps = [
30+
"//executorch/backends/vulkan:vulkan_preprocess",
31+
],
32+
)

‎backends/vulkan/test/test_vulkan_delegate_header.py

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
import unittest
8+
9+
from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
10+
VulkanDelegateHeader,
11+
)
12+
13+
EXAMPLE_FLATBUFFER_OFFSET: int = 0x11223344
14+
EXAMPLE_FLATBUFFER_SIZE: int = 0x55667788
15+
EXAMPLE_CONSTANTS_OFFSET: int = EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE
16+
EXAMPLE_CONSTANTS_SIZE: int = 0x99AABB
17+
EXAMPLE_SHADERS_OFFSET: int = EXAMPLE_CONSTANTS_OFFSET + EXAMPLE_CONSTANTS_SIZE
18+
EXAMPLE_SHADERS_SIZE: int = 0x11223344
19+
20+
# If header layout or magic changes, this test must change too.
21+
# The layout of the header is a contract, not an implementation detail
22+
EXAMPLE_HEADER_DATA: bytes = (
23+
# zeros
24+
b"\x00\x00\x00\x00"
25+
# magic
26+
+ b"VKDG"
27+
# All Values below are littl Endian
28+
# header length
29+
+ b"\x2A\x00"
30+
# Flatbuffer Offset
31+
+ b"\x44\x33\x22\x11"
32+
# Flatbuffer Size
33+
+ b"\x88\x77\x66\x55"
34+
# Constant Data Offset
35+
+ b"\xCC\xAA\x88\x66"
36+
# Constant Data Size
37+
+ b"\xBB\xAA\x99\x00\x00\x00\x00\x00"
38+
# Shader Offset
39+
+ b"\x87\x55\x22\x67"
40+
# Shader Size
41+
+ b"\x44\x33\x22\x11\x00\x00\x00\x00"
42+
)
43+
44+
45+
class TestVulkanDelegateHeader(unittest.TestCase):
46+
def test_to_bytes(self) -> None:
47+
header = VulkanDelegateHeader(
48+
EXAMPLE_FLATBUFFER_OFFSET,
49+
EXAMPLE_FLATBUFFER_SIZE,
50+
EXAMPLE_CONSTANTS_OFFSET,
51+
EXAMPLE_CONSTANTS_SIZE,
52+
EXAMPLE_SHADERS_OFFSET,
53+
EXAMPLE_SHADERS_SIZE,
54+
)
55+
self.assertEqual(header.to_bytes(), EXAMPLE_HEADER_DATA)
56+
self.assertTrue(header.is_valid())
57+
58+
def test_from_bytes(self) -> None:
59+
header = VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA)
60+
self.assertEqual(header.flatbuffer_offset, EXAMPLE_FLATBUFFER_OFFSET)
61+
self.assertEqual(header.flatbuffer_size, EXAMPLE_FLATBUFFER_SIZE)
62+
self.assertEqual(header.constants_offset, EXAMPLE_CONSTANTS_OFFSET)
63+
self.assertEqual(header.constants_size, EXAMPLE_CONSTANTS_SIZE)
64+
self.assertEqual(header.shaders_offset, EXAMPLE_SHADERS_OFFSET)
65+
self.assertEqual(header.shaders_size, EXAMPLE_SHADERS_SIZE)
66+
67+
def test_invalid_metadata(self) -> None:
68+
WRONG_MAGIC_DATA = EXAMPLE_HEADER_DATA[0:4] + b"YT01" + EXAMPLE_HEADER_DATA[8:]
69+
with self.assertRaisesRegex(
70+
ValueError,
71+
"Expected magic bytes to be b'VKDG', but got b'YT01'",
72+
):
73+
VulkanDelegateHeader.from_bytes(WRONG_MAGIC_DATA)
74+
75+
WRONG_LENGTH_DATA = (
76+
EXAMPLE_HEADER_DATA[0:8] + b"\x1D\x00" + EXAMPLE_HEADER_DATA[10:]
77+
)
78+
with self.assertRaisesRegex(
79+
ValueError, "Expected header to be 42 bytes, but got 29 bytes."
80+
):
81+
VulkanDelegateHeader.from_bytes(WRONG_LENGTH_DATA)
82+
83+
with self.assertRaisesRegex(
84+
ValueError, "Expected header to be 42 bytes, but got 43 bytes."
85+
):
86+
VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA + b"\x00")
87+
88+
def test_invalid_flatbuffer_size(self) -> None:
89+
header = VulkanDelegateHeader(
90+
EXAMPLE_FLATBUFFER_OFFSET,
91+
0,
92+
EXAMPLE_CONSTANTS_OFFSET,
93+
EXAMPLE_CONSTANTS_SIZE,
94+
EXAMPLE_SHADERS_OFFSET,
95+
EXAMPLE_SHADERS_SIZE,
96+
)
97+
98+
with self.assertRaises(ValueError):
99+
header.to_bytes()
100+
101+
def test_invalid_constants_offset(self) -> None:
102+
header = VulkanDelegateHeader(
103+
EXAMPLE_FLATBUFFER_OFFSET,
104+
EXAMPLE_FLATBUFFER_SIZE,
105+
EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE - 1,
106+
EXAMPLE_CONSTANTS_SIZE,
107+
EXAMPLE_SHADERS_OFFSET,
108+
EXAMPLE_SHADERS_SIZE,
109+
)
110+
111+
with self.assertRaises(ValueError):
112+
header.to_bytes()
113+
114+
def test_to_bytes_same_as_from_bytes(self) -> None:
115+
header = VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA)
116+
117+
to_bytes = header.to_bytes()
118+
self.assertEquals(EXAMPLE_HEADER_DATA, to_bytes)

0 commit comments

Comments
 (0)