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

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].

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

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