[ET-VK][AOT] Serialize constant tensors via NamedDataMap

Differential Revision: D80460034

Pull Request resolved: #13473

Author: SS-JIA

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
 

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(

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(),
         )