From eb4d8320057b0ec62e509756e310061dd5e37acd Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 5 Aug 2026 12:54:37 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../serialization/vulkan_graph_builder.py | 48 +++++- backends/vulkan/test/test_serialization.py | 138 +++++++++++++++++- .../test/test_vulkan_compile_options.py | 25 +++- backends/vulkan/vulkan_preprocess.py | 5 + 4 files changed, 209 insertions(+), 7 deletions(-) diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index ca5ab196dd2..46e01e701b1 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -9,7 +9,7 @@ import logging import operator from types import NoneType -from typing import cast, List, Optional, Union +from typing import cast, Dict, List, Optional, Union import executorch.backends.vulkan.serialization.vulkan_graph_schema as vk_graph_schema import torch @@ -28,6 +28,7 @@ ) from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.utils import DelegateMappingBuilder +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.tensor import TensorSpec from torch._export.utils import get_buffer, get_param, is_buffer, is_param from torch.export import ExportedProgram @@ -49,11 +50,41 @@ def __init__( delegate_mapping_builder: DelegateMappingBuilder, downcast_64_bit: bool = True, force_fp16: bool = False, + alias_buffer_mutations: bool = False, ) -> None: self.program = program self.delegate_mapping_builder = delegate_mapping_builder self.downcast_64_bit = downcast_64_bit self.force_fp16 = force_fp16 + self.buffer_mutation_inputs: Dict[str, Node] = {} + self.buffer_mutation_user_outputs: set[str] = set() + if alias_buffer_mutations: + nodes_by_name = { + node.name: node for node in program.graph_module.graph.nodes + } + buffer_inputs_by_target: Dict[str, Node] = {} + for name, target in program.graph_signature.inputs_to_buffers.items(): + if name not in nodes_by_name: + continue + buffer_input = nodes_by_name[name] + prepack = next( + ( + user + for user in buffer_input.users + if user.op == "call_function" + and user.target == exir_ops.edge.et_vk.prepack.default + ), + None, + ) + buffer_inputs_by_target[target] = prepack or buffer_input + self.buffer_mutation_inputs = { + output_name: buffer_inputs_by_target[target] + for output_name, target in program.graph_signature.buffers_to_mutate.items() + if target in buffer_inputs_by_target + } + self.buffer_mutation_user_outputs = set( + program.graph_signature.user_outputs + ) self.chain = [] self.values = [] self.input_ids = [] @@ -160,6 +191,16 @@ def maybe_add_constant_tensor(self, node: Node) -> int: return constant_id def create_node_value(self, node: Node) -> int: + if node.name in self.buffer_mutation_inputs: + input_node = self.buffer_mutation_inputs[node.name] + if input_node not in self.node_to_value_ids: + raise AssertionError( + "Cannot alias a buffer mutation before its input is serialized" + ) + value_id = self.node_to_value_ids[input_node] + self.node_to_value_ids[node] = value_id + return value_id + # If the node has been marked as a scalar tensor, create a SymInt instead of a tensor if is_symint_node(node) or node.meta.get("etvk_is_scalar_tensor", False): new_id = self.create_symint_value() @@ -448,7 +489,10 @@ def process_output_node(self, node: Node) -> None: ) # Mutable buffers outputs are not included as an output to the # delegate call. Skip marking them as an output. - if is_mutable_buffer_node(out_node, self.program): + if out_node.name in self.buffer_mutation_inputs: + if out_node.name not in self.buffer_mutation_user_outputs: + continue + elif is_mutable_buffer_node(out_node, self.program): continue self.output_ids.append(self.node_to_value_ids[out_node]) diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index c373f5216d2..71a6980635a 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -9,9 +9,14 @@ import ctypes import random import unittest -from typing import List +from types import SimpleNamespace +from typing import List, Tuple +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 import torch +from executorch.backends.vulkan.serialization import ( + vulkan_graph_builder as graph_builder_module, +) from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( IntList, @@ -30,6 +35,137 @@ class TestSerialization(unittest.TestCase): + def _build_mutation_program( + self, prepack: bool, shared_user_output: bool = False + ) -> Tuple[SimpleNamespace, torch.fx.Node, torch.fx.Node, torch.fx.Node]: + graph = torch.fx.Graph() + state = graph.placeholder("state") + user_input = graph.placeholder("user_input") + state.meta["spec"] = graph_builder_module.TensorSpec.from_tensor(torch.zeros(4)) + user_input.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + + state_value = state + if prepack: + state_value = graph.call_function( + graph_builder_module.exir_ops.edge.et_vk.prepack.default, + (state,), + ) + state_value.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.zeros(4) + ) + + mutation = graph.call_function( + torch.ops.aten.add.Tensor, (state_value, user_input) + ) + mutation.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + user_output = mutation + if not shared_user_output: + user_output = graph.call_function( + torch.ops.aten.mul.Tensor, (user_input, 2.0) + ) + user_output.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + graph.output((mutation, user_output)) + + graph_module = torch.fx.GraphModule({}, graph) + signature = SimpleNamespace( + buffers_to_mutate={mutation.name: "state"}, + inputs_to_buffers={state.name: "state"}, + inputs_to_lifted_tensor_constants={}, + inputs_to_parameters={}, + non_persistent_buffers=set(), + user_outputs=(user_output.name,), + ) + program = SimpleNamespace( + constants={}, + graph_module=graph_module, + graph_signature=signature, + state_dict={"state": torch.zeros(4)}, + ) + return program, state_value, mutation, user_output + + def test_alias_buffer_mutations_is_opt_in(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, user_output = ( + self._build_mutation_program(prepack) + ) + + default_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + ) + default_graph = default_builder.build_graph() + self.assertNotEqual( + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + default_graph.output_ids, + [ + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[user_output], + ], + ) + + explicit_false_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=False, + ) + self.assertEqual(default_graph, explicit_false_builder.build_graph()) + + aliasing_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + aliasing_graph = aliasing_builder.build_graph() + self.assertEqual( + aliasing_builder.node_to_value_ids[mutation], + aliasing_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + aliasing_graph.output_ids, + [aliasing_builder.node_to_value_ids[user_output]], + ) + + def test_alias_buffer_mutations_preserves_shared_user_output(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, _ = self._build_mutation_program( + prepack, shared_user_output=True + ) + builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + + graph = builder.build_graph() + + self.assertEqual( + builder.node_to_value_ids[mutation], + builder.node_to_value_ids[state_value], + ) + self.assertEqual( + graph.output_ids, + [builder.node_to_value_ids[mutation]], + ) + def _generate_random_const_tensors(self, num_tensors: int) -> List[torch.Tensor]: """ Helper function to generate `num_tensor` buffers of random sizes and random contents, diff --git a/backends/vulkan/test/test_vulkan_compile_options.py b/backends/vulkan/test/test_vulkan_compile_options.py index f45cfdf12d9..ff507bb3197 100644 --- a/backends/vulkan/test/test_vulkan_compile_options.py +++ b/backends/vulkan/test/test_vulkan_compile_options.py @@ -41,6 +41,10 @@ def test_skip_memory_planning_round_trips(self) -> None: round_tripped = self._round_trip({"skip_memory_planning": True}) self.assertTrue(round_tripped.get("skip_memory_planning")) + def test_alias_buffer_mutations_round_trips(self) -> None: + round_tripped = self._round_trip({"alias_buffer_mutations": True}) + self.assertTrue(round_tripped.get("alias_buffer_mutations")) + def test_force_fp16_round_trips(self) -> None: round_tripped = self._round_trip({"force_fp16": True}) self.assertTrue(round_tripped.get("force_fp16")) @@ -105,15 +109,15 @@ def build_graph(): ), patch( "executorch.backends.vulkan.vulkan_preprocess.VkGraphBuilder", return_value=graph_builder, - ), patch( + ) as graph_builder_factory, 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 + return result.data_store_output, externalize_pte_data, graph_builder_factory def test_external_constants_default_keeps_constants_inline(self) -> None: - output, externalize_pte_data = self._preprocess_named_data({}) + output, externalize_pte_data, _ = self._preprocess_named_data({}) self.assertEqual(output.buffers, [b"constant"]) self.assertEqual(output.pte_data, {"constant": DataEntry(0, 16, None)}) @@ -121,7 +125,7 @@ def test_external_constants_default_keeps_constants_inline(self) -> None: externalize_pte_data.assert_not_called() def test_external_constants_option_externalizes_constants(self) -> None: - output, externalize_pte_data = self._preprocess_named_data( + output, externalize_pte_data, _ = self._preprocess_named_data( {"external_constants_max_data_bytes": 16} ) @@ -131,8 +135,21 @@ def test_external_constants_option_externalizes_constants(self) -> None: self.assertEqual(list(next(iter(output.external_data.values()))), ["constant"]) externalize_pte_data.assert_called_once_with(16, "vulkan_constants") + def test_alias_buffer_mutations_reaches_graph_builder(self) -> None: + for options, expected in ( + ({}, False), + ({"alias_buffer_mutations": True}, True), + ): + with self.subTest(options=options): + _, _, graph_builder_factory = self._preprocess_named_data(options) + self.assertIs( + graph_builder_factory.call_args.kwargs["alias_buffer_mutations"], + expected, + ) + def test_unset_options_are_absent(self) -> None: round_tripped = self._round_trip({}) + self.assertNotIn("alias_buffer_mutations", round_tripped) self.assertNotIn("small_texture_limits", round_tripped) self.assertNotIn("skip_memory_planning", round_tripped) self.assertNotIn("external_constants_max_data_bytes", round_tripped) diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index d3954240880..f7d6955ce26 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -130,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 == "alias_buffer_mutations": + 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) @@ -172,6 +175,7 @@ def preprocess( # noqa: C901 ) downcast_64_bit = compile_options.get("downcast_64_bit", True) force_fp16 = compile_options.get("force_fp16", False) + alias_buffer_mutations = compile_options.get("alias_buffer_mutations", False) program = unsafe_remove_auto_functionalized_pass(program) @@ -258,6 +262,7 @@ def preprocess( # noqa: C901 DelegateMappingBuilder(generated_identifiers=True), downcast_64_bit=downcast_64_bit, force_fp16=force_fp16, + alias_buffer_mutations=alias_buffer_mutations, ) vk_graph = graph_builder.build_graph() external_constants_max_data_bytes = compile_options.get(