From e71822f964554893af4a19a7a7886e9fa1ead207 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 12:05:49 -0800 Subject: [PATCH] OpenVINO: fix wrong outputs when a partition returns a copy of its input When an OpenVINO partition returned a copy of one of its inputs, for example `x.clone()` or `x[:, :]`, the copy came back with stale memory, and so did anything else the partition computed from that input. OpenVINO converts these copies to nothing, so the converted model returns its input directly. ExecuTorch's runtime binds the output buffers after the inputs, and in that order OpenVINO's CPU plugin does not write such an output and can overwrite the input with the output buffer's old contents. Every OpenVINO release tried behaves this way, from 2025.2 to the current nightly. Binding the outputs first avoids it, but only because of an undocumented detail of the plugin. The partitioner now leaves such a copy (aten.clone, dim_order_ops._clone_dim_order or aten.alias_copy) out of the partition when its result leaves the partition, so it runs on the CPU kernels. Copies used only inside the partition stay there, since OpenVINO skips them anyway. The OpenVINO op tests can now pass an EdgeCompileConfig, to also cover lowering with dim order turned off, which gives aten.clone. --- backends/openvino/partitioner.py | 46 ++++++++++++++++++- .../tests/ops/base_openvino_op_test.py | 9 +++- backends/openvino/tests/ops/test_copy.py | 44 ++++++++++++++++++ 3 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 backends/openvino/tests/ops/test_copy.py 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),))