diff --git a/backends/vulkan/partitioner/vulkan_partitioner.py b/backends/vulkan/partitioner/vulkan_partitioner.py index fb51a0edfad..298581ebef7 100644 --- a/backends/vulkan/partitioner/vulkan_partitioner.py +++ b/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)) diff --git a/backends/vulkan/test/TARGETS b/backends/vulkan/test/TARGETS index 7517f7d66f3..196c41d3ec3 100644 --- a/backends/vulkan/test/TARGETS +++ b/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 = [ diff --git a/backends/vulkan/test/test_vulkan_compile_options.py b/backends/vulkan/test/test_vulkan_compile_options.py index f44850d2915..f45cfdf12d9 100644 --- a/backends/vulkan/test/test_vulkan_compile_options.py +++ b/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__": diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index 34c68fe5828..d3954240880 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/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( diff --git a/exir/_serialize/_named_data_store.py b/exir/_serialize/_named_data_store.py index c2c67b4bc28..d0b18f9d6c2 100644 --- a/exir/_serialize/_named_data_store.py +++ b/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} diff --git a/exir/_serialize/test/test_named_data_store.py b/exir/_serialize/test/test_named_data_store.py index a9e32afd0f1..929725dda6d 100644 --- a/exir/_serialize/test/test_named_data_store.py +++ b/exir/_serialize/test/test_named_data_store.py @@ -6,7 +6,10 @@ # pyre-strict +import copy +import hashlib import unittest +from typing import Any, cast import torch @@ -16,7 +19,157 @@ from executorch.exir.tensor_layout import TensorLayout +class _SizedBuffer: + def __init__(self, size: int) -> None: + self.size = size + + def __len__(self) -> int: + return self.size + + class TestNamedDataStore(unittest.TestCase): + def test_externalize_pte_data_counts_aliased_buffer_once(self) -> None: + store = NamedDataStore() + layout = TensorLayout(ScalarType.FLOAT, [1], [0]) + store.add_named_data("key_a", b"aaaaaa", 16, None, layout) + store.add_named_data("key_a_alias", b"aaaaaa", 32, None, layout) + store.add_named_data("key_b", b"bbbb", 16, None, layout) + expected_buffers = list(store.buffers) + expected_entries = copy.deepcopy(store.pte_data) + + store.externalize_pte_data(10, "test_constants") + output = store.get_named_data_store_output() + + self.assertEqual(output.buffers, expected_buffers) + self.assertEqual(output.pte_data, {}) + self.assertEqual(len(output.external_data), 1) + self.assertEqual(next(iter(output.external_data.values())), expected_entries) + + def test_externalize_pte_data_rollover_is_insertion_order_independent( + self, + ) -> None: + def externalize( + order: list[str], + ) -> dict[str, list[tuple[str, bytes, int, TensorLayout | None]]]: + data = { + "key_a": (b"a" * 6, 16), + "key_a_alias": (b"a" * 6, 32), + "key_b": (b"b" * 4, 32), + "key_c": (b"c" * 5, 64), + } + store = NamedDataStore() + for key in order: + payload, alignment = data[key] + store.add_named_data(key, payload, alignment) + store.externalize_pte_data(10, "test_constants") + output = store.get_named_data_store_output() + return { + tag: [ + ( + key, + output.buffers[entry.buffer_index], + entry.alignment, + entry.tensor_layout, + ) + for key, entry in entries.items() + ] + for tag, entries in output.external_data.items() + } + + forward = externalize(["key_a", "key_a_alias", "key_b", "key_c"]) + reverse = externalize(["key_c", "key_b", "key_a_alias", "key_a"]) + + self.assertEqual(len(forward), 2) + self.assertEqual(forward, reverse) + + def test_externalize_pte_data_rejects_oversized_buffer_atomically(self) -> None: + store = NamedDataStore() + store.add_named_data("external", b"ext", 8, "existing") + store.add_named_data("key_a", b"aaaa", 16) + store.add_named_data("key_z", b"z" * 11, 32) + before = ( + list(store.buffers), + copy.deepcopy(store.pte_data), + copy.deepcopy(store.external_data), + dict(store.key_to_buffer_idx), + ) + + with self.assertRaisesRegex(ValueError, "exceeds external data shard cap"): + store.externalize_pte_data(10, "test_constants") + + after = ( + list(store.buffers), + store.pte_data, + store.external_data, + store.key_to_buffer_idx, + ) + self.assertEqual(after, before) + + def test_externalize_pte_data_accepts_production_max_buffer_metadata( + self, + ) -> None: + store = NamedDataStore() + store.buffers = [cast(bytes, _SizedBuffer(1_174_405_120))] + store.pte_data = {"largest": DataEntry(0, 16, None)} + store.key_to_buffer_idx = {"largest": 0} + + store.externalize_pte_data(1_500_000_000, "vulkan_constants") + + output = store.get_named_data_store_output() + self.assertEqual(output.pte_data, {}) + self.assertEqual( + [list(entries) for entries in output.external_data.values()], + [["largest"]], + ) + + def test_externalize_pte_data_rejects_invalid_arguments(self) -> None: + store = NamedDataStore() + invalid_caps: list[Any] = [True, 0, -1, 1.5, "10"] + for cap in invalid_caps: + with self.subTest(cap=cap), self.assertRaisesRegex( + ValueError, "positive integer" + ): + store.externalize_pte_data(cap, "test_constants") + + with self.assertRaisesRegex(ValueError, "prefix must be nonempty"): + store.externalize_pte_data(10, "") + + def test_externalize_pte_data_rejects_tag_collision_atomically(self) -> None: + store = NamedDataStore() + key = "inline" + digest = hashlib.sha256(key.encode("utf-8")).hexdigest() + tag = f"test_constants_{digest}" + store.add_named_data("external", b"ext", 8, tag) + store.add_named_data(key, b"inline", 16) + before = ( + list(store.buffers), + copy.deepcopy(store.pte_data), + copy.deepcopy(store.external_data), + ) + + with self.assertRaisesRegex(ValueError, "external data tag collision"): + store.externalize_pte_data(10, "test_constants") + + self.assertEqual( + (list(store.buffers), store.pte_data, store.external_data), before + ) + + def test_externalize_pte_data_preserves_existing_external_data_and_is_idempotent( + self, + ) -> None: + store = NamedDataStore() + store.add_named_data("external", b"ext", 8, "existing") + store.add_named_data("inline", b"inline", 16) + + store.externalize_pte_data(10, "test_constants") + first = copy.deepcopy(store.get_named_data_store_output()) + store.externalize_pte_data(10, "test_constants") + second = store.get_named_data_store_output() + + self.assertEqual(first, second) + self.assertEqual(list(second.external_data["existing"]), ["external"]) + self.assertEqual(len(second.external_data), 2) + def test_add(self) -> None: store = NamedDataStore() store.add_named_data("key1", b"data1", None, None)