xref: /aosp_15_r20/external/executorch/backends/vulkan/test/test_serialization.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 ctypes
8import random
9import unittest
10from typing import List
11
12import torch
13
14from executorch.backends.vulkan.serialization.vulkan_graph_schema import VkGraph
15
16from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
17    serialize_vulkan_graph,
18    VulkanDelegateHeader,
19)
20
21
22class TestSerialization(unittest.TestCase):
23    def _generate_random_const_tensors(self, num_tensors: int) -> List[torch.Tensor]:
24        """
25        Helper function to generate `num_tensor` buffers of random sizes and random contents,
26        we return a tuple of (list_of_buffers, list_of_mem_sizes),
27        """
28        tensors = []
29        for _ in range(num_tensors):
30            width = random.randint(4, 100)
31            height = random.randint(4, 100)
32            channels = random.randint(2, 8)
33
34            tensor = torch.randn(channels, width, height)
35            tensors.append(tensor)
36
37        return tensors
38
39    def test_serialize_vulkan_binary(self):
40        vk_graph = VkGraph(
41            version="0",
42            chain=[],
43            values=[],
44            input_ids=[],
45            output_ids=[],
46            constants=[],
47            shaders=[],
48        )
49        const_tensors = self._generate_random_const_tensors(5)
50
51        serialized_binary = serialize_vulkan_graph(vk_graph, const_tensors, [])
52
53        # Check header
54        self.assertEqual(serialized_binary[0:4], b"\x00\x00\x00\x00")
55        self.assertEqual(serialized_binary[VulkanDelegateHeader.MAGIC_IX], b"VH00")
56        flatbuffer_offset = int.from_bytes(
57            serialized_binary[VulkanDelegateHeader.FLATBUFFER_OFFSET_IX],
58            byteorder="little",
59        )
60        constants_offset = int.from_bytes(
61            serialized_binary[VulkanDelegateHeader.BYTES_OFFSET_IX],
62            byteorder="little",
63        )
64        constants_size = int.from_bytes(
65            serialized_binary[VulkanDelegateHeader.BYTES_SIZE_IX],
66            byteorder="little",
67        )
68
69        # Flatbuffer magic should be in the same spot as the Header's magic
70        self.assertEqual(
71            serialized_binary[flatbuffer_offset:][VulkanDelegateHeader.MAGIC_IX],
72            b"VK00",
73        )
74
75        constant_data_payload = serialized_binary[
76            constants_offset : constants_offset + constants_size
77        ]
78
79        # We check that constant data indexes stored in the vk_graph correctly index
80        # into the correct buffer in the constant data section
81        self.assertEqual(len(vk_graph.constants), len(const_tensors))
82        for bytes_range, tensor in zip(vk_graph.constants, const_tensors):
83            offset = bytes_range.offset
84            length = bytes_range.length
85
86            constant_data_bytes = constant_data_payload[offset : offset + length]
87
88            array_type = ctypes.c_char * tensor.untyped_storage().nbytes()
89            array = ctypes.cast(
90                tensor.untyped_storage().data_ptr(),
91                ctypes.POINTER(array_type),
92            ).contents
93
94            tensor_bytes = bytes(array)
95            self.assertEqual(constant_data_bytes, tensor_bytes)
96