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