OpenVINO: fix wrong outputs when a partition returns a copy of its input

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

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

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