diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index 46e01e701b1..e5dfde9d865 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -418,15 +418,26 @@ def get_or_create_value_for(self, arg: _Argument): raise RuntimeError(f"Cannot create value for arg of type {type(arg)}") def process_placeholder_node(self, node: Node) -> None: - # ignores any tensors that don't get used in any ops - if len(node.users) == 0: + # A non-param placeholder occupies a slot in the delegate call's + # argument list whether or not this graph goes on to use it, and + # VulkanBackend::execute matches `args` to graph inputs positionally. + # Dropping an unused one from input_ids desynchronises the two, and the + # runtime then rejects the call because it was handed more arguments + # than the graph declares inputs and outputs. That happens in practice + # when a placeholder's only consumers are folded away by the passes + # that run after partitioning, so the graph the partitioner tagged and + # the graph serialized here disagree about which inputs are live. + if is_param_node(self.program, node): + # Params are serialized into the blob rather than passed at call + # time, so an unused one costs nothing to skip. + if len(node.users) > 0: + self.create_node_value(node) return None ids = self.create_node_value(node) - if not is_param_node(self.program, node): - if isinstance(ids, int): - self.input_ids.append(ids) - else: - self.input_ids += ids + if isinstance(ids, int): + self.input_ids.append(ids) + else: + self.input_ids += ids def process_getitem_node(self, node: Node) -> None: # Find ValueList id from the collection node. diff --git a/backends/vulkan/test/test_vulkan_graph_builder.py b/backends/vulkan/test/test_vulkan_graph_builder.py new file mode 100644 index 00000000000..346f67b7612 --- /dev/null +++ b/backends/vulkan/test/test_vulkan_graph_builder.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.vulkan.serialization.vulkan_graph_builder import ( + VkGraphBuilder, +) +from executorch.backends.vulkan.vulkan_preprocess import apply_passes +from executorch.exir import to_edge +from executorch.exir.backend.utils import DelegateMappingBuilder +from executorch.exir.passes import SpecPropPass + + +class TestVkGraphBuilderInputIds(unittest.TestCase): + """The serialized input list has to match the delegate call's arguments. + + VulkanBackend::execute walks `args` positionally against + ComputeGraph::inputs() and rejects the call when the counts disagree, so + every placeholder that the delegate call passes must appear in input_ids, + including ones this graph happens not to use. Unused placeholders are not + hypothetical: passes that run after partitioning can fold away a + placeholder's only consumers, leaving the argument list and the serialized + graph out of step. + """ + + def _build(self, module: torch.nn.Module, inputs) -> VkGraphBuilder: + edge = to_edge(torch.export.export(module, inputs, strict=True)) + # The builder reads node specs, which the backend's own preprocess + # populates before it gets here. + program = apply_passes(edge.exported_program(), [SpecPropPass()]) + builder = VkGraphBuilder( + program, DelegateMappingBuilder(generated_identifiers=True) + ) + builder.build_graph() + return builder + + def test_unused_placeholder_is_still_declared_as_an_input(self) -> None: + class UsesOnlyTheFirstInput(torch.nn.Module): + def forward(self, used, unused): + return used + used + + builder = self._build( + UsesOnlyTheFirstInput(), (torch.randn(2, 3), torch.randn(2, 3)) + ) + self.assertEqual(len(builder.input_ids), 2) + + def test_used_placeholders_are_declared_in_order(self) -> None: + class UsesBothInputs(torch.nn.Module): + def forward(self, first, second): + return first + second + + builder = self._build(UsesBothInputs(), (torch.randn(2, 3), torch.randn(2, 3))) + self.assertEqual(len(builder.input_ids), 2) + + +if __name__ == "__main__": + unittest.main()