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

Pull Request resolved: #13473



When exporting models to Vulkan backend, save constant tensors in the NamedDataMap instead of the constant data section of the delegate header.

## Motivation

Prevent screen blackout (Llama 3.2 1B) / device crash (Llama 3.2 3B) when running Llama 3.2 models on Samsung Galaxy S24. This behaviour is related to high peak memory usage when loading the model. For more information, see the top diff/PR in the stack.

## Context

This change is based on the equivalent change D70315207/#9153 in XNNPACK.
ghstack-source-id: 303830114

Differential Revision: [D80460034](https://our.internmc.facebook.com/intern/diff/D80460034/)

Author: ssjia

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