diff --git a/backends/openvino/partitioner.py b/backends/openvino/partitioner.py index 5b01995daa5..4d6cfd2f93e 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,42 @@ 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 _remove_copies_of_inputs(partition: Partition) -> 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. + """ + + def copies_input(node: torch.fx.Node) -> bool: + while (source := _copy_source(node)) is not None: + if source not in partition.nodes: + return True + 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, @@ -252,6 +295,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: partition_tags = {} for partition in partition_list: + _remove_copies_of_inputs(partition) for node in partition.nodes: tag = f"tag{partition.id}" node.meta["delegation_tag"] = tag diff --git a/backends/openvino/tests/ops/base_openvino_op_test.py b/backends/openvino/tests/ops/base_openvino_op_test.py index 4e20b41cbf6..df921c4c4ac 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 diff --git a/backends/openvino/tests/ops/test_copy.py b/backends/openvino/tests/ops/test_copy.py new file mode 100644 index 00000000000..c65b1d35272 --- /dev/null +++ b/backends/openvino/tests/ops/test_copy.py @@ -0,0 +1,44 @@ +import torch +from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( + BaseOpenvinoOpTest, +) +from executorch.exir import EdgeCompileConfig + + +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),))