xref: /aosp_15_r20/external/executorch/backends/vulkan/test/test_vulkan_delegate_header.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
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
7import unittest
8
9from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
10    VulkanDelegateHeader,
11)
12
13EXAMPLE_FLATBUFFER_OFFSET: int = 0x11223344
14EXAMPLE_FLATBUFFER_SIZE: int = 0x55667788
15EXAMPLE_BYTES_OFFSET: int = EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE
16EXAMPLE_BYTES_SIZE: int = 0x99AABBCC99AABBCC
17
18# If header layout or magic changes, this test must change too.
19# The layout of the header is a contract, not an implementation detail
20EXAMPLE_HEADER_DATA: bytes = (
21    # zeros
22    b"\x00\x00\x00\x00"
23    # magic
24    + b"VH00"
25    # All Values below are littl Endian
26    # header length
27    + b"\x1E\x00"
28    # Flatbuffer Offset
29    + b"\x44\x33\x22\x11"
30    # Flatbuffer Size
31    + b"\x88\x77\x66\x55"
32    # Bytes Data Offset
33    + b"\xCC\xAA\x88\x66"
34    # Bytes Data Size
35    + b"\xCC\xBB\xAA\x99\xCC\xBB\xAA\x99"
36)
37
38
39class TestVulkanDelegateHeader(unittest.TestCase):
40    def test_to_bytes(self) -> None:
41        header = VulkanDelegateHeader(
42            EXAMPLE_FLATBUFFER_OFFSET,
43            EXAMPLE_FLATBUFFER_SIZE,
44            EXAMPLE_BYTES_OFFSET,
45            EXAMPLE_BYTES_SIZE,
46        )
47        self.assertEqual(header.to_bytes(), EXAMPLE_HEADER_DATA)
48        self.assertTrue(header.is_valid())
49
50    def test_from_bytes(self) -> None:
51        header = VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA)
52        self.assertEqual(header.flatbuffer_offset, EXAMPLE_FLATBUFFER_OFFSET)
53        self.assertEqual(header.flatbuffer_size, EXAMPLE_FLATBUFFER_SIZE)
54        self.assertEqual(header.bytes_offset, EXAMPLE_BYTES_OFFSET)
55        self.assertEqual(header.bytes_size, EXAMPLE_BYTES_SIZE)
56
57    def test_invalid_metadata(self) -> None:
58        WRONG_MAGIC_DATA = EXAMPLE_HEADER_DATA[0:4] + b"YT01" + EXAMPLE_HEADER_DATA[8:]
59        with self.assertRaisesRegex(
60            ValueError,
61            "Expected magic bytes to be b'VH00', but got b'YT01'",
62        ):
63            VulkanDelegateHeader.from_bytes(WRONG_MAGIC_DATA)
64
65        WRONG_LENGTH_DATA = (
66            EXAMPLE_HEADER_DATA[0:8] + b"\x1D\x00" + EXAMPLE_HEADER_DATA[10:]
67        )
68        with self.assertRaisesRegex(
69            ValueError, "Expected header to be 30 bytes, but got 29 bytes."
70        ):
71            VulkanDelegateHeader.from_bytes(WRONG_LENGTH_DATA)
72
73        with self.assertRaisesRegex(
74            ValueError, "Expected header to be 30 bytes, but got 31 bytes."
75        ):
76            VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA + b"\x00")
77
78    def test_invalid_flatbuffer_size(self) -> None:
79        header = VulkanDelegateHeader(
80            EXAMPLE_FLATBUFFER_OFFSET,
81            0,
82            EXAMPLE_BYTES_OFFSET,
83            EXAMPLE_BYTES_SIZE,
84        )
85
86        with self.assertRaises(ValueError):
87            header.to_bytes()
88
89    def test_invalid_constants_offset(self) -> None:
90        header = VulkanDelegateHeader(
91            EXAMPLE_FLATBUFFER_OFFSET,
92            EXAMPLE_FLATBUFFER_SIZE,
93            EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE - 1,
94            EXAMPLE_BYTES_SIZE,
95        )
96
97        with self.assertRaises(ValueError):
98            header.to_bytes()
99
100    def test_to_bytes_same_as_from_bytes(self) -> None:
101        header = VulkanDelegateHeader.from_bytes(EXAMPLE_HEADER_DATA)
102
103        to_bytes = header.to_bytes()
104        self.assertEqual(EXAMPLE_HEADER_DATA, to_bytes)
105