diff --git a/backends/openvino/partitioner.py b/backends/openvino/partitioner.py index 5b01995daa5..4941c2b9548 100644 --- a/backends/openvino/partitioner.py +++ b/backends/openvino/partitioner.py @@ -22,9 +22,16 @@ ) from torch.export.exported_program import ExportedProgram -from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner +from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition from torch.fx.passes.operator_support import OperatorSupportBase +# OpenVINO converts these to nothing, so their output is their input. +_COPY_OPS = { + "aten.alias_copy.default", + "aten.clone.default", + "dim_order_ops._clone_dim_order.default", +} + class PatternNode: op_types: dict[str, Optional[list]] = {} @@ -41,6 +48,67 @@ def _reads_only_bool(node: torch.fx.Node) -> bool: return bool(tensors) and all(t.dtype == torch.bool for t in tensors) +def _copy_source(node: torch.fx.Node) -> Optional[torch.fx.Node]: + """The node's input, if OpenVINO returns that input as the node's output.""" + if isinstance(node.target, str) or node.target.__name__ not in _COPY_OPS: + return None + source = node.args[0] + return source if isinstance(source, torch.fx.Node) else None + + +def _constants(program: ExportedProgram) -> set[torch.fx.Node]: + """The placeholders that each partition reading them stores as its own + data (see tag_constant_data), so OpenVINO gets them as constants, not as + inputs.""" + signature = program.graph_signature + mutated = set(signature.buffers_to_mutate.values()) + names = { + *signature.inputs_to_parameters, + *signature.inputs_to_lifted_tensor_constants, + *( + name + for name, buffer in signature.inputs_to_buffers.items() + if buffer not in mutated + ), + } + return { + node + for node in program.graph.nodes + if node.op == "placeholder" and node.name in names + } + + +def _remove_copies_of_inputs( + partition: Partition, constants: set[torch.fx.Node] +) -> None: + """Leave out of the partition each copy that would return one of the + partition's inputs, unchanged, as one of its outputs. + + OpenVINO would return the input itself, and when the runtime binds the + output buffers after the inputs, its CPU plugin does not write such an + output and can overwrite the input with the output buffer's old contents. + Copies used only inside the partition stay, because OpenVINO reads every + input in the default dim order, which a copy left out could have changed. + Copies of constants stay too, since constants are not partition inputs. + """ + + def copies_input(node: torch.fx.Node) -> bool: + while (source := _copy_source(node)) is not None: + if source not in partition.nodes: + return source not in constants + node = source + return False + + # Taking out a copy can make the copy it reads an output. + while leaving := [ + node + for node in partition.nodes + if copies_input(node) and any(u not in partition.nodes for u in node.users) + ]: + for node in leaving: + partition.remove_node(node) + + class OpenvinoOperatorsSupport(OperatorSupportBase): extended_support_dict = { "torch.ops.dim_order_ops._clone_dim_order.default": None, @@ -251,7 +319,9 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: partition_list = partitioner.propose_partitions() partition_tags = {} + constants = _constants(exported_program) for partition in partition_list: + _remove_copies_of_inputs(partition, constants) for node in partition.nodes: tag = f"tag{partition.id}" node.meta["delegation_tag"] = tag diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index 809993f960a..677e9797d03 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -22,7 +22,7 @@ ) from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.passes.memory_format_ops_pass import DimOrderOpsRevertPass -from openvino import compile_model, opset13 # type: ignore[import-untyped] +from openvino import compile_model, opset13, opset16 # type: ignore[import-untyped] from openvino.frontend import FrontEndManager # type: ignore[import-untyped] from openvino.frontend.pytorch.fx_decoder import ( # type: ignore[import-untyped] TorchFXPythonDecoder, @@ -63,6 +63,21 @@ def _match_partition_inputs( ) +def _copy_returned_inputs(model) -> None: + """Make the converted model compute each output that is one of its inputs. + + The converter turns copies like aten.clone into nothing, so such an output + is the input itself. The CPU plugin does not write it, and can overwrite + the input with the output buffer's old contents. OpenVINO keeps an + Identity between an input and an output, and the CPU plugin runs it as a + copy, while it can remove a Convert to the same type. + """ + for result in model.get_results(): + source = result.input_value(0) + if source.get_node().get_type_name() == "Parameter": + result.input(0).replace_source_output(opset16.identity(source).output(0)) + + @final class OpenvinoBackend(BackendDetails): @@ -142,6 +157,7 @@ def preprocess( frontend = FrontEndManager().load_by_framework("pytorch") model = frontend.convert(frontend.load(decoder)) _match_partition_inputs(model, module, input_shapes, input_types) + _copy_returned_inputs(model) ppp = PrePostProcessor(model) # The converter takes an int input, like a SymInt size, as a scalar, but # the runtime passes it as an int64 tensor of shape [1]. diff --git a/backends/openvino/tests/ops/base_openvino_op_test.py b/backends/openvino/tests/ops/base_openvino_op_test.py index 4e20b41cbf6..c7a00bdcbf6 100644 --- a/backends/openvino/tests/ops/base_openvino_op_test.py +++ b/backends/openvino/tests/ops/base_openvino_op_test.py @@ -1,12 +1,17 @@ import os import tempfile import unittest +from typing import Optional import executorch import torch from executorch.backends.openvino.partitioner import OpenvinoPartitioner from executorch.backends.openvino.preprocess import OpenvinoBackend -from executorch.exir import EdgeProgramManager, to_edge_transform_and_lower +from executorch.exir import ( + EdgeCompileConfig, + EdgeProgramManager, + to_edge_transform_and_lower, +) from executorch.exir.backend.backend_details import CompileSpec from executorch.runtime import Runtime @@ -27,6 +32,7 @@ def execute_layer_test( assert_output_equal: bool = True, compile_specs: tuple[CompileSpec, ...] = (), more_inputs: tuple[tuple[torch.Tensor, ...], ...] = (), + compile_config: Optional[EdgeCompileConfig] = None, ): module = module.eval() @@ -40,6 +46,7 @@ def execute_layer_test( partitioner=[ OpenvinoPartitioner(compile_spec), ], + compile_config=compile_config, ) # Apply backend-specific passes @@ -98,3 +105,5 @@ def execute_layer_test( ), msg=f"ref_output:\n{ref_output[i]}\n\ntest_output:\n{outputs[i]}", ) + + return exec_prog diff --git a/backends/openvino/tests/ops/test_copy.py b/backends/openvino/tests/ops/test_copy.py new file mode 100644 index 00000000000..a9b0060c2ff --- /dev/null +++ b/backends/openvino/tests/ops/test_copy.py @@ -0,0 +1,118 @@ +import torch +from executorch.backends.openvino.preprocess import OpenvinoBackend +from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( + BaseOpenvinoOpTest, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.backend.backend_api import to_backend +from executorch.exir.backend.backend_details import CompileSpec +from executorch.runtime import Runtime +from torch.export import export + + +class TestCopyOperator(BaseOpenvinoOpTest): + + def test_copy_of_input_as_output(self): + class CopyAndSin(torch.nn.Module): + def __init__(self, copy): + super().__init__() + self.copy = copy + + def forward(self, x): + copy = self.copy(x) + return copy, torch.sin(copy) + + copies = { + "clone": lambda x: x.clone(), + "full slice": lambda x: x[:, :], + "clone of a clone": lambda x: x.clone().clone(), + } + inputs = (torch.randn(2, 3, 4, 5),) + for name, copy in copies.items(): + with self.subTest(copy=name): + self.execute_layer_test(CopyAndSin(copy), inputs) + + # LLM exports turn dim order off, which makes a clone aten.clone. + self.execute_layer_test( + CopyAndSin(lambda x: x.clone()), + inputs, + compile_config=EdgeCompileConfig(_skip_dim_order=True), + ) + + def test_copy_of_input_used_inside(self): + class ChannelsLastMean(torch.nn.Module): + def forward(self, x): + return x.clone(memory_format=torch.channels_last).mean(dim=(2, 3)) + + # Run outside, the clone would hand OpenVINO a channels-last input, + # which it reads as if it were contiguous. + self.execute_layer_test(ChannelsLastMean(), (torch.randn(2, 3, 4, 5),)) + + def test_copy_of_constant_as_output(self): + class CopyOfConstant(torch.nn.Module): + def __init__(self, kind): + super().__init__() + value = torch.randn(2, 3, 4, 5) + if kind == "parameter": + self.constant = torch.nn.Parameter(value) + elif kind == "buffer": + self.register_buffer("constant", value) + else: + self.constant = value + + def forward(self, x): + return self.constant.clone(), x + self.constant + + for kind in ("parameter", "buffer", "lifted tensor"): + with self.subTest(kind=kind): + program = self.execute_layer_test( + CopyOfConstant(kind), (torch.randn(2, 3, 4, 5),) + ) + # OpenVINO gets a constant as its own data, not as an input, so + # the copy stays. Run outside, it would store the constant twice. + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.operators, []) + + def test_copy_of_mutated_buffer_as_output(self): + class CopyOfMutatedBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("buffer", torch.randn(2, 3, 4, 5)) + + def forward(self, x): + copy = self.buffer.clone() + self.buffer.add_(x) + return copy, torch.sin(copy) + + # The program does not keep the starting value of a mutated buffer, so + # there is no reference output to compare with. + program = self.execute_layer_test( + CopyOfMutatedBuffer(), + (torch.randn(2, 3, 4, 5),), + assert_output_equal=False, + ) + # A mutated buffer is passed to the partition as an input, so its copy + # must run outside. + operators = program.executorch_program.execution_plan[0].operators + self.assertIn("dim_order_ops::_clone_dim_order", [op.name for op in operators]) + + def test_whole_program_returning_copy_of_input(self): + class CopyAndSin(torch.nn.Module): + def forward(self, x): + copy = x.clone() + return copy, torch.sin(copy) + + x = torch.randn(2, 3, 4, 5) + # Without a partitioner the copy cannot run on the CPU, so OpenVINO + # has to make it. + lowered = to_backend( + OpenvinoBackend.__name__, + to_edge(export(CopyAndSin(), (x,))).exported_program(), + [CompileSpec("device", self.device.encode())], + ) + program = Runtime.get().load_program(lowered.buffer()) + method = program.load_method("forward") + assert method is not None + torch.testing.assert_close( + method.execute((x,)), [x, torch.sin(x)], atol=self.atol, rtol=self.rtol + )