[ExecuTorch][Vulkan] Preserve persistent buffer mutations

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])

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,

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,23 +109,23 @@ 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)})
         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(
+        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)

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(