diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index b74a7fb1f8e..78ac51c8808 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import ctypes +import hashlib import logging import operator from types import NoneType @@ -25,6 +27,7 @@ is_symint_node, TensorRepr, ) +from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.utils import DelegateMappingBuilder from executorch.exir.tensor import TensorSpec @@ -56,6 +59,7 @@ def __init__( self.input_ids = [] self.output_ids = [] self.const_tensors = [] + self.named_data_store = NamedDataStore() # Mapping from Node to VkValue id self.node_to_value_ids = {} @@ -129,8 +133,36 @@ def get_param_tensor(self, node: Node) -> torch.Tensor: def maybe_add_constant_tensor(self, node: Node) -> int: constant_id = -1 if is_param_node(self.program, node): - constant_id = len(self.const_tensors) - self.const_tensors.append(self.get_param_tensor(node)) + tensor = self.get_param_tensor(node) + + # Serialize tensor data to bytes + tensor = tensor.contiguous() + size = tensor.untyped_storage().nbytes() + + if size > 0: + array_type = ctypes.c_char * size + array = ctypes.cast( + tensor.untyped_storage().data_ptr(), + ctypes.POINTER(array_type), + ).contents + + # Generate SHA256 hash as the named key + tensor_bytes = bytes(array) + sha256_hash = hashlib.sha256(tensor_bytes) + named_key = sha256_hash.hexdigest() + + # Add to named data store with 16-byte alignment (matching XNNPACK) + self.named_data_store.add_named_data( + named_key, tensor_bytes, alignment=16 + ) + + # Create VkBytes entry with named_key and set offset to indicate named data usage + constant_id = len(self.const_tensors) + self.const_tensors.append((named_key, size)) + else: + # Handle empty tensors + constant_id = len(self.const_tensors) + self.const_tensors.append(None) return constant_id diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index 2ceedf73d10..db682f4e67e 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -191,10 +191,21 @@ def serialize_constant_tensors( current_offset = len(raw_bytes) for tensor in const_tensors: - if tensor.numel() == 0: + # The tensor data is stored in the named data map + if isinstance(tensor, tuple): + named_key, size = tensor + vk_graph.constants.append( + VkBytes( + offset=18446744073709551615, # UINT64_MAX to indicate named data + length=size, + named_key=named_key, + ) + ) + elif tensor is None or ( + isinstance(tensor, torch.Tensor) and tensor.numel() == 0 + ): vk_graph.constants.append(VkBytes(current_offset, 0)) - continue - else: + elif isinstance(tensor, torch.Tensor): array_type = ctypes.c_char * tensor.untyped_storage().nbytes() array = ctypes.cast( tensor.untyped_storage().data_ptr(), @@ -208,6 +219,8 @@ def serialize_constant_tensors( vk_graph.constants.append(VkBytes(current_offset, len(tensor_bytes))) current_offset += aligned_size(len(tensor_bytes)) + else: + raise ValueError(f"Unsupported constant tensor type: {type(tensor)}") def serialize_custom_shaders( diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index 8c1165a89df..1816d9b12de 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -229,4 +229,5 @@ def preprocess( # noqa: C901 vk_graph, graph_builder.const_tensors, [] ), debug_handle_map=graph_builder.delegate_mapping_builder.get_delegate_mapping(), + data_store_output=graph_builder.named_data_store.get_named_data_store_output(), )