[executorch][vulkan] Add payload-bounded constant sharding

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -287,6 +287,21 @@ def parse_compile_options(compile_options: Dict[str, Any]) -> List[CompileSpec]:
     compile_specs = []
 
     for key, value in compile_options.items():
+        if key == "external_constants_max_data_bytes":
+            # Validate at the user-facing option boundary. Preprocess and the
+            # data store repeat validation because they can be called directly.
+            if (
+                isinstance(value, bool)
+                or not isinstance(value, int)
+                or value <= 0
+                or value >= 1 << 64
+            ):
+                raise ValueError(
+                    "external_constants_max_data_bytes must be a positive uint64"
+                )
+            compile_specs.append(CompileSpec(key, value.to_bytes(8, "little")))
+            continue
+
         if isinstance(value, (VkStorageType, VkMemoryLayout)):
             value_bytes = int(value).to_bytes(4, byteorder="little")
             compile_specs.append(CompileSpec(key, value_bytes))

backends/vulkan/test/TARGETS

@@ -50,6 +50,20 @@ python_unittest(
     ],
 )
 
+python_unittest(
+    name = "test_vulkan_compile_options",
+    srcs = [
+        "test_vulkan_compile_options.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/vulkan:vulkan_preprocess",
+        "//executorch/backends/vulkan/partitioner:vulkan_partitioner",
+        "//executorch/exir/_serialize:lib",
+        "//executorch/exir:lib",
+    ],
+)
+
 python_unittest(
     name = "test_serialization",
     srcs = [

backends/vulkan/test/test_vulkan_compile_options.py

@@ -6,11 +6,18 @@
 
 import unittest
 from typing import Any, Dict
+from unittest.mock import MagicMock, patch
 
 from executorch.backends.vulkan.partitioner.vulkan_partitioner import (
     parse_compile_options,
 )
-from executorch.backends.vulkan.vulkan_preprocess import parse_compile_spec
+from executorch.backends.vulkan.vulkan_preprocess import (
+    parse_compile_spec,
+    VulkanBackend,
+)
+from executorch.exir._serialize._named_data_store import NamedDataStore
+from executorch.exir._serialize.data_serializer import DataEntry
+from executorch.exir.backend.compile_spec_schema import CompileSpec
 
 
 class TestVulkanCompileOptions(unittest.TestCase):
@@ -38,10 +45,97 @@ def test_force_fp16_round_trips(self) -> None:
         round_tripped = self._round_trip({"force_fp16": True})
         self.assertTrue(round_tripped.get("force_fp16"))
 
+    def test_external_constants_max_data_bytes_round_trips_uint64_bounds(
+        self,
+    ) -> None:
+        for value in (1, (1 << 64) - 1):
+            with self.subTest(value=value):
+                self.assertEqual(
+                    self._round_trip({"external_constants_max_data_bytes": value}).get(
+                        "external_constants_max_data_bytes"
+                    ),
+                    value,
+                )
+
+    def test_external_constants_max_data_bytes_rejects_invalid_values(self) -> None:
+        invalid_values: list[Any] = [True, 0, -1, 1 << 64, 1.5, "10"]
+        for value in invalid_values:
+            with self.subTest(value=value), self.assertRaisesRegex(
+                ValueError, "positive uint64"
+            ):
+                parse_compile_options({"external_constants_max_data_bytes": value})
+
+    def test_external_constants_max_data_bytes_rejects_invalid_encoding(
+        self,
+    ) -> None:
+        for payload in (b"", b"\x01", b"\x01" * 7, b"\x01" * 9):
+            with self.subTest(payload=payload), self.assertRaisesRegex(
+                ValueError, "encoded as uint64"
+            ):
+                parse_compile_spec(
+                    [CompileSpec("external_constants_max_data_bytes", payload)]
+                )
+        with self.assertRaisesRegex(ValueError, "positive uint64"):
+            parse_compile_spec(
+                [CompileSpec("external_constants_max_data_bytes", b"\x00" * 8)]
+            )
+
+    def _preprocess_named_data(self, options: Dict[str, Any]):
+        store = NamedDataStore()
+        graph_builder = MagicMock()
+        graph_builder.named_data_store = store
+
+        def build_graph():
+            store.add_named_data("constant", b"constant", 16)
+            return MagicMock()
+
+        graph_builder.build_graph.side_effect = build_graph
+        graph_builder.delegate_mapping_builder.get_delegate_mapping.return_value = {}
+        program = MagicMock()
+
+        with patch.object(
+            store, "externalize_pte_data", wraps=store.externalize_pte_data
+        ) as externalize_pte_data, patch(
+            "executorch.backends.vulkan.vulkan_preprocess."
+            "unsafe_remove_auto_functionalized_pass",
+            side_effect=lambda value: value,
+        ), patch(
+            "executorch.backends.vulkan.vulkan_preprocess.apply_passes",
+            side_effect=lambda value, _passes: value,
+        ), patch(
+            "executorch.backends.vulkan.vulkan_preprocess.VkGraphBuilder",
+            return_value=graph_builder,
+        ), patch(
+            "executorch.backends.vulkan.vulkan_preprocess.serialize_vulkan_graph",
+            return_value=b"vk_graph",
+        ):
+            result = VulkanBackend.preprocess(program, parse_compile_options(options))
+        return result.data_store_output, externalize_pte_data
+
+    def test_external_constants_default_keeps_constants_inline(self) -> None:
+        output, externalize_pte_data = self._preprocess_named_data({})
+
+        self.assertEqual(output.buffers, [b"constant"])
+        self.assertEqual(output.pte_data, {"constant": DataEntry(0, 16, None)})
+        self.assertEqual(output.external_data, {})
+        externalize_pte_data.assert_not_called()
+
+    def test_external_constants_option_externalizes_constants(self) -> None:
+        output, externalize_pte_data = self._preprocess_named_data(
+            {"external_constants_max_data_bytes": 16}
+        )
+
+        self.assertEqual(output.buffers, [b"constant"])
+        self.assertEqual(output.pte_data, {})
+        self.assertEqual(len(output.external_data), 1)
+        self.assertEqual(list(next(iter(output.external_data.values()))), ["constant"])
+        externalize_pte_data.assert_called_once_with(16, "vulkan_constants")
+
     def test_unset_options_are_absent(self) -> None:
         round_tripped = self._round_trip({})
         self.assertNotIn("small_texture_limits", round_tripped)
         self.assertNotIn("skip_memory_planning", round_tripped)
+        self.assertNotIn("external_constants_max_data_bytes", round_tripped)
 
 
 if __name__ == "__main__":

backends/vulkan/vulkan_preprocess.py

@@ -90,6 +90,17 @@ def apply_passes(program: ExportedProgram, passes) -> ExportedProgram:
     return program
 
 
+def _parse_external_constants_max_data_bytes(value_bytes: bytes) -> int:
+    # CompileSpec values can bypass parse_compile_options, so validate this
+    # serialized boundary independently.
+    if len(value_bytes) != 8:
+        raise ValueError("external_constants_max_data_bytes must be encoded as uint64")
+    value = int.from_bytes(value_bytes, byteorder="little")
+    if value <= 0:
+        raise ValueError("external_constants_max_data_bytes must be a positive uint64")
+    return value
+
+
 def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]:
     options = {}
     for spec in compile_specs:
@@ -119,6 +130,9 @@ def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]:
         if spec.key == "skip_memory_planning":
             options[spec.key] = bool.from_bytes(spec.value, byteorder="little")
 
+        if spec.key == "external_constants_max_data_bytes":
+            options[spec.key] = _parse_external_constants_max_data_bytes(spec.value)
+
         # Unhandled options are ignored
 
     return options
@@ -246,6 +260,16 @@ def preprocess(  # noqa: C901
             force_fp16=force_fp16,
         )
         vk_graph = graph_builder.build_graph()
+        external_constants_max_data_bytes = compile_options.get(
+            "external_constants_max_data_bytes"
+        )
+        if external_constants_max_data_bytes is not None:
+            # VkGraphBuilder populates pte_data only from constant tensors;
+            # already-tagged named data remains in external_data.
+            graph_builder.named_data_store.externalize_pte_data(
+                external_constants_max_data_bytes,
+                "vulkan_constants",
+            )
 
         return PreprocessResult(
             processed_bytes=serialize_vulkan_graph(

exir/_serialize/_named_data_store.py

@@ -210,6 +210,66 @@ def add_named_data(
                 tensor_layout,
             )
 
+    def externalize_pte_data(
+        self,
+        max_data_bytes: int,
+        tag_prefix: str,
+    ) -> None:
+        # Keep this generic API defensive because callers can bypass backend
+        # option parsing and serialized compile-spec validation.
+        if (
+            isinstance(max_data_bytes, bool)
+            or not isinstance(max_data_bytes, int)
+            or max_data_bytes <= 0
+        ):
+            raise ValueError("external data shard cap must be a positive integer")
+        if not tag_prefix:
+            raise ValueError("external data tag prefix must be nonempty")
+
+        entries_by_buffer: Dict[int, Dict[str, DataEntry]] = {}
+        for key, entry in self.pte_data.items():
+            entries_by_buffer.setdefault(entry.buffer_index, {})[key] = entry
+
+        shards: List[Dict[str, DataEntry]] = []
+        current_shard: Dict[str, DataEntry] = {}
+        current_bytes = 0
+        # Prefer stable key order over size-based packing so equivalent stores
+        # always produce the same shards.
+        ordered_groups = sorted(
+            entries_by_buffer.items(), key=lambda item: tuple(sorted(item[1]))
+        )
+        for buffer_index, entries in ordered_groups:
+            buffer_size = len(self.buffers[buffer_index])
+            if buffer_size > max_data_bytes:
+                raise ValueError(
+                    f"buffer {buffer_index} has {buffer_size} bytes and exceeds "
+                    f"external data shard cap {max_data_bytes}"
+                )
+            if current_shard and current_bytes + buffer_size > max_data_bytes:
+                shards.append(current_shard)
+                current_shard = {}
+                current_bytes = 0
+            current_shard.update(entries)
+            current_bytes += buffer_size
+        if current_shard:
+            shards.append(current_shard)
+
+        external_data = {
+            tag: dict(entries) for tag, entries in self.external_data.items()
+        }
+        for entries in shards:
+            keys = sorted(entries)
+            digest = hashlib.sha256("\0".join(keys).encode("utf-8")).hexdigest()
+            tag = f"{tag_prefix}_{digest}"
+            canonical_entries = {key: entries[key] for key in keys}
+            existing = external_data.get(tag)
+            if existing is not None and existing != canonical_entries:
+                raise ValueError(f"external data tag collision for {tag}")
+            external_data[tag] = canonical_entries
+
+        self.external_data = external_data
+        self.pte_data = {}
+
     def get_named_data_store_output(self) -> NamedDataStoreOutput:
         # Clean up empty maps inside self.external_data
         self.external_data = {k: v for k, v in self.external_data.items() if len(v) > 0}