From 4b224395c0efc3910f4b1ad7e30d0123991d9e0a Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 23 Oct 2026 02:19:29 -0700
Subject: [PATCH 01/26] Normalize the dim order of parameters before
 delegation, like constants

The emitter stores a constant tensor that is neither contiguous nor
channels-last as contiguous. convert_constant_dim_order_pass makes lifted
constants and buffers contiguous in to_edge, before partitioning, so a
delegate is lowered with the layout the program will store. Parameters
were left out.

So a non-contiguous parameter that a delegate reads as an input, such as a
transposed weight left outside the delegate, was lowered with its original
strides but stored contiguous. A delegate that compiles those strides in,
like the CUDA backend, would read the wrong elements from a copy of the
stored bytes.

The pass now converts such parameters too, keeping them nn.Parameter.
---
 .../passes/convert_constant_dim_order_pass.py | 13 +++-
 exir/program/_program.py                      | 11 ++-
 exir/tests/test_passes.py                     | 70 +++++++++++++++++++
 3 files changed, 88 insertions(+), 6 deletions(-)

diff --git a/exir/passes/convert_constant_dim_order_pass.py b/exir/passes/convert_constant_dim_order_pass.py
index dd98fc9ece6..9cc00a66498 100644
--- a/exir/passes/convert_constant_dim_order_pass.py
+++ b/exir/passes/convert_constant_dim_order_pass.py
@@ -54,9 +54,9 @@ def convert_constant_dim_order_pass(
     """
     Normalize lifted tensors to default or channels-last dim order.
 
-    Tensor constants and buffers with other dim orders or striding are converted
-    to contiguous tensors. This pass acts in-place on the unlifted exported
-    program.
+    Tensor constants, buffers and parameters with other dim orders or striding
+    are converted to contiguous tensors. This pass acts in-place on the unlifted
+    exported program.
 
     Args:
         exported_program: The ExportedProgram to transform.
@@ -84,6 +84,13 @@ def convert_constant_dim_order_pass(
             )
             _update_placeholder_meta(exported_program, key, kind)
 
+    for key, param in exported_program.named_parameters():
+        if _should_transform(param):
+            exported_program.state_dict[key] = torch.nn.Parameter(
+                param.contiguous(), requires_grad=param.requires_grad
+            )
+            _update_placeholder_meta(exported_program, key, InputKind.PARAMETER)
+
     # Convert buffers.
     for key, buffer in exported_program.named_buffers():
         if (
diff --git a/exir/program/_program.py b/exir/program/_program.py
index ed55b345362..f2084235bda 100644
--- a/exir/program/_program.py
+++ b/exir/program/_program.py
@@ -823,8 +823,10 @@ def _generate_edge_program(
     program = remove_unused_parameters_pass(program)
 
     if not config._skip_dim_order:
-        # Normalize constants and buffers before MemoryFormatOpsPass rewrites
-        # preserve-format copies, which cannot represent arbitrary dim orders.
+        # Normalize constants, buffers and parameters before partitioning, so a
+        # delegate sees the layout the program stores, and before
+        # MemoryFormatOpsPass rewrites preserve-format copies, which cannot
+        # represent arbitrary dim orders.
         program = convert_constant_dim_order_pass.convert_constant_dim_order_pass(
             program
         )
@@ -870,7 +872,10 @@ def _generate_edge_program(
     # Lift the tensor constants created in ScalarToTensorPass
     edge_program = lift_constant_tensor_pass(edge_program)
 
-    # Normalize tensor constants introduced by the preceding passes.
+    # Normalize tensor constants introduced by the preceding passes. With
+    # _skip_dim_order this is the only call, so it fixes the layout of
+    # parameters, buffers and constants, but ops computed from them keep the
+    # old strides.
     edge_program = convert_constant_dim_order_pass.convert_constant_dim_order_pass(
         edge_program
     )
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 1afaece0f3a..1a7ba243950 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -3115,6 +3115,76 @@ def test_convert_constant_dim_order_to_contiguous(self):
                         f"Constant should be contiguous after pass, got strides {modified_const.stride()}",
                     )
 
+    def test_convert_constant_dim_order_parameter(self):
+        """
+        Verify that the convert_constant_dim_order_pass converts a parameter
+        outside the allowed dim orders to contiguous, keeps it a parameter, and
+        leaves a channels-last parameter as it is.
+        """
+
+        class AddParameter(torch.nn.Module):
+            def __init__(self, weight: torch.Tensor):
+                super().__init__()
+                self.weight = torch.nn.Parameter(weight)
+
+            def forward(self, x):
+                return x + self.weight
+
+        weights = {
+            "transposed": (torch.randn(5, 4).t(), True),
+            "channels_last": (
+                torch.randn(2, 3, 4, 5).to(memory_format=torch.channels_last),
+                False,
+            ),
+        }
+        for description, (weight, converted) in weights.items():
+            with self.subTest(description=description):
+                inputs = (torch.randn(weight.shape),)
+                ep = torch.export.export(AddParameter(weight), inputs)
+                modified_ep = (
+                    convert_constant_dim_order_pass.convert_constant_dim_order_pass(
+                        copy.deepcopy(ep)
+                    )
+                )
+
+                torch.testing.assert_close(
+                    ep.module()(*inputs), modified_ep.module()(*inputs)
+                )
+                param = modified_ep.state_dict["weight"]
+                self.assertIsInstance(param, torch.nn.Parameter)
+                self.assertEqual(param.is_contiguous(), converted)
+                placeholder = next(
+                    n for n in modified_ep.graph.nodes if n.name == "p_weight"
+                )
+                self.assertEqual(placeholder.meta["val"].stride(), param.stride())
+
+    def test_to_edge_makes_transposed_parameter_contiguous(self):
+        """
+        Verify that to_edge stores a transposed parameter contiguous with and
+        without _skip_dim_order, so a partitioner sees the stored layout.
+        """
+
+        class AddParameter(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.randn(5, 4).t())
+
+            def forward(self, x):
+                return self.weight + x
+
+        model = AddParameter()
+        inputs = (torch.randn(4, 5),)
+        for skip_dim_order in (False, True):
+            with self.subTest(skip_dim_order=skip_dim_order):
+                ep = to_edge(
+                    torch.export.export(model, inputs),
+                    compile_config=EdgeCompileConfig(_skip_dim_order=skip_dim_order),
+                ).exported_program()
+                self.assertTrue(ep.state_dict["weight"].is_contiguous())
+                placeholder = next(n for n in ep.graph.nodes if n.name == "p_weight")
+                self.assertTrue(placeholder.meta["val"].is_contiguous())
+                torch.testing.assert_close(ep.module()(*inputs), model(*inputs))
+
 
 class TestMemoryFormatOpsPassPreserveFormat(unittest.TestCase):
     """

From 6a8207dc60b7ca375d56f975cbc28b0df960fe67 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 23 Oct 2026 13:57:07 -0700
Subject: [PATCH 02/26] Keep a device round-trip when a host op writes the host
 copy

When two delegates on the same device are connected, PropagateDevicePass
removes the device -> host -> device copy pair between them, and the
second delegate reads the first one's device output directly.

That is wrong when a host op writes the host copy in place in between.
A mutable custom op that stays on the host (like a KV cache update)
does this: to_executorch removes auto_functionalized before this pass,
so the op writes its argument in place. The second delegate then reads
the device tensor from before the write and gets a stale value.

Now the pass keeps the round-trip when any op that reads the host copy
also writes it, according to the op's schema. The check reuses
_schema_writes_input from the buffer write-back pass. It does not look at
where the write happens, so it also keeps a round-trip whose second
delegate runs before the write.

The new test runs a mutable host op between two delegates on cuda:0
and checks that the second delegate reads a fresh copy of the written
host tensor. It fails without the change.
---
 exir/passes/BUCK                         |  1 +
 exir/passes/propagate_device_pass.py     | 25 +++++++-
 exir/tests/test_propagate_device_pass.py | 80 ++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 2 deletions(-)

diff --git a/exir/passes/BUCK b/exir/passes/BUCK
index 6c896ea9f5d..53a3cc199a8 100644
--- a/exir/passes/BUCK
+++ b/exir/passes/BUCK
@@ -524,6 +524,7 @@ fbcode_target(_kind = runtime.python_library,
     ],
     deps = [
         ":device_copy_ops_registry",
+        ":insert_write_back_for_buffers_pass",
         ":propagate_device_config",
         "//caffe2:torch",
         "//executorch/exir:delegate",
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index aac2de75593..a5feb0cced8 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -18,6 +18,9 @@
 import torch.utils._pytree as pytree
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.lowered_backend_module import LoweredBackendModule
+from executorch.exir.passes.insert_write_back_for_buffers_pass import (
+    _schema_writes_input,
+)
 
 # Re-exported for backward compatibility; the dataclass lives in a lightweight
 # module so that ExecutorchBackendConfig can reference it without importing the
@@ -147,6 +150,21 @@ def _clone_spec_with_device(
     return new_spec
 
 
+def _is_written_in_place(node: torch.fx.Node) -> bool:
+    """Whether an op that reads *node* also writes it, per the op's schema.
+
+    A mutable op left on the host does this once to_executorch has removed
+    auto_functionalized, so a copy of *node* can miss the write. Only direct
+    readers count: a write through a view of *node* is not seen. Custom op
+    schemas are trusted, as functionalization trusts them.
+    """
+    return any(
+        _schema_writes_input(user, user.target._schema, node)  # pyre-ignore[16]
+        for user in node.users
+        if hasattr(user.target, "_schema")
+    )
+
+
 class PropagateDevicePass(PassBase):
     """
     After to_backend, walk the graph and insert H2D/D2H copy ops at delegate
@@ -327,8 +345,10 @@ def _elide_redundant_same_device_copies(
         rewires the consumer straight to that device tensor and drops both copies.
         The device comparison is the correctness guard: a genuine cross-device
         move (e.g. ``cuda:0 -> host -> cuda:1``) has differing devices and is left
-        intact. A D2H feeding other (host) consumers keeps those users and is
-        removed only once it becomes dead.
+        intact. So is a round-trip whose host copy any op writes in place,
+        wherever that write is, because the device tensor never sees it. A D2H
+        feeding other (host) consumers keeps those users and is removed only once
+        it becomes dead.
         """
         h2d = torch.ops.et_copy._h2d_copy.default
         d2h = torch.ops.et_copy._d2h_copy.default
@@ -356,6 +376,7 @@ def _elide_redundant_same_device_copies(
             if (
                 x_spec.device != h2d_spec.device
                 or x_spec.device_index != h2d_spec.device_index
+                or _is_written_in_place(src)
             ):
                 continue
 
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 41902ddced4..c269d755989 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -142,6 +142,55 @@ def partition(self, exported_program) -> PartitionResult:
         )
 
 
+# A mutable op that stays on the host, like a KV cache update. Once
+# to_executorch removes auto_functionalized, it adds other into x in place.
+_test_lib = torch.library.Library("propagate_device_test", "DEF")
+_test_lib.define("accumulate(Tensor(a!) x, Tensor other) -> Tensor")
+_test_lib.define(
+    "accumulate.out(Tensor(a!) x, Tensor other, *, Tensor(b!) out) -> Tensor(b!)"
+)
+
+
+@torch.library.impl(_test_lib, "accumulate", "CompositeExplicitAutograd")
+def _accumulate(x: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
+    x.add_(other)
+    return x.clone()
+
+
+@torch.library.impl(_test_lib, "accumulate.out", "CompositeExplicitAutograd")
+def _accumulate_out(
+    x: torch.Tensor, other: torch.Tensor, out: torch.Tensor
+) -> torch.Tensor:
+    x.add_(other)
+    return out.copy_(x)
+
+
+def _lower_with_host_writer(
+    model: torch.nn.Module,
+    inputs: tuple,
+    devices: List[str],
+    et_config: ExecutorchBackendConfig,
+) -> torch.fx.GraphModule:
+    """Delegate each add of a model that also calls accumulate on the host."""
+    edge = to_edge(
+        export(model, inputs),
+        compile_config=EdgeCompileConfig(_check_ir_validity=False),
+    )
+    et = edge.to_backend(PerAddDeviceDelegatePartitioner(devices)).to_executorch(
+        et_config
+    )
+    return et.exported_program().graph_module
+
+
+def _host_writer(gm: torch.fx.GraphModule) -> torch.fx.Node:
+    (writer,) = [
+        n
+        for n in gm.graph.nodes
+        if n.target == torch.ops.propagate_device_test.accumulate.out
+    ]
+    return writer
+
+
 class TestPropagateDevicePass(unittest.TestCase):
     @staticmethod
     def _collect_tensor_specs(node: torch.fx.Node) -> List[TensorSpec]:
@@ -493,6 +542,37 @@ def forward(self, a, b, c):
                     f"(1 boundary D2H), got {len(nodes.d2h_nodes)}",
                 )
 
+    def test_same_device_copies_kept_when_host_op_writes_copy(self):
+        """A host op that writes the host copy of a delegate output in place
+        changes what the next delegate on the same device reads, so that
+        round-trip stays."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                a = x + y
+                torch.ops.propagate_device_test.accumulate(a, y)
+                return a + y
+
+        gm = _lower_with_host_writer(
+            Model(),
+            (torch.randn(2, 2), torch.randn(2, 2)),
+            ["cuda:0", "cuda:0"],
+            ExecutorchBackendConfig(
+                emit_stacktrace=False, enable_non_cpu_memory_planning=True
+            ),
+        )
+        written = _host_writer(gm).args[0]
+        self.assertEqual(written.target, torch.ops.et_copy._d2h_copy.out)
+        second = _collect_device_copy_nodes(gm).delegate_nodes[1]
+        self.assertIn(
+            written,
+            [
+                arg.args[0]
+                for arg in second.args[1:]
+                if arg.target == torch.ops.et_copy._h2d_copy.out
+            ],
+        )
+
     def test_cross_device_copies_preserved(self):
         """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two
         directly-connected delegates must be preserved -- the round-trip is a

From 9c26c3e9ea3837db39d560bad9f5c035d80cc7bb Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 23 Oct 2026 01:34:20 -0700
Subject: [PATCH 03/26] Copy a device delegate output to the host when it is
 used without getitem

A lowered module that returns one tensor and is called from another module
has no getitem after its call. The device pass only placed copies after
getitems, so such an output got no _d2h_copy. A host op that read it got
device memory, and another delegate on the same device got an _h2d_copy
with a device source, which the copy kernel rejects at run time.

The device-to-host copy is now inserted after the call itself in that case.
The existing elision then removes the round-trip between two calls on the
same device.
---
 exir/passes/propagate_device_pass.py     |  40 ++++++---
 exir/tests/test_propagate_device_pass.py | 110 +++++++++++++++++++++++
 2 files changed, 140 insertions(+), 10 deletions(-)

diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index a5feb0cced8..656bab8996f 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -173,16 +173,17 @@ class PropagateDevicePass(PassBase):
     When a delegate has a target_device CompileSpec (e.g., "cuda:0"):
     - For each delegate input: insert et_copy._h2d_copy before the delegate call.
       The original input node stays CPU; the h2d_copy output is tagged as device.
-    - For each delegate output: insert et_copy._d2h_copy after each getitem.
-      The getitem stays device; the d2h_copy output is tagged as CPU.
+    - For each delegate output: insert et_copy._d2h_copy after each getitem,
+      or after the call itself when it returns a single tensor.
+      The delegate output stays device; the d2h_copy output is tagged as CPU.
     - Getitem nodes that extract from a delegate call inherit the device.
 
     Skip-copy optimizations:
     - skip_h2d_for_method_inputs: If the input is a graph-level placeholder
       feeding directly to a delegate, don't insert H2D — tag the placeholder
       as device instead (user provides GPU tensor at runtime).
-    - skip_d2h_for_method_outputs: If the getitem feeds directly to graph
-      output, don't insert D2H — the output stays on device.
+    - skip_d2h_for_method_outputs: If a delegate output feeds directly to
+      graph output, don't insert D2H; the output stays on device.
     """
 
     def __init__(
@@ -298,9 +299,18 @@ def _insert_d2h_for_getitem(
             return False
 
         _set_device_on_spec(spec, source_spec.device, source_spec.device_index)
+        self._insert_d2h_after(graph_module, node, spec)
+        return True
 
+    def _insert_d2h_after(
+        self,
+        graph_module: torch.fx.GraphModule,
+        node: torch.fx.Node,
+        spec: TensorSpec,
+    ) -> None:
+        """Insert a D2H copy after *node*, a device tensor, for all of its users."""
         if self.skip_d2h_for_method_outputs and self._feeds_directly_to_output(node):
-            return True
+            return
 
         with graph_module.graph.inserting_after(node):
             d2h_node = graph_module.graph.call_function(
@@ -325,7 +335,6 @@ def _insert_d2h_for_getitem(
                         lambda value: d2h_spec if value is spec else value,
                         user.meta["spec"],
                     )
-        return True
 
     def _elide_redundant_same_device_copies(
         self,
@@ -398,7 +407,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         #   Pass 2 – For each getitem that extracts from a device-tagged delegate
         #            (tracked in device_delegates), propagate the device onto the
         #            getitem spec and insert a D2H copy after it so downstream
-        #            non-delegated ops receive CPU tensors.
+        #            non-delegated ops receive CPU tensors. A delegate call that
+        #            returns a single tensor gets the D2H copy itself.
         #   Pass 3 – Elide the redundant device->host->device round-trips left
         #            between consecutive same-device delegates by passes 1 and 2.
         changed = False
@@ -449,10 +459,20 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                     lowered_module.backend_id,
                 )
 
-        # Second pass: propagate device through getitem nodes and insert D2H
-        # only for delegates that have a target_device.
+        # Second pass: propagate device through getitem nodes and insert D2H,
+        # after each getitem or single-tensor call, only for delegates that have
+        # a target_device.
         for node in list(graph_module.graph.nodes):
-            if node.op == "call_function" and node.target == operator.getitem:
+            # A lowered module that returns one tensor and is called from another
+            # module has no getitem after its call, so copy its output directly.
+            if (
+                node in device_delegates
+                and self.enable_non_cpu_memory_planning
+                and isinstance(node.meta.get("spec"), TensorSpec)
+            ):
+                self._insert_d2h_after(graph_module, node, node.meta["spec"])
+                changed = True
+            elif node.op == "call_function" and node.target == operator.getitem:
                 source = node.args[0]
                 if isinstance(source, torch.fx.Node) and source in device_delegates:
                     if self.enable_non_cpu_memory_planning:
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index c269d755989..8e67699d71d 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -13,6 +13,7 @@
 import executorch.exir.passes._device_copy_ops_registry  # noqa: F401
 import torch
 from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower
+from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import (
     DelegationSpec,
@@ -102,6 +103,23 @@ def _collect_device_copy_nodes(gm: torch.fx.GraphModule) -> DeviceCopyNodes:
     )
 
 
+def _lower_add_to_device(inputs: tuple, device: str) -> torch.nn.Module:
+    """Lower a module that returns ``x + y`` to a delegate on ``device``."""
+
+    class Add(torch.nn.Module):
+        def forward(self, x, y):
+            return torch.add(x, y)
+
+    return to_backend(
+        BackendWithCompilerDemo.__name__,
+        to_edge(export(Add(), inputs)).exported_program(),
+        [
+            CompileSpec("max_value", bytes([4])),
+            CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")),
+        ],
+    )
+
+
 class PerAddDeviceDelegatePartitioner(Partitioner):
     """Delegates each ``aten.add.Tensor`` as its OWN delegate, each with a
     per-op ``target_device``.
@@ -573,6 +591,98 @@ def forward(self, x, y):
             ],
         )
 
+    def test_single_tensor_delegate_output_read_on_host(self):
+        """A host op that reads the output of a lowered module returning one
+        tensor reads a host copy of it."""
+
+        inputs = (torch.randn(2, 2), torch.randn(2, 2))
+        lowered = _lower_add_to_device(inputs, "cuda:0")
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.lowered = lowered
+
+            def forward(self, x, y):
+                return torch.sin(self.lowered(x, y))
+
+        et = to_edge(export(Model(), inputs)).to_executorch(
+            ExecutorchBackendConfig(
+                emit_stacktrace=False, enable_non_cpu_memory_planning=True
+            )
+        )
+        gm = et.exported_program().graph_module
+        nodes = _collect_device_copy_nodes(gm)
+        (d2h,) = nodes.d2h_nodes
+        self.assertIs(d2h.args[0], nodes.delegate_nodes[0])
+        (sin,) = d2h.users
+        self.assertEqual(sin.target, torch.ops.aten.sin.out)
+        self._assert_nodes_buffer_device(
+            [d2h], et._emitter_output.program, DeviceType.CPU, "host op", "D2H copy"
+        )
+
+    def test_single_tensor_delegate_output_copied_to_host(self):
+        """A lowered module that returns one tensor, called from another module,
+        has no getitem after its call. Its device output still gets a _d2h_copy,
+        and the next call on the same device reads it without a round-trip."""
+
+        inputs = (torch.randn(2, 2), torch.randn(2, 2))
+        lowered = _lower_add_to_device(inputs, "cuda:0")
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.lowered = lowered
+
+            def forward(self, x, y):
+                return self.lowered(self.lowered(x, y), y)
+
+        for skip_d2h in (False, True):
+            et = to_edge(export(Model(), inputs)).to_executorch(
+                ExecutorchBackendConfig(
+                    emit_stacktrace=False,
+                    propagate_device_config=PropagateDeviceConfig(
+                        skip_d2h_for_method_outputs=skip_d2h
+                    ),
+                    enable_non_cpu_memory_planning=True,
+                )
+            )
+            program = et._emitter_output.program
+            gm = et.exported_program().graph_module
+            with self.subTest(skip_d2h=skip_d2h):
+                nodes = _collect_device_copy_nodes(gm)
+                self.assertEqual(len(nodes.delegate_nodes), 2)
+                # Only method inputs are copied to the device: x and y, then y again.
+                self.assertEqual(len(nodes.h2d_nodes), 3)
+                self.assertTrue(
+                    all(n.args[0].op == "placeholder" for n in nodes.h2d_nodes)
+                )
+                # Only the method output can be copied back to the host.
+                self.assertEqual(len(nodes.d2h_nodes), 0 if skip_d2h else 1)
+                (output,) = gm.graph.output_node().args[0]
+                self.assertEqual(
+                    output.target,
+                    (
+                        executorch_call_delegate
+                        if skip_d2h
+                        else torch.ops.et_copy._d2h_copy.out
+                    ),
+                )
+                self._assert_nodes_buffer_device(
+                    nodes.delegate_nodes,
+                    program,
+                    DeviceType.CUDA,
+                    f"skip_d2h={skip_d2h}",
+                    "Delegate",
+                )
+                self._assert_nodes_buffer_device(
+                    nodes.d2h_nodes,
+                    program,
+                    DeviceType.CPU,
+                    f"skip_d2h={skip_d2h}",
+                    "D2H copy",
+                )
+
     def test_cross_device_copies_preserved(self):
         """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two
         directly-connected delegates must be preserved -- the round-trip is a

From 8f40af71a65ee6a8a167e829f88847f166223b39 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 23 Oct 2026 10:42:52 -0700
Subject: [PATCH 04/26] Return delegate outputs on the device whenever
 skip_d2h_for_method_outputs is set

With skip_d2h_for_method_outputs, a delegate output that the method returns
should stay on the device. The pass skipped the device-to-host copy only
when the graph output was the copy's only reader, and it decided this while
it was still inserting copies. So a returned output that something else
also read was returned on the host:

- A delegate on the same device: its input copy was still a reader at that
  point. The round-trip was removed later, but the host copy stayed.
- A host op, or a delegate on another device: these need the host copy, so
  the graph output kept reading it too.

For example, `a = f(x, y); return a, torch.sin(a)` returned `a` on the host.

The pass now inserts every device-to-host copy, removes the round-trips,
and then points the graph output at the device tensor. A copy stays only
for the readers that still need it.

When a host op writes the host copy in place, the device tensor never sees
the write, so the output reads a new device copy of the written host copy.
---
 exir/passes/propagate_device_config.py   |  10 +-
 exir/passes/propagate_device_pass.py     |  87 +++++++++---
 exir/tests/test_propagate_device_pass.py | 165 +++++++++++++++++++++++
 3 files changed, 239 insertions(+), 23 deletions(-)

diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index 6071209ffeb..aa53e717576 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -56,10 +56,12 @@ class PropagateDeviceConfig:
     # A dict can be used to set per-method values, keyed by method name.
     skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False
 
-    # When True, device delegate outputs that are directly method outputs
-    # are NOT wrapped with _d2h_copy. The method outputs stay on device.
-    # Useful for cross-method GPU pipelines where the next method consumes
-    # GPU tensors directly.
+    # When True, device delegate outputs that the method returns stay on
+    # device: the method output never reads a _d2h_copy of them. When a host op
+    # writes the host copy of such an output in place, the method returns a
+    # device copy of the written tensor, which costs a copy on every call.
+    # Useful for cross-method GPU pipelines where the next method consumes GPU
+    # tensors directly.
     #
     # Pair with MemoryPlanningPass(alloc_graph_output=False) for the same reason as the input flag.
     # A dict can be used to set per-method values, keyed by method name.
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 656bab8996f..22a08df8947 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -165,6 +165,22 @@ def _is_written_in_place(node: torch.fx.Node) -> bool:
     )
 
 
+def _replace_output_spec(
+    node: torch.fx.Node, old_spec: TensorSpec, new_spec: TensorSpec
+) -> None:
+    """Replace *old_spec* with *new_spec* in the spec of the output node reading *node*.
+
+    Memory planning finds the graph outputs through the output node's spec, so
+    it has to name the tensor that is actually returned.
+    """
+    for user in node.users:
+        if user.op == "output" and "spec" in user.meta:
+            user.meta["spec"] = pytree.tree_map(
+                lambda value: new_spec if value is old_spec else value,
+                user.meta["spec"],
+            )
+
+
 class PropagateDevicePass(PassBase):
     """
     After to_backend, walk the graph and insert H2D/D2H copy ops at delegate
@@ -182,8 +198,10 @@ class PropagateDevicePass(PassBase):
     - skip_h2d_for_method_inputs: If the input is a graph-level placeholder
       feeding directly to a delegate, don't insert H2D — tag the placeholder
       as device instead (user provides GPU tensor at runtime).
-    - skip_d2h_for_method_outputs: If a delegate output feeds directly to
-      graph output, don't insert D2H; the output stays on device.
+    - skip_d2h_for_method_outputs: The graph output reads a delegate output
+      on device, not its D2H copy; the copy stays only for other readers. When
+      one of them writes the copy in place, the graph output reads a device
+      copy of the written tensor instead.
     """
 
     def __init__(
@@ -212,10 +230,6 @@ def _is_placeholder(self, node: torch.fx.Node) -> bool:
         """Check if a node is a graph-level input (placeholder)."""
         return node.op == "placeholder"
 
-    def _feeds_directly_to_output(self, node: torch.fx.Node) -> bool:
-        """Check if all users of a node are output nodes."""
-        return all(user.op == "output" for user in node.users)
-
     def _insert_h2d_copies(
         self,
         graph_module: torch.fx.GraphModule,
@@ -309,9 +323,6 @@ def _insert_d2h_after(
         spec: TensorSpec,
     ) -> None:
         """Insert a D2H copy after *node*, a device tensor, for all of its users."""
-        if self.skip_d2h_for_method_outputs and self._feeds_directly_to_output(node):
-            return
-
         with graph_module.graph.inserting_after(node):
             d2h_node = graph_module.graph.call_function(
                 torch.ops.et_copy._d2h_copy.default,
@@ -327,14 +338,47 @@ def _insert_d2h_after(
                 d2h_node,
                 delete_user_cb=lambda user, _d2h=d2h_node: user != _d2h,
             )
-            # Memory planning finds the graph outputs through the output node's
-            # spec, so it has to name the copy now rather than the device tensor.
-            for user in d2h_node.users:
-                if user.op == "output" and "spec" in user.meta:
-                    user.meta["spec"] = pytree.tree_map(
-                        lambda value: d2h_spec if value is spec else value,
-                        user.meta["spec"],
+            _replace_output_spec(d2h_node, spec, d2h_spec)
+
+    def _keep_method_outputs_on_device(
+        self,
+        graph_module: torch.fx.GraphModule,
+    ) -> None:
+        """Make the graph output read the device tensor instead of its D2H copy.
+
+        The copy stays for its other users, such as a host op. When one of them
+        writes the copy in place, the device tensor never sees that write, so the
+        output reads a new device copy of the written host copy instead.
+        """
+        d2h = torch.ops.et_copy._d2h_copy.default
+        for node in list(graph_module.graph.nodes):
+            if (
+                node.op != "call_function"
+                or node.target != d2h
+                or not any(user.op == "output" for user in node.users)
+            ):
+                continue
+            returned = node.args[0]
+            if _is_written_in_place(node):
+                device_spec = returned.meta["spec"]
+                with graph_module.graph.inserting_before(
+                    graph_module.graph.output_node()
+                ):
+                    returned = graph_module.graph.call_function(
+                        torch.ops.et_copy._h2d_copy.default, (node,)
                     )
+                returned.meta["spec"] = _clone_spec_with_device(
+                    node.meta["spec"], device_spec.device, device_spec.device_index
+                )
+                returned.meta["val"] = node.meta.get("val")
+                if "tensor_meta" in node.meta:
+                    returned.meta["tensor_meta"] = node.meta["tensor_meta"]
+            _replace_output_spec(node, node.meta["spec"], returned.meta["spec"])
+            node.replace_all_uses_with(
+                returned, delete_user_cb=lambda user: user.op == "output"
+            )
+            if not node.users:
+                graph_module.graph.erase_node(node)
 
     def _elide_redundant_same_device_copies(
         self,
@@ -411,6 +455,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         #            returns a single tensor gets the D2H copy itself.
         #   Pass 3 – Elide the redundant device->host->device round-trips left
         #            between consecutive same-device delegates by passes 1 and 2.
+        #            With skip_d2h_for_method_outputs, then make the graph output
+        #            read device tensors instead of D2H copies.
         changed = False
         device_delegates: set[torch.fx.Node] = set()
 
@@ -498,11 +544,14 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
 
         # Third pass: elide the redundant device->host->device round-trips left
         # between consecutive same-device delegates by the per-boundary H2D/D2H
-        # insertion above. Only meaningful when passes 1/2 modified the graph
-        # (there was a device delegate); when nothing changed there are no copies
-        # to elide, so skip the scan entirely.
+        # insertion above, then, with skip_d2h_for_method_outputs, make the graph
+        # output read device tensors. Only meaningful when passes 1/2 modified the
+        # graph (there was a device delegate); when nothing changed there are no
+        # copies to elide, so skip the scan entirely.
         if changed:
             self._elide_redundant_same_device_copies(graph_module)
+            if self.skip_d2h_for_method_outputs:
+                self._keep_method_outputs_on_device(graph_module)
 
         graph_module.recompile()
         return PassResult(graph_module, changed)
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 8e67699d71d..a90be18fd4f 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -683,6 +683,171 @@ def forward(self, x, y):
                     "D2H copy",
                 )
 
+    def test_skip_d2h_keeps_output_also_read_on_device(self):
+        """With skip_d2h_for_method_outputs, a returned delegate output that a
+        delegate on the same device also reads stays on the device, like the
+        other outputs."""
+
+        inputs = (torch.randn(2, 2), torch.randn(2, 2))
+        lowered = _lower_add_to_device(inputs, "cuda:0")
+
+        class Calls(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.lowered = lowered
+
+            def forward(self, x, y):
+                a = self.lowered(x, y)
+                return a, self.lowered(a, y)
+
+        class Adds(torch.nn.Module):
+            def forward(self, x, y):
+                a = torch.add(x, y)
+                return a, torch.add(a, y)
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_d2h_for_method_outputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        cases = {
+            "single tensor": (Calls(), None),
+            "getitem": (Adds(), PerAddDeviceDelegatePartitioner(["cuda:0"])),
+        }
+        for case, (model, partitioner) in cases.items():
+            edge = to_edge(export(model, inputs))
+            if partitioner is not None:
+                edge = edge.to_backend(partitioner)
+            et = edge.to_executorch(deepcopy(et_config))
+            gm = et.exported_program().graph_module
+            with self.subTest(case=case):
+                self.assertEqual(len(_collect_device_copy_nodes(gm).d2h_nodes), 0)
+                outputs = gm.graph.output_node().args[0]
+                self._assert_nodes_device(outputs, DeviceType.CUDA, case, "Output")
+                self._assert_nodes_buffer_device(
+                    outputs, et._emitter_output.program, DeviceType.CUDA, case, "Output"
+                )
+
+    def test_skip_d2h_keeps_output_also_read_off_device(self):
+        """With skip_d2h_for_method_outputs, a returned delegate output that a
+        host op or a delegate on another device also reads stays on the device.
+        Only that other reader gets the host copy."""
+
+        inputs = (torch.randn(2, 2), torch.randn(2, 2))
+        lowered = _lower_add_to_device(inputs, "cuda:0")
+        lowered_cuda_1 = _lower_add_to_device(inputs, "cuda:1")
+
+        class HostReader(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.lowered = lowered
+
+            def forward(self, x, y):
+                a = self.lowered(x, y)
+                return a, torch.sin(a)
+
+        class OtherDeviceReader(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.lowered = lowered
+                self.lowered_cuda_1 = lowered_cuda_1
+
+            def forward(self, x, y):
+                a = self.lowered(x, y)
+                return a, self.lowered_cuda_1(a, y)
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_d2h_for_method_outputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        for case, model in {
+            "host op": HostReader(),
+            "delegate on cuda:1": OtherDeviceReader(),
+        }.items():
+            et = to_edge(export(model, inputs)).to_executorch(deepcopy(et_config))
+            gm = et.exported_program().graph_module
+            with self.subTest(case=case):
+                output = gm.graph.output_node().args[0][0]
+                self.assertEqual(output.target, executorch_call_delegate)
+                self._assert_nodes_device([output], DeviceType.CUDA, case, "Output", 0)
+                self._assert_nodes_buffer_device(
+                    [output],
+                    et._emitter_output.program,
+                    DeviceType.CUDA,
+                    case,
+                    "Output",
+                )
+                (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes
+                self.assertIs(d2h.args[0], output)
+                self.assertNotIn(gm.graph.output_node(), d2h.users)
+
+    def test_skip_d2h_leaves_returned_output_unplanned(self):
+        """With skip_d2h_for_method_outputs and unplanned graph outputs, the
+        caller provides the returned delegate output, while the host copy that
+        a host op reads is still planned."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                a = x + y
+                return a, torch.sin(a)
+
+        et = (
+            to_edge(export(Model(), (torch.randn(2, 2), torch.randn(2, 2))))
+            .to_backend(PerAddDeviceDelegatePartitioner(["cuda:0"]))
+            .to_executorch(
+                ExecutorchBackendConfig(
+                    emit_stacktrace=False,
+                    propagate_device_config=PropagateDeviceConfig(
+                        skip_d2h_for_method_outputs=True
+                    ),
+                    enable_non_cpu_memory_planning=True,
+                    memory_planning_pass=MemoryPlanningPass(alloc_graph_output=False),
+                )
+            )
+        )
+        gm = et.exported_program().graph_module
+        output = gm.graph.output_node().args[0][0]
+        (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes
+        self.assertIs(d2h.args[0], output)
+        self.assertIsNone(output.meta["spec"].mem_id)
+        self.assertIsNotNone(d2h.meta["spec"].mem_id)
+
+    def test_skip_d2h_returns_device_copy_of_output_written_on_host(self):
+        """With skip_d2h_for_method_outputs, a returned delegate output whose
+        host copy a host op writes in place is returned as a device copy of the
+        written host copy, because the device tensor never sees the write."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                a = x + y
+                torch.ops.propagate_device_test.accumulate(a, y)
+                return a
+
+        gm = _lower_with_host_writer(
+            Model(),
+            (torch.randn(2, 2), torch.randn(2, 2)),
+            ["cuda:0"],
+            ExecutorchBackendConfig(
+                emit_stacktrace=False,
+                propagate_device_config=PropagateDeviceConfig(
+                    skip_d2h_for_method_outputs=True
+                ),
+                enable_non_cpu_memory_planning=True,
+            ),
+        )
+        (output,) = gm.graph.output_node().args[0]
+        self.assertEqual(output.target, torch.ops.et_copy._h2d_copy.out)
+        self.assertEqual(output.meta["spec"].device, DeviceType.CUDA)
+        writer = _host_writer(gm)
+        self.assertIs(output.args[0], writer.args[0])
+        nodes = list(gm.graph.nodes)
+        self.assertLess(nodes.index(writer), nodes.index(output))
+
     def test_cross_device_copies_preserved(self):
         """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two
         directly-connected delegates must be preserved -- the round-trip is a

From ccab255c010a77ddd0b6da8e1a8f382a93915e75 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 22 Oct 2026 12:08:26 -0700
Subject: [PATCH 05/26] Copy constants that feed a device delegate into device
 memory

A constant that feeds a delegate on a device, such as a buffer the
partitioner leaves outside the delegate, broke device memory planning.

- Without skip_h2d_for_method_inputs, its _h2d_copy cloned the constant's
  spec, constant flag included, so memory planning skipped the copy's
  output and emission failed.
- With it, the constant was treated as a method input: its placeholder was
  tagged with the device while its data stays in the program on the host,
  so the delegate got a host pointer labeled as device memory.

A copy's output is now never a constant, and skip_h2d_for_method_inputs
applies only to placeholders that are not constants. The config comment and
the CUDA docs say that this copy runs on every call, and the pass logs a
warning that names the constant and its size, so the cost is visible at
export time.

The copy moves the constant's bytes as they are stored. The program stores
a constant that is neither contiguous nor channels-last as contiguous, so
such a constant, which a pass that runs after to_edge can create, is now
refused with a clear error instead of getting a copy with the wrong layout.
---
 docs/source/backends/cuda/cuda-overview.md |   6 +
 exir/passes/BUCK                           |   1 +
 exir/passes/propagate_device_config.py     |   5 +-
 exir/passes/propagate_device_pass.py       |  54 +++++++--
 exir/tests/test_propagate_device_pass.py   | 135 ++++++++++++++++++++-
 5 files changed, 190 insertions(+), 11 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index ef57aca7ae5..938a864e257 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -169,6 +169,12 @@ out works, and the snippet above sets it explicitly only to make the requirement
 it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens
 during device-aware memory planning.
 
+**A constant that a delegate reads from outside is still copied.** Constants, such as weights and
+buffers the method does not change, are not passed by the caller. Their data stays in host memory,
+in the program or in its external data file, so `to_executorch` adds a copy to the device that runs
+on every call, one for each delegate that reads the constant, and logs a warning. Keep constants
+inside the delegate to avoid that cost.
+
 **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the
 program reserves its own input buffer and the runtime copies the caller's memory into it, so the
 copy comes back at run time. On device memory that copy is a host memcpy into a device pointer,
diff --git a/exir/passes/BUCK b/exir/passes/BUCK
index 53a3cc199a8..7a25cbaa134 100644
--- a/exir/passes/BUCK
+++ b/exir/passes/BUCK
@@ -523,6 +523,7 @@ fbcode_target(_kind = runtime.python_library,
         "propagate_device_pass.py",
     ],
     deps = [
+        ":convert_constant_dim_order_pass",
         ":device_copy_ops_registry",
         ":insert_write_back_for_buffers_pass",
         ":propagate_device_config",
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index aa53e717576..265aece65e5 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -49,7 +49,10 @@ class PropagateDeviceConfig:
     # When True, method-level input tensors that feed directly into a device
     # delegate are NOT wrapped with _h2d_copy. The user must provide tensors
     # already on the target device. Useful for pipelines where inputs are
-    # pre-staged on GPU.
+    # pre-staged on GPU. Constants are not passed by the caller, so they are
+    # still copied: their data stays on the host, in the program or in its
+    # external data file. That copy runs on every call, once for each delegate
+    # that reads the constant.
     #
     # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer
     # for the input and copies the caller's memory into it, which is what this exists to avoid.
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 22a08df8947..39613eda1ed 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -18,6 +18,7 @@
 import torch.utils._pytree as pytree
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.lowered_backend_module import LoweredBackendModule
+from executorch.exir.passes.convert_constant_dim_order_pass import _should_transform
 from executorch.exir.passes.insert_write_back_for_buffers_pass import (
     _schema_writes_input,
 )
@@ -138,14 +139,19 @@ def _tag_specs_with_device(
     return False
 
 
-def _clone_spec_with_device(
+def _copy_output_spec(
     spec: TensorSpec,
     device_type: schema.DeviceType,
     device_index: int = 0,
 ) -> TensorSpec:
-    """Create a copy of a TensorSpec with a different device."""
+    """Create the spec of an _h2d_copy or _d2h_copy output from the spec it copies.
+
+    The output lives on ``device_type``. The copy writes it at run time, so it is
+    never a constant, even when the copied tensor is one.
+    """
     new_spec = copy.copy(spec)
     new_spec.init_mem_planning_fields()
+    new_spec.const = False
     _set_device_on_spec(new_spec, device_type, device_index)
     return new_spec
 
@@ -197,7 +203,9 @@ class PropagateDevicePass(PassBase):
     Skip-copy optimizations:
     - skip_h2d_for_method_inputs: If the input is a graph-level placeholder
       feeding directly to a delegate, don't insert H2D — tag the placeholder
-      as device instead (user provides GPU tensor at runtime).
+      as device instead (user provides GPU tensor at runtime). Constants are
+      placeholders too, but their data is stored on the host, so they still
+      get an H2D copy.
     - skip_d2h_for_method_outputs: The graph output reads a delegate output
       on device, not its D2H copy; the copy stays only for other readers. When
       one of them writes the copy in place, the graph output reads a device
@@ -247,7 +255,11 @@ def _insert_h2d_copies(
             if not isinstance(arg_spec, TensorSpec):
                 continue
 
-            if self.skip_h2d_for_method_inputs and self._is_placeholder(arg):
+            if (
+                self.skip_h2d_for_method_inputs
+                and self._is_placeholder(arg)
+                and not arg_spec.const
+            ):
                 # TODO(gasoonjia): support skip_h2d_for_method_inputs for
                 # multiple-user placeholder inputs.
                 if len(arg.users) != 1:
@@ -264,14 +276,38 @@ def _insert_h2d_copies(
                 changed = True
                 continue
 
+            val = arg.meta.get("val")
+            if (
+                arg_spec.const
+                and isinstance(val, torch.Tensor)
+                and _should_transform(val)
+            ):
+                raise RuntimeError(
+                    f"Constant '{arg.name}' feeds a delegate on "
+                    f"{target_device_type.name.lower()}:{device_index}, but it is "
+                    f"neither contiguous nor channels-last. The program stores "
+                    f"such a constant contiguous, so its device copy would not "
+                    f"match the stored data. Make the constant contiguous where "
+                    f"it is created, for example in the pass that adds it."
+                )
+            if arg_spec.const:
+                logger.warning(
+                    "Constant '%s' (%d bytes) feeds a delegate on %s:%d from "
+                    "outside it, so it is copied to the device on every call. "
+                    "Have the partitioner take it into the delegate to avoid "
+                    "that copy.",
+                    arg.name,
+                    arg_spec.nbytes(),
+                    target_device_type.name.lower(),
+                    device_index,
+                )
+
             with graph_module.graph.inserting_before(node):
                 h2d_node = graph_module.graph.call_function(
                     torch.ops.et_copy._h2d_copy.default,
                     (arg,),
                 )
-                h2d_spec = _clone_spec_with_device(
-                    arg_spec, target_device_type, device_index
-                )
+                h2d_spec = _copy_output_spec(arg_spec, target_device_type, device_index)
                 h2d_node.meta["spec"] = h2d_spec
                 h2d_node.meta["val"] = arg.meta.get("val")
                 if "tensor_meta" in arg.meta:
@@ -328,7 +364,7 @@ def _insert_d2h_after(
                 torch.ops.et_copy._d2h_copy.default,
                 (node,),
             )
-            d2h_spec = _clone_spec_with_device(spec, schema.DeviceType.CPU, 0)
+            d2h_spec = _copy_output_spec(spec, schema.DeviceType.CPU, 0)
             d2h_node.meta["spec"] = d2h_spec
             d2h_node.meta["val"] = node.meta.get("val")
             if "tensor_meta" in node.meta:
@@ -367,7 +403,7 @@ def _keep_method_outputs_on_device(
                     returned = graph_module.graph.call_function(
                         torch.ops.et_copy._h2d_copy.default, (node,)
                     )
-                returned.meta["spec"] = _clone_spec_with_device(
+                returned.meta["spec"] = _copy_output_spec(
                     node.meta["spec"], device_spec.device, device_spec.device_index
                 )
                 returned.meta["val"] = node.meta.get("val")
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index a90be18fd4f..2ceefe7dc28 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -30,7 +30,8 @@
 from executorch.exir.capture._config import ExecutorchBackendConfig
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.passes import MemoryPlanningPass
+from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult
+from executorch.exir.passes import MemoryPlanningPass, propagate_device_pass
 from executorch.exir.passes.propagate_device_pass import (
     _get_target_device_from_compile_specs,
     _parse_device_spec_value,
@@ -1620,6 +1621,138 @@ def forward(self, a, b):
                     "CPU placeholder",
                 )
 
+    def test_constant_feeding_device_delegate_is_copied(self):
+        """A constant feeding a device delegate gets an _h2d_copy into planned device
+        memory, with or without skip_h2d_for_method_inputs: its data lives in the
+        program on the host, so it is not a method input the caller can place."""
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("offset", torch.randn(2, 2))
+
+            def forward(self, a):
+                return torch.add(a, self.offset)
+
+        model = Model()
+        inputs = (torch.randn(2, 2),)
+        for skip in (False, True):
+            et_config = ExecutorchBackendConfig(
+                emit_stacktrace=False,
+                propagate_device_config=PropagateDeviceConfig(
+                    skip_h2d_for_method_inputs=skip
+                ),
+                enable_non_cpu_memory_planning=True,
+            )
+            with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs:
+                results = self._get_executorch_program(
+                    model, inputs, DeviceAwarePartitioner("cuda:0"), et_config
+                )
+            self.assertIn("Constant 'b_offset' (16 bytes)", logs.output[0])
+            for pipeline, program, gm in results:
+                with self.subTest(pipeline=pipeline, skip=skip):
+                    nodes = _collect_device_copy_nodes(gm)
+                    # Input a is copied only when the caller cannot place it.
+                    self.assertEqual(len(nodes.h2d_nodes), 1 if skip else 2)
+                    cuda_ph, _ = self._collect_placeholders_by_device(gm)
+                    self.assertEqual(len(cuda_ph), int(skip))
+                    self.assertTrue(all(not n.meta["spec"].const for n in cuda_ph))
+                    copies = [n.meta["spec"] for n in nodes.h2d_nodes]
+                    self.assertTrue(all(not spec.const for spec in copies))
+                    self._assert_nodes_buffer_device(
+                        nodes.h2d_nodes, program, DeviceType.CUDA, pipeline, "H2D copy"
+                    )
+
+    def test_constant_copy_keeps_stored_layout(self):
+        """The copy of a constant moves its bytes as they are, so it must use the
+        layout the constant is stored in: contiguous for a transposed parameter,
+        which to_edge makes contiguous, and channels-last for a channels-last
+        one."""
+
+        class Model(torch.nn.Module):
+            def __init__(self, weight):
+                super().__init__()
+                self.weight = torch.nn.Parameter(weight)
+
+            def forward(self, a):
+                return torch.add(a, self.weight)
+
+        weights = {
+            "transposed": torch.randn(3, 2).t(),
+            "channels_last": torch.randn(1, 2, 3, 3).to(
+                memory_format=torch.channels_last
+            ),
+        }
+        for name, weight in weights.items():
+            et_config = ExecutorchBackendConfig(
+                emit_stacktrace=False,
+                propagate_device_config=PropagateDeviceConfig(
+                    skip_h2d_for_method_inputs=True
+                ),
+                enable_non_cpu_memory_planning=True,
+            )
+            for pipeline, _program, gm in self._get_executorch_program(
+                Model(weight),
+                (torch.randn(weight.shape),),
+                DeviceAwarePartitioner("cuda:0"),
+                et_config,
+            ):
+                with self.subTest(pipeline=pipeline, weight=name):
+                    nodes = _collect_device_copy_nodes(gm)
+                    self.assertEqual(len(nodes.h2d_nodes), 1)
+                    h2d = nodes.h2d_nodes[0]
+                    # The emitter rewrites the constant's spec to the stored layout.
+                    stored = h2d.args[0].meta["spec"]
+                    self.assertEqual(h2d.meta["spec"].dim_order, stored.dim_order)
+                    self.assertEqual(
+                        tuple(h2d.meta["spec"].stride), tuple(stored.stride)
+                    )
+
+    def test_constant_with_unstored_layout_is_refused(self):
+        """A constant that is neither contiguous nor channels-last is stored
+        contiguous, so its copy would not match the stored data. A pass that runs
+        after to_edge normalized the constants can create one, and lowering then
+        fails instead of producing wrong results."""
+
+        class StoreWeightTransposed(ExportedProgramPassBase):
+            def call(self, ep):
+                ep.state_dict["w"] = torch.nn.Parameter(
+                    ep.state_dict["w"].t().contiguous().t()
+                )
+                for node in ep.graph.nodes:
+                    if node.name == "p_w":
+                        node.meta["val"] = node.meta["val"].t().contiguous().t()
+                return ExportedProgramPassResult(ep, True)
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = torch.nn.Parameter(torch.randn(2, 3))
+
+            def forward(self, a):
+                return torch.add(a, self.w)
+
+        for skip in (False, True):
+            edge = to_edge_transform_and_lower(
+                export(Model(), (torch.randn(2, 3),)),
+                transform_passes=[StoreWeightTransposed()],
+                partitioner=[DeviceAwarePartitioner("cuda:0")],
+            )
+            et_config = ExecutorchBackendConfig(
+                emit_stacktrace=False,
+                propagate_device_config=PropagateDeviceConfig(
+                    skip_h2d_for_method_inputs=skip
+                ),
+                enable_non_cpu_memory_planning=True,
+            )
+            with self.subTest(skip=skip):
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    "Constant 'p_w' feeds a delegate on cuda:0, but it is neither "
+                    "contiguous nor channels-last",
+                ):
+                    edge.to_executorch(et_config)
+
     def test_tensorspec_repr_includes_device(self):
         spec = TensorSpec(dtype=torch.float32, shape=torch.Size([2, 3]))
         repr_str = repr(spec)

From 73dc9f19a9076d03687760284d9dc62785193a63 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 22 Oct 2026 19:25:36 -0700
Subject: [PATCH 06/26] Let a method input feed several delegates on one device
 without a copy

With skip_h2d_for_method_inputs, a method input had to have exactly one
user, because its spec is changed in place to the delegate's device. A
program that calls one lowered module more than once, or passes one input
to two delegates, could not keep its inputs on the device.

Changing the spec is safe when every user is a delegate on that same
device, so that case is now allowed. An input that a host op or a delegate
on another device also reads is still refused, and the error names that
reader and the device the input was going to. When one of the readers is
the write-back of a mutated input or mutable buffer, the error names it
and says so, and for a mutable buffer it points to tag_mutated_buffer,
which lets the partitioner keep the buffer inside the delegate. The CUDA
docs now list these rules.
---
 docs/source/backends/cuda/cuda-overview.md |   9 +
 exir/passes/propagate_device_config.py     |  11 +-
 exir/passes/propagate_device_pass.py       |  66 +++++--
 exir/tests/test_propagate_device_pass.py   | 203 +++++++++++++++++++--
 4 files changed, 252 insertions(+), 37 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index 938a864e257..d0f9568fa99 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -169,6 +169,15 @@ out works, and the snippet above sets it explicitly only to make the requirement
 it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens
 during device-aware memory planning.
 
+**With `skip_h2d_for_method_inputs`, an input that a device delegate reads may only be read by
+delegates on that device.** If anything else reads it, even only its shape, such as an operator
+outside those delegates, a delegate on another device, or the method output, `to_executorch` raises
+an error instead of adding a copy. Turn the flag off, or have the partitioner put those operators in
+a delegate on that device. `to_executorch` also refuses an input that the method changes, because
+the change is written back on the host. For a mutable buffer, have the partitioner put the buffer
+and every operator that reads or writes it in the delegate. `tag_mutated_buffer` tags the buffer
+once those operators are there, and the CUDA partitioner already calls it.
+
 **A constant that a delegate reads from outside is still copied.** Constants, such as weights and
 buffers the method does not change, are not passed by the caller. Their data stays in host memory,
 in the program or in its external data file, so `to_executorch` adds a copy to the device that runs
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index 265aece65e5..54d51851db9 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -48,11 +48,12 @@ class PropagateDeviceConfig:
 
     # When True, method-level input tensors that feed directly into a device
     # delegate are NOT wrapped with _h2d_copy. The user must provide tensors
-    # already on the target device. Useful for pipelines where inputs are
-    # pre-staged on GPU. Constants are not passed by the caller, so they are
-    # still copied: their data stays on the host, in the program or in its
-    # external data file. That copy runs on every call, once for each delegate
-    # that reads the constant.
+    # already on the target device. Every reader of such an input must then be
+    # a delegate on that device, or to_executorch raises. Useful for pipelines
+    # where inputs are pre-staged on GPU. Constants are not passed by the
+    # caller, so they are still copied: their data stays on the host, in the
+    # program or in its external data file. That copy runs on every call, once
+    # for each delegate that reads the constant.
     #
     # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer
     # for the input and copies the caller's memory into it, which is what this exists to avoid.
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 39613eda1ed..f6451b96bed 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -202,10 +202,11 @@ class PropagateDevicePass(PassBase):
 
     Skip-copy optimizations:
     - skip_h2d_for_method_inputs: If the input is a graph-level placeholder
-      feeding directly to a delegate, don't insert H2D — tag the placeholder
-      as device instead (user provides GPU tensor at runtime). Constants are
-      placeholders too, but their data is stored on the host, so they still
-      get an H2D copy.
+      feeding a delegate, don't insert H2D; tag the placeholder as device
+      instead (user provides GPU tensor at runtime). Every reader of the
+      placeholder must then be a delegate on that device, or the pass raises.
+      Constants are placeholders too, but their data is stored on the host,
+      so they still get an H2D copy.
     - skip_d2h_for_method_outputs: The graph output reads a delegate output
       on device, not its D2H copy; the copy stays only for other readers. When
       one of them writes the copy in place, the graph output reads a device
@@ -238,6 +239,17 @@ def _is_placeholder(self, node: torch.fx.Node) -> bool:
         """Check if a node is a graph-level input (placeholder)."""
         return node.op == "placeholder"
 
+    def _delegate_device(
+        self, graph_module: torch.fx.GraphModule, node: torch.fx.Node
+    ) -> Optional[tuple[schema.DeviceType, int]]:
+        """The target device of a delegate call node, or None when there is none."""
+        if node.op != "call_function" or node.target != executorch_call_delegate:
+            return None
+        lowered_module = _get_lowered_module(graph_module, node)
+        if lowered_module is None:
+            return None
+        return _get_target_device_from_compile_specs(lowered_module)
+
     def _insert_h2d_copies(
         self,
         graph_module: torch.fx.GraphModule,
@@ -260,17 +272,43 @@ def _insert_h2d_copies(
                 and self._is_placeholder(arg)
                 and not arg_spec.const
             ):
-                # TODO(gasoonjia): support skip_h2d_for_method_inputs for
-                # multiple-user placeholder inputs.
-                if len(arg.users) != 1:
+                # The spec is shared by every reader, so all of them must be
+                # delegates on this same device.
+                other_readers = [
+                    user
+                    for user in arg.users
+                    if self._delegate_device(graph_module, user)
+                    != (target_device_type, device_index)
+                ]
+                if other_readers:
+                    write_back = next(
+                        (
+                            user
+                            for user in other_readers
+                            if user.target == torch.ops.aten.copy_.default
+                            and user.args[0] is arg
+                        ),
+                        None,
+                    )
+                    refusal = (
+                        f"skip_h2d_for_method_inputs=True cannot place placeholder "
+                        f"'{arg.name}' on {target_device_type.name.lower()}:"
+                        f"{device_index}, because "
+                        f"'{(write_back or other_readers[0]).name}'"
+                    )
+                    if write_back is not None:
+                        raise RuntimeError(
+                            f"{refusal} writes it back on the host (the method "
+                            f"mutates it). Either disable "
+                            f"skip_h2d_for_method_inputs, or, for a mutable buffer, "
+                            f"have the partitioner put the buffer and every op that "
+                            f"reads or writes it in the delegate; tag_mutated_buffer "
+                            f"tags the buffer once those ops are there."
+                        )
                     raise RuntimeError(
-                        f"skip_h2d_for_method_inputs=True requires placeholder "
-                        f"'{arg.name}' to have exactly one user, but it has "
-                        f"{len(arg.users)} users. The placeholder is shared by "
-                        f"multiple consumers, so its TensorSpec cannot be safely "
-                        f"mutated in-place to the delegate's device. Either disable "
-                        f"skip_h2d_for_method_inputs, or ensure the placeholder is "
-                        f"used exclusively by this delegate."
+                        f"{refusal} also reads it and is not a delegate on that "
+                        f"device. Either disable skip_h2d_for_method_inputs, or use "
+                        f"the placeholder only from delegates on one device."
                     )
                 _set_device_on_spec(arg_spec, target_device_type, device_index)
                 changed = True
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 2ceefe7dc28..5c1df35b19e 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -625,7 +625,8 @@ def forward(self, x, y):
     def test_single_tensor_delegate_output_copied_to_host(self):
         """A lowered module that returns one tensor, called from another module,
         has no getitem after its call. Its device output still gets a _d2h_copy,
-        and the next call on the same device reads it without a round-trip."""
+        and the next call on the same device reads it without a round-trip. With
+        skip_h2d_for_method_inputs, y, read by both calls, is placed once."""
 
         inputs = (torch.randn(2, 2), torch.randn(2, 2))
         lowered = _lower_add_to_device(inputs, "cuda:0")
@@ -638,26 +639,35 @@ def __init__(self):
             def forward(self, x, y):
                 return self.lowered(self.lowered(x, y), y)
 
-        for skip_d2h in (False, True):
+        configs = (
+            PropagateDeviceConfig(),
+            PropagateDeviceConfig(skip_h2d_for_method_inputs=True),
+            PropagateDeviceConfig(skip_d2h_for_method_outputs=True),
+        )
+        for device_config in configs:
+            skip_h2d = device_config.skip_h2d_for_method_inputs
+            skip_d2h = device_config.skip_d2h_for_method_outputs
             et = to_edge(export(Model(), inputs)).to_executorch(
                 ExecutorchBackendConfig(
                     emit_stacktrace=False,
-                    propagate_device_config=PropagateDeviceConfig(
-                        skip_d2h_for_method_outputs=skip_d2h
-                    ),
+                    propagate_device_config=device_config,
                     enable_non_cpu_memory_planning=True,
                 )
             )
             program = et._emitter_output.program
             gm = et.exported_program().graph_module
-            with self.subTest(skip_d2h=skip_d2h):
+            case = f"skip_h2d={skip_h2d}, skip_d2h={skip_d2h}"
+            with self.subTest(case=case):
                 nodes = _collect_device_copy_nodes(gm)
                 self.assertEqual(len(nodes.delegate_nodes), 2)
-                # Only method inputs are copied to the device: x and y, then y again.
-                self.assertEqual(len(nodes.h2d_nodes), 3)
+                # Only method inputs are copied to the device: x and y, then y
+                # again, unless the caller places them.
+                self.assertEqual(len(nodes.h2d_nodes), 0 if skip_h2d else 3)
                 self.assertTrue(
                     all(n.args[0].op == "placeholder" for n in nodes.h2d_nodes)
                 )
+                cuda_ph, _ = self._collect_placeholders_by_device(gm)
+                self.assertEqual(len(cuda_ph), 2 if skip_h2d else 0)
                 # Only the method output can be copied back to the host.
                 self.assertEqual(len(nodes.d2h_nodes), 0 if skip_d2h else 1)
                 (output,) = gm.graph.output_node().args[0]
@@ -670,18 +680,10 @@ def forward(self, x, y):
                     ),
                 )
                 self._assert_nodes_buffer_device(
-                    nodes.delegate_nodes,
-                    program,
-                    DeviceType.CUDA,
-                    f"skip_d2h={skip_d2h}",
-                    "Delegate",
+                    nodes.delegate_nodes, program, DeviceType.CUDA, case, "Delegate"
                 )
                 self._assert_nodes_buffer_device(
-                    nodes.d2h_nodes,
-                    program,
-                    DeviceType.CPU,
-                    f"skip_d2h={skip_d2h}",
-                    "D2H copy",
+                    nodes.d2h_nodes, program, DeviceType.CPU, case, "D2H copy"
                 )
 
     def test_skip_d2h_keeps_output_also_read_on_device(self):
@@ -1753,6 +1755,171 @@ def forward(self, a):
                 ):
                     edge.to_executorch(et_config)
 
+    def test_skip_h2d_for_input_shared_by_device_delegates(self):
+        """With skip_h2d_for_method_inputs=True, a method input read by two
+        delegates on the same device is placed on that device once, with no
+        _h2d_copy."""
+
+        class Model(torch.nn.Module):
+            def forward(self, a, b):
+                c = torch.sin(torch.add(a, b))
+                return torch.add(a, c)
+
+        model = Model()
+        inputs = (torch.randn(2, 2), torch.randn(2, 2))
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        for pipeline, program, gm in self._get_executorch_program(
+            model, inputs, DeviceAwarePartitioner("cuda:0"), et_config
+        ):
+            with self.subTest(pipeline=pipeline):
+                nodes = _collect_device_copy_nodes(gm)
+                self.assertEqual(len(nodes.delegate_nodes), 2)
+                # Only sin's output, computed on the host, is copied.
+                self.assertEqual(len(nodes.h2d_nodes), 1)
+                cuda_ph, _ = self._collect_placeholders_by_device(gm)
+                self.assertEqual(sorted(n.name for n in cuda_ph), ["a", "b"])
+                self._assert_nodes_buffer_device(
+                    cuda_ph, program, DeviceType.CUDA, pipeline, "Placeholder"
+                )
+
+    def test_skip_h2d_refuses_input_read_off_device(self):
+        """A method input cannot move to a device when a host op, or a delegate
+        on another device, also reads it."""
+
+        class HostOp(torch.nn.Module):
+            def forward(self, a, b):
+                return torch.add(a, b), torch.sin(a)
+
+        class TwoDelegates(torch.nn.Module):
+            def forward(self, a, b, c):
+                return torch.add(a, b), torch.add(a, c)
+
+        class CopiedIntoBuffer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(2, 2))
+
+            def forward(self, a, b):
+                self.state.copy_(a)
+                return torch.add(a, b)
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        cases = {
+            "host op": (
+                HostOp(),
+                2,
+                DeviceAwarePartitioner("cuda:0"),
+                "aten_sin_default",
+            ),
+            "other device index": (
+                TwoDelegates(),
+                3,
+                PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1"]),
+                "executorch_call_delegate_1",
+            ),
+            "other device type": (
+                TwoDelegates(),
+                3,
+                PerAddDeviceDelegatePartitioner(["cuda:0", "cpu"]),
+                "executorch_call_delegate_1",
+            ),
+            # The write-back of the buffer only reads the input.
+            "buffer write-back": (
+                CopiedIntoBuffer(),
+                2,
+                DeviceAwarePartitioner("cuda:0"),
+                "copy_",
+            ),
+        }
+        for case, (model, num_inputs, partitioner, reader) in cases.items():
+            ep = export(model, tuple(torch.randn(2, 2) for _ in range(num_inputs)))
+            # Each pipeline gets its own check, so the first refusal cannot hide
+            # the second pipeline.
+            for pipeline in ("to_edge+to_backend", "to_edge_transform_and_lower"):
+                if pipeline == "to_edge+to_backend":
+                    edge = to_edge(
+                        deepcopy(ep),
+                        compile_config=EdgeCompileConfig(_check_ir_validity=False),
+                    ).to_backend(partitioner)
+                else:
+                    edge = to_edge_transform_and_lower(
+                        deepcopy(ep), partitioner=[partitioner]
+                    )
+                with self.subTest(case=case, pipeline=pipeline):
+                    with self.assertRaisesRegex(
+                        RuntimeError,
+                        f"placeholder 'a' on cuda:0, because '{reader}' also reads it",
+                    ):
+                        edge.to_executorch(deepcopy(et_config))
+
+    def test_skip_h2d_refuses_placeholder_written_on_host(self):
+        """A mutated input or mutable buffer is written back on the host, so the
+        refusal says so, also when a host op reads it too."""
+
+        class MutableBuffer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(2, 2))
+
+            def forward(self, a):
+                out = torch.add(a, self.state)
+                self.state.add_(a)
+                return out
+
+        class MutableBufferReadOnHost(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(2, 2))
+
+            def forward(self, a):
+                out = torch.add(a, self.state)
+                read = torch.sin(self.state)
+                self.state.add_(a)
+                return out, read
+
+        class MutatedInput(torch.nn.Module):
+            def forward(self, a, x):
+                out = torch.add(a, x)
+                x.add_(a)
+                return out
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        cases = {
+            "mutable buffer": (MutableBuffer(), 1, "b_state"),
+            "mutable buffer read on host": (MutableBufferReadOnHost(), 1, "b_state"),
+            "mutated input": (MutatedInput(), 2, "x"),
+        }
+        for case, (model, num_inputs, name) in cases.items():
+            inputs = tuple(torch.randn(2, 2) for _ in range(num_inputs))
+            edge = to_edge(export(model, inputs)).to_backend(
+                DeviceAwarePartitioner("cuda:0")
+            )
+            with self.subTest(case=case):
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    f"placeholder '{name}' on cuda:0, because 'copy_' writes it "
+                    r"back on the host \(the method mutates it\)",
+                ):
+                    edge.to_executorch(deepcopy(et_config))
+
     def test_tensorspec_repr_includes_device(self):
         spec = TensorSpec(dtype=torch.float32, shape=torch.Size([2, 3]))
         repr_str = repr(spec)

From fd1ac9ee76c9a937f13b35eb0570e19077bcd49d Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 23 Oct 2026 11:59:40 -0700
Subject: [PATCH 07/26] Copy a tensor to a device once when several delegates
 there read it

When several delegates on one device read the same tensor, the pass made a
separate host-to-device copy for each of them. A constant read by three
delegates on cuda:0, such as a lookup table, was copied three times on every
call. Method inputs were copied the same way when skip_h2d_for_method_inputs
is off.

The pass visits delegates in graph order, so a copy made for an earlier
delegate on the same device already comes before the later ones. It now
reuses that copy. A tensor read on two devices still gets one copy for each
device. The warning for a constant copied on every call is logged once per
copy.

The reused copy stays in device memory from the first delegate that reads it
to the last, so memory planning cannot give that space to other tensors in
between. That is the cost of saving the copy on every call.

A tensor that any op writes in place, such as an argument of a mutable
custom op or a mutable buffer with its write-back, is not reused, because
an earlier copy can hold the value from before the write. The check does
not look at where the write happens, so it also covers writes before the
first reader or after the last one, where reuse would be safe. Those
tensors get one copy for each delegate, as before.

For example, a constant read by delegates on cuda:0, cuda:1 and cuda:0 now
gets two copies instead of three.
---
 docs/source/backends/cuda/cuda-overview.md |   5 +-
 exir/passes/propagate_device_config.py     |   3 +-
 exir/passes/propagate_device_pass.py       |  80 ++++++++++------
 exir/tests/test_propagate_device_pass.py   | 101 +++++++++++++++++++--
 4 files changed, 151 insertions(+), 38 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index d0f9568fa99..ba681b84dea 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -181,8 +181,9 @@ once those operators are there, and the CUDA partitioner already calls it.
 **A constant that a delegate reads from outside is still copied.** Constants, such as weights and
 buffers the method does not change, are not passed by the caller. Their data stays in host memory,
 in the program or in its external data file, so `to_executorch` adds a copy to the device that runs
-on every call, one for each delegate that reads the constant, and logs a warning. Keep constants
-inside the delegate to avoid that cost.
+on every call, one for each device that reads the constant, and logs a warning. Each copy stays in
+device memory until the last delegate on that device has read it. Keep constants inside the
+delegate to avoid that cost.
 
 **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the
 program reserves its own input buffer and the runtime copies the caller's memory into it, so the
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index 54d51851db9..ab2cfe1d8f0 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -53,7 +53,8 @@ class PropagateDeviceConfig:
     # where inputs are pre-staged on GPU. Constants are not passed by the
     # caller, so they are still copied: their data stays on the host, in the
     # program or in its external data file. That copy runs on every call, once
-    # for each delegate that reads the constant.
+    # for each device that reads the constant, and stays in device memory until
+    # the last delegate there has read it.
     #
     # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer
     # for the input and copies the caller's memory into it, which is what this exists to avoid.
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index f6451b96bed..90f6f7b88cc 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -193,8 +193,12 @@ class PropagateDevicePass(PassBase):
     boundaries based on partitioner-assigned device info.
 
     When a delegate has a target_device CompileSpec (e.g., "cuda:0"):
-    - For each delegate input: insert et_copy._h2d_copy before the delegate call.
-      The original input node stays CPU; the h2d_copy output is tagged as device.
+    - For each delegate input: insert et_copy._h2d_copy before the delegate call,
+      or reuse the one an earlier delegate on the same device reads when no op
+      writes the input in place. A reused copy stays in device memory until its
+      last reader.
+      The original input node keeps its device; the h2d_copy output is tagged as
+      device.
     - For each delegate output: insert et_copy._d2h_copy after each getitem,
       or after the call itself when it returns a single tensor.
       The delegate output stays device; the d2h_copy output is tagged as CPU.
@@ -257,7 +261,13 @@ def _insert_h2d_copies(
         target_device_type: schema.DeviceType,
         device_index: int,
     ) -> bool:
-        """Insert H2D copy nodes for each tensor input to a delegate call."""
+        """Make each tensor input of a delegate call read device memory.
+
+        A method input that is not a constant moves to the device under
+        skip_h2d_for_method_inputs.
+        Any other input reads an H2D copy, reused from an earlier delegate on
+        the same device when possible.
+        """
         changed = False
         new_args = list(node.args)
         for i, arg in enumerate(node.args[1:], start=1):
@@ -328,30 +338,47 @@ def _insert_h2d_copies(
                     f"match the stored data. Make the constant contiguous where "
                     f"it is created, for example in the pass that adds it."
                 )
-            if arg_spec.const:
-                logger.warning(
-                    "Constant '%s' (%d bytes) feeds a delegate on %s:%d from "
-                    "outside it, so it is copied to the device on every call. "
-                    "Have the partitioner take it into the delegate to avoid "
-                    "that copy.",
-                    arg.name,
-                    arg_spec.nbytes(),
-                    target_device_type.name.lower(),
-                    device_index,
-                )
 
-            with graph_module.graph.inserting_before(node):
-                h2d_node = graph_module.graph.call_function(
-                    torch.ops.et_copy._h2d_copy.default,
-                    (arg,),
+            # Delegates are visited in graph order, so a copy made for an earlier
+            # delegate on this device already comes before this one. It can be
+            # stale when an op writes the tensor in place.
+            h2d_node = None
+            if not _is_written_in_place(arg):
+                h2d_node = next(
+                    (
+                        user
+                        for user in arg.users
+                        if user.target == torch.ops.et_copy._h2d_copy.default
+                        and user.meta["spec"].device == target_device_type
+                        and user.meta["spec"].device_index == device_index
+                    ),
+                    None,
                 )
-                h2d_spec = _copy_output_spec(arg_spec, target_device_type, device_index)
-                h2d_node.meta["spec"] = h2d_spec
-                h2d_node.meta["val"] = arg.meta.get("val")
-                if "tensor_meta" in arg.meta:
-                    h2d_node.meta["tensor_meta"] = arg.meta["tensor_meta"]
-                new_args[i] = h2d_node
-                changed = True
+            if h2d_node is None:
+                with graph_module.graph.inserting_before(node):
+                    h2d_node = graph_module.graph.call_function(
+                        torch.ops.et_copy._h2d_copy.default,
+                        (arg,),
+                    )
+                    h2d_node.meta["spec"] = _copy_output_spec(
+                        arg_spec, target_device_type, device_index
+                    )
+                    h2d_node.meta["val"] = arg.meta.get("val")
+                    if "tensor_meta" in arg.meta:
+                        h2d_node.meta["tensor_meta"] = arg.meta["tensor_meta"]
+                if arg_spec.const:
+                    logger.warning(
+                        "Constant '%s' (%d bytes) feeds a delegate on %s:%d from "
+                        "outside it, so it is copied to the device on every call. "
+                        "Have the partitioner take it into the delegate to avoid "
+                        "that copy.",
+                        arg.name,
+                        arg_spec.nbytes(),
+                        target_device_type.name.lower(),
+                        device_index,
+                    )
+            new_args[i] = h2d_node
+            changed = True
 
         node.args = tuple(new_args)
         return changed
@@ -519,7 +546,8 @@ def _elide_redundant_same_device_copies(
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         # Three-pass approach:
         #   Pass 1 – For each delegate with a target_device CompileSpec, insert
-        #            H2D copy nodes before delegate inputs and tag the delegate
+        #            H2D copy nodes before delegate inputs, or reuse the copy an
+        #            earlier delegate on the same device reads, and tag the delegate
         #            output specs with the target device.  Delegates without a
         #            target_device are left untouched (no copies, specs stay CPU).
         #   Pass 2 – For each getitem that extracts from a device-tagged delegate
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 5c1df35b19e..89ddf099601 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -660,9 +660,9 @@ def forward(self, x, y):
             with self.subTest(case=case):
                 nodes = _collect_device_copy_nodes(gm)
                 self.assertEqual(len(nodes.delegate_nodes), 2)
-                # Only method inputs are copied to the device: x and y, then y
-                # again, unless the caller places them.
-                self.assertEqual(len(nodes.h2d_nodes), 0 if skip_h2d else 3)
+                # Only method inputs are copied to the device, x and y once
+                # each, unless the caller places them.
+                self.assertEqual(len(nodes.h2d_nodes), 0 if skip_h2d else 2)
                 self.assertTrue(
                     all(n.args[0].op == "placeholder" for n in nodes.h2d_nodes)
                 )
@@ -921,14 +921,14 @@ def forward(self, a, b):
                     4,
                     f"[{pipeline}] expected four delegates",
                 )
-                # a -> {D0, D1} and b -> {D0, D2} give 4 graph-input H2D copies;
-                # the single graph output gives 1 D2H. All inter-delegate
-                # round-trips (including delegate 0's output feeding D1 and D2)
-                # are elided.
+                # a -> {D0, D1} and b -> {D0, D2} give one graph-input H2D copy
+                # each, shared by the delegates that read it; the single graph
+                # output gives 1 D2H. All inter-delegate round-trips (including
+                # delegate 0's output feeding D1 and D2) are elided.
                 self.assertEqual(
                     len(nodes.h2d_nodes),
-                    4,
-                    f"[{pipeline}] expected 4 boundary H2D copies, got "
+                    2,
+                    f"[{pipeline}] expected 2 boundary H2D copies, got "
                     f"{len(nodes.h2d_nodes)}",
                 )
                 self.assertEqual(
@@ -1755,6 +1755,89 @@ def forward(self, a):
                 ):
                     edge.to_executorch(et_config)
 
+    def test_constant_read_by_several_delegates_is_copied_once_per_device(self):
+        """A constant read by several delegates gets one _h2d_copy for each
+        device, not one for each delegate."""
+
+        inputs = (torch.randn(2, 2), torch.randn(2, 2))
+        lowered = {
+            device: _lower_add_to_device(inputs, device)
+            for device in ("cuda:0", "cuda:1")
+        }
+
+        class Model(torch.nn.Module):
+            def __init__(self, devices):
+                super().__init__()
+                self.register_buffer("offset", torch.randn(2, 2))
+                self.first, self.second, self.third = (lowered[d] for d in devices)
+
+            def forward(self, x):
+                return (
+                    self.first(x, self.offset),
+                    self.second(x, self.offset),
+                    self.third(x, self.offset),
+                )
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False, enable_non_cpu_memory_planning=True
+        )
+        cases = {
+            "one device": (("cuda:0", "cuda:0", "cuda:0"), [0]),
+            "two devices": (("cuda:0", "cuda:1", "cuda:0"), [0, 1]),
+        }
+        for case, (devices, copy_indices) in cases.items():
+            with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs:
+                et = to_edge(export(Model(devices), inputs[:1])).to_executorch(
+                    deepcopy(et_config)
+                )
+            gm = et.exported_program().graph_module
+            with self.subTest(case=case):
+                self.assertEqual(len(logs.output), len(copy_indices))
+                nodes = _collect_device_copy_nodes(gm)
+                offset_copies = [
+                    n for n in nodes.h2d_nodes if n.args[0].name == "b_offset"
+                ]
+                self.assertEqual(
+                    sorted(n.meta["spec"].device_index for n in offset_copies),
+                    copy_indices,
+                )
+                for delegate, device in zip(nodes.delegate_nodes, devices):
+                    offset = delegate.args[2]
+                    self.assertIn(offset, offset_copies)
+                    self.assertEqual(f"cuda:{offset.meta['spec'].device_index}", device)
+
+    def test_tensor_written_on_host_is_copied_for_each_delegate(self):
+        """A tensor that a host op writes in place between two delegates on one
+        device gets a new copy for the second delegate, so it reads the new
+        value."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                t = x * 2
+                a = t + y
+                torch.ops.propagate_device_test.accumulate(t, a)
+                return a, t + y
+
+        gm = _lower_with_host_writer(
+            Model(),
+            (torch.randn(2, 2), torch.randn(2, 2)),
+            ["cuda:0", "cuda:0"],
+            ExecutorchBackendConfig(
+                emit_stacktrace=False, enable_non_cpu_memory_planning=True
+            ),
+        )
+        writer = _host_writer(gm)
+        first, second = _collect_device_copy_nodes(gm).delegate_nodes
+        (second_copy,) = [
+            arg
+            for arg in second.args[1:]
+            if arg.target == torch.ops.et_copy._h2d_copy.out
+            and arg.args[0] is writer.args[0]
+        ]
+        nodes = list(gm.graph.nodes)
+        self.assertLess(nodes.index(first), nodes.index(writer))
+        self.assertLess(nodes.index(writer), nodes.index(second_copy))
+
     def test_skip_h2d_for_input_shared_by_device_delegates(self):
         """With skip_h2d_for_method_inputs=True, a method input read by two
         delegates on the same device is placed on that device once, with no

From 9b508c796c5b0aa7569a9cbb902b5a41ec270124 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 23 Oct 2026 23:41:35 -0700
Subject: [PATCH 08/26] Keep in-place results on the host copy of a device
 delegate output

SpecPropPass gives the result of an in-place op the TensorSpec of the
tensor it writes, so memory planning puts both in one buffer. When a host
op writes a device delegate output in place, for example relu_ or
index_put_ made by reinplace_pass, PropagateDevicePass moves the op to the
host copy of the output, but the result kept the spec of the device
output. Memory planning then refused the program:

  In-place node aten_relu__default ... result 0 does not share the
  TensorSpec of its mutated input _d2h_copy_default.

This happened both for an output read through getitem and for a single
tensor returned by the call.

When the pass moves readers to the host copy, it already pointed the graph
output at the copy's spec. It now does the same for in-place results and
their own readers, so they share the copy's spec. A returned in-place
result is the result of a host op, so the method returns it from the host,
also with skip_d2h_for_method_outputs.

The new test runs relu_ and then mul_ on a delegate output, for getitem
and for a single tensor, with skip_d2h_for_method_outputs off and on. It
fails without this change.
---
 exir/passes/propagate_device_pass.py     | 33 ++++++++-----
 exir/tests/test_propagate_device_pass.py | 61 ++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 12 deletions(-)

diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 90f6f7b88cc..bb85224fee5 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -171,20 +171,25 @@ def _is_written_in_place(node: torch.fx.Node) -> bool:
     )
 
 
-def _replace_output_spec(
-    node: torch.fx.Node, old_spec: TensorSpec, new_spec: TensorSpec
+def _replace_spec(
+    readers: list[torch.fx.Node], old_spec: TensorSpec, new_spec: TensorSpec
 ) -> None:
-    """Replace *old_spec* with *new_spec* in the spec of the output node reading *node*.
+    """Replace *old_spec* with *new_spec* in *readers*, which now read the
+    tensor with *new_spec*, and in the readers of their results.
 
-    Memory planning finds the graph outputs through the output node's spec, so
-    it has to name the tensor that is actually returned.
+    SpecPropPass gives the output node the spec of each tensor it returns, and
+    an in-place result the spec of the tensor it writes. Memory planning relies
+    on both, so these specs have to follow the tensor that is read now.
     """
-    for user in node.users:
-        if user.op == "output" and "spec" in user.meta:
-            user.meta["spec"] = pytree.tree_map(
-                lambda value: new_spec if value is old_spec else value,
-                user.meta["spec"],
+    pending = list(readers)
+    while pending:
+        node = pending.pop()
+        spec = node.meta.get("spec")
+        if any(value is old_spec for value in pytree.tree_leaves(spec)):
+            node.meta["spec"] = pytree.tree_map(
+                lambda value: new_spec if value is old_spec else value, spec
             )
+            pending.extend(node.users)
 
 
 class PropagateDevicePass(PassBase):
@@ -439,7 +444,7 @@ def _insert_d2h_after(
                 d2h_node,
                 delete_user_cb=lambda user, _d2h=d2h_node: user != _d2h,
             )
-            _replace_output_spec(d2h_node, spec, d2h_spec)
+            _replace_spec(list(d2h_node.users), spec, d2h_spec)
 
     def _keep_method_outputs_on_device(
         self,
@@ -474,7 +479,11 @@ def _keep_method_outputs_on_device(
                 returned.meta["val"] = node.meta.get("val")
                 if "tensor_meta" in node.meta:
                     returned.meta["tensor_meta"] = node.meta["tensor_meta"]
-            _replace_output_spec(node, node.meta["spec"], returned.meta["spec"])
+            _replace_spec(
+                [graph_module.graph.output_node()],
+                node.meta["spec"],
+                returned.meta["spec"],
+            )
             node.replace_all_uses_with(
                 returned, delete_user_cb=lambda user: user.op == "output"
             )
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 89ddf099601..98231e59437 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -686,6 +686,67 @@ def forward(self, x, y):
                     nodes.d2h_nodes, program, DeviceType.CPU, case, "D2H copy"
                 )
 
+    def test_host_inplace_ops_on_delegate_output_write_host_copy(self):
+        """Host ops that write a delegate output in place, here relu_ and then
+        mul_ from reinplace_pass, write its host copy. Their results share the
+        copy's spec, as memory planning requires, and the method returns the
+        last one from the host."""
+
+        inputs = (torch.randn(2, 2), torch.randn(2, 2))
+        lowered = _lower_add_to_device(inputs, "cuda:0")
+
+        class Call(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.lowered = lowered
+
+            def forward(self, x, y):
+                return torch.relu(self.lowered(x, y)) * 2
+
+        class Add(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.relu(x + y) * 2
+
+        cases = {
+            "single tensor": (Call(), None),
+            "getitem": (Add(), PerAddDeviceDelegatePartitioner(["cuda:0"])),
+        }
+        reinplaced = frozenset(
+            {exir_ops.edge.aten.relu.default, exir_ops.edge.aten.mul.Tensor}
+        )
+        for case, (model, partitioner) in cases.items():
+            for skip_d2h in (False, True):
+                edge = to_edge(export(model, inputs))
+                if partitioner is not None:
+                    edge = edge.to_backend(partitioner)
+                et = edge.to_executorch(
+                    ExecutorchBackendConfig(
+                        emit_stacktrace=False,
+                        reinplace_extra_ops=reinplaced,
+                        propagate_device_config=PropagateDeviceConfig(
+                            skip_d2h_for_method_outputs=skip_d2h
+                        ),
+                        enable_non_cpu_memory_planning=True,
+                    )
+                )
+                gm = et.exported_program().graph_module
+                with self.subTest(case=case, skip_d2h=skip_d2h):
+                    (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes
+                    (relu,) = d2h.users
+                    (mul,) = relu.users
+                    self.assertEqual(relu.target, exir_ops.edge.aten.relu_.default)
+                    self.assertEqual(mul.target, exir_ops.edge.aten.mul_.Tensor)
+                    self.assertIs(relu.meta["spec"], d2h.meta["spec"])
+                    self.assertIs(mul.meta["spec"], d2h.meta["spec"])
+                    self.assertEqual(gm.graph.output_node().args[0], (mul,))
+                    self._assert_nodes_buffer_device(
+                        [mul],
+                        et._emitter_output.program,
+                        DeviceType.CPU,
+                        case,
+                        "Output",
+                    )
+
     def test_skip_d2h_keeps_output_also_read_on_device(self):
         """With skip_d2h_for_method_outputs, a returned delegate output that a
         delegate on the same device also reads stays on the device, like the

From ba5a303c269877c751cdb6f855510d6150eb95c3 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 23 Oct 2026 23:49:33 -0700
Subject: [PATCH 09/26] Count a host write through a view as a write to the
 tensor

Before PropagateDevicePass reuses a device copy of a tensor, removes a
device -> host -> device round-trip, or returns a delegate output on the
device instead of its host copy, it checks that no host op writes that
tensor in place. The check only looked at ops that write the tensor itself.

A mutable custom op on the host can also write a view of the tensor. When
to_executorch removes auto_functionalized, it makes that view again, for
example with aten.as_strided, and the op writes the view. The check missed
this, so a later delegate, or the method output, could get the value from
before the write. For example:

  t = x * y
  a = t + y                           # delegate on cuda:0
  accumulate(t.view(-1), y.view(-1))  # host op, writes t
  return a, t + y                     # delegate on cuda:0

Here the second delegate read the device copy of t made before accumulate.

The check now also counts writes to views of the tensor, and to views of
those, as the op schemas declare them. The result of an in-place op counts
as a view of its input. A view_copy does not, because it still makes its
own tensor at this point, and ReplaceViewCopyWithViewPass turns it into a
view later only where no write can tell the difference.

The tests for round-trip removal, returned outputs and copy reuse each get
a case that writes through a view. These cases fail without this change.
aten.as_strided has no out variant yet, so in these cases to_executorch is
told to keep it.
---
 exir/passes/propagate_device_pass.py     |  51 +++++--
 exir/tests/test_propagate_device_pass.py | 184 ++++++++++++++---------
 2 files changed, 158 insertions(+), 77 deletions(-)

diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index bb85224fee5..e0aea8f3a63 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -20,6 +20,9 @@
 from executorch.exir.lowered_backend_module import LoweredBackendModule
 from executorch.exir.passes.convert_constant_dim_order_pass import _should_transform
 from executorch.exir.passes.insert_write_back_for_buffers_pass import (
+    _alias_sets,
+    _contains_node,
+    _schema_paired_args,
     _schema_writes_input,
 )
 
@@ -156,19 +159,49 @@ def _copy_output_spec(
     return new_spec
 
 
+def _schema_aliases_input(
+    node: torch.fx.Node, schema: torch.FunctionSchema, input_node: torch.fx.Node
+) -> bool:
+    """Whether schema marks the result of node as an alias of input_node, as for
+    a view or an in-place op."""
+    returned = set().union(*(_alias_sets(ret.alias_info) for ret in schema.returns))
+    return any(
+        _alias_sets(schema_arg.alias_info) & returned
+        for arg, schema_arg in _schema_paired_args(node, schema)
+        if _contains_node(arg, input_node)
+    )
+
+
 def _is_written_in_place(node: torch.fx.Node) -> bool:
-    """Whether an op that reads *node* also writes it, per the op's schema.
+    """Whether an op writes *node* in place, directly or through a view of it,
+    per the op's schema.
 
     A mutable op left on the host does this once to_executorch has removed
-    auto_functionalized, so a copy of *node* can miss the write. Only direct
-    readers count: a write through a view of *node* is not seen. Custom op
-    schemas are trusted, as functionalization trusts them.
+    auto_functionalized, so a copy of *node* can miss the write. If the op
+    wrote a view of *node* in the model, that removal makes the view again,
+    for example with as_strided. The result of an in-place op is a view of
+    its input too. A view_copy is not: it still makes its own tensor here,
+    and ReplaceViewCopyWithViewPass turns it into a view later only where no
+    write can tell the difference. Custom op schemas are trusted, as
+    functionalization trusts them. The write-back pass trusts only aten
+    schemas, but here that would keep a copy for every custom op on the host,
+    such as a KV cache update. Users without a schema, such as delegate calls
+    and getitem, do not write their inputs.
     """
-    return any(
-        _schema_writes_input(user, user.target._schema, node)  # pyre-ignore[16]
-        for user in node.users
-        if hasattr(user.target, "_schema")
-    )
+    aliases = {node}
+    pending = [node]
+    while pending:
+        alias = pending.pop()
+        for user in alias.users:
+            schema = getattr(user.target, "_schema", None)
+            if schema is None:
+                continue
+            if _schema_writes_input(user, schema, alias):
+                return True
+            if user not in aliases and _schema_aliases_input(user, schema, alias):
+                aliases.add(user)
+                pending.append(user)
+    return False
 
 
 def _replace_spec(
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 98231e59437..f00819a9eab 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -31,7 +31,11 @@
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult
-from executorch.exir.passes import MemoryPlanningPass, propagate_device_pass
+from executorch.exir.passes import (
+    MemoryPlanningPass,
+    propagate_device_pass,
+    ToOutVarPass,
+)
 from executorch.exir.passes.propagate_device_pass import (
     _get_target_device_from_compile_specs,
     _parse_device_spec_value,
@@ -210,6 +214,25 @@ def _host_writer(gm: torch.fx.GraphModule) -> torch.fx.Node:
     return writer
 
 
+def _host_written(gm: torch.fx.GraphModule) -> torch.fx.Node:
+    """The tensor that accumulate writes, looking through the view it may write."""
+    written = _host_writer(gm).args[0]
+    if written.target == torch.ops.aten.as_strided.default:
+        return written.args[0]
+    return written
+
+
+def _host_writer_config(write_view: bool, **kwargs) -> ExecutorchBackendConfig:
+    return ExecutorchBackendConfig(
+        emit_stacktrace=False,
+        enable_non_cpu_memory_planning=True,
+        # A write through a view leaves the aten.as_strided that
+        # to_executorch makes for it, which has no out variant.
+        to_out_var_pass=ToOutVarPass(ignore_to_out_var_failure=write_view),
+        **kwargs,
+    )
+
+
 class TestPropagateDevicePass(unittest.TestCase):
     @staticmethod
     def _collect_tensor_specs(node: torch.fx.Node) -> List[TensorSpec]:
@@ -562,35 +585,45 @@ def forward(self, a, b, c):
                 )
 
     def test_same_device_copies_kept_when_host_op_writes_copy(self):
-        """A host op that writes the host copy of a delegate output in place
-        changes what the next delegate on the same device reads, so that
-        round-trip stays."""
+        """A host op that writes the host copy of a delegate output in place,
+        directly or through a view, changes what the next delegate on the same
+        device reads, so that round-trip stays. The input it only reads still
+        gets one copy for both delegates."""
 
         class Model(torch.nn.Module):
+            def __init__(self, write_view):
+                super().__init__()
+                self.write_view = write_view
+
             def forward(self, x, y):
                 a = x + y
-                torch.ops.propagate_device_test.accumulate(a, y)
+                if self.write_view:
+                    torch.ops.propagate_device_test.accumulate(a.view(-1), y.view(-1))
+                else:
+                    torch.ops.propagate_device_test.accumulate(a, y)
                 return a + y
 
-        gm = _lower_with_host_writer(
-            Model(),
-            (torch.randn(2, 2), torch.randn(2, 2)),
-            ["cuda:0", "cuda:0"],
-            ExecutorchBackendConfig(
-                emit_stacktrace=False, enable_non_cpu_memory_planning=True
-            ),
-        )
-        written = _host_writer(gm).args[0]
-        self.assertEqual(written.target, torch.ops.et_copy._d2h_copy.out)
-        second = _collect_device_copy_nodes(gm).delegate_nodes[1]
-        self.assertIn(
-            written,
-            [
-                arg.args[0]
-                for arg in second.args[1:]
-                if arg.target == torch.ops.et_copy._h2d_copy.out
-            ],
-        )
+        for write_view in (False, True):
+            gm = _lower_with_host_writer(
+                Model(write_view),
+                (torch.randn(2, 2), torch.randn(2, 2)),
+                ["cuda:0", "cuda:0"],
+                _host_writer_config(write_view),
+            )
+            with self.subTest(write_view=write_view):
+                written = _host_written(gm)
+                self.assertEqual(written.target, torch.ops.et_copy._d2h_copy.out)
+                nodes = _collect_device_copy_nodes(gm)
+                self.assertIn(
+                    written,
+                    [
+                        arg.args[0]
+                        for arg in nodes.delegate_nodes[1].args[1:]
+                        if arg.target == torch.ops.et_copy._h2d_copy.out
+                    ],
+                )
+                y_copies = [n for n in nodes.h2d_nodes if n.args[0].name == "y"]
+                self.assertEqual(len(y_copies), 1)
 
     def test_single_tensor_delegate_output_read_on_host(self):
         """A host op that reads the output of a lowered module returning one
@@ -883,34 +916,42 @@ def forward(self, x, y):
 
     def test_skip_d2h_returns_device_copy_of_output_written_on_host(self):
         """With skip_d2h_for_method_outputs, a returned delegate output whose
-        host copy a host op writes in place is returned as a device copy of the
-        written host copy, because the device tensor never sees the write."""
+        host copy a host op writes in place, directly or through a view, is
+        returned as a device copy of the written host copy, because the device
+        tensor never sees the write."""
 
         class Model(torch.nn.Module):
+            def __init__(self, write_view):
+                super().__init__()
+                self.write_view = write_view
+
             def forward(self, x, y):
                 a = x + y
-                torch.ops.propagate_device_test.accumulate(a, y)
+                if self.write_view:
+                    torch.ops.propagate_device_test.accumulate(a.view(-1), y.view(-1))
+                else:
+                    torch.ops.propagate_device_test.accumulate(a, y)
                 return a
 
-        gm = _lower_with_host_writer(
-            Model(),
-            (torch.randn(2, 2), torch.randn(2, 2)),
-            ["cuda:0"],
-            ExecutorchBackendConfig(
-                emit_stacktrace=False,
-                propagate_device_config=PropagateDeviceConfig(
-                    skip_d2h_for_method_outputs=True
+        for write_view in (False, True):
+            gm = _lower_with_host_writer(
+                Model(write_view),
+                (torch.randn(2, 2), torch.randn(2, 2)),
+                ["cuda:0"],
+                _host_writer_config(
+                    write_view,
+                    propagate_device_config=PropagateDeviceConfig(
+                        skip_d2h_for_method_outputs=True
+                    ),
                 ),
-                enable_non_cpu_memory_planning=True,
-            ),
-        )
-        (output,) = gm.graph.output_node().args[0]
-        self.assertEqual(output.target, torch.ops.et_copy._h2d_copy.out)
-        self.assertEqual(output.meta["spec"].device, DeviceType.CUDA)
-        writer = _host_writer(gm)
-        self.assertIs(output.args[0], writer.args[0])
-        nodes = list(gm.graph.nodes)
-        self.assertLess(nodes.index(writer), nodes.index(output))
+            )
+            with self.subTest(write_view=write_view):
+                (output,) = gm.graph.output_node().args[0]
+                self.assertEqual(output.target, torch.ops.et_copy._h2d_copy.out)
+                self.assertEqual(output.meta["spec"].device, DeviceType.CUDA)
+                self.assertIs(output.args[0], _host_written(gm))
+                nodes = list(gm.graph.nodes)
+                self.assertLess(nodes.index(_host_writer(gm)), nodes.index(output))
 
     def test_cross_device_copies_preserved(self):
         """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two
@@ -1868,36 +1909,43 @@ def forward(self, x):
                     self.assertEqual(f"cuda:{offset.meta['spec'].device_index}", device)
 
     def test_tensor_written_on_host_is_copied_for_each_delegate(self):
-        """A tensor that a host op writes in place between two delegates on one
-        device gets a new copy for the second delegate, so it reads the new
-        value."""
+        """A tensor that a host op writes in place, directly or through a view,
+        between two delegates on one device gets a new copy for the second
+        delegate, so it reads the new value."""
 
         class Model(torch.nn.Module):
+            def __init__(self, write_view):
+                super().__init__()
+                self.write_view = write_view
+
             def forward(self, x, y):
                 t = x * 2
                 a = t + y
-                torch.ops.propagate_device_test.accumulate(t, a)
+                if self.write_view:
+                    torch.ops.propagate_device_test.accumulate(t.view(-1), a.view(-1))
+                else:
+                    torch.ops.propagate_device_test.accumulate(t, a)
                 return a, t + y
 
-        gm = _lower_with_host_writer(
-            Model(),
-            (torch.randn(2, 2), torch.randn(2, 2)),
-            ["cuda:0", "cuda:0"],
-            ExecutorchBackendConfig(
-                emit_stacktrace=False, enable_non_cpu_memory_planning=True
-            ),
-        )
-        writer = _host_writer(gm)
-        first, second = _collect_device_copy_nodes(gm).delegate_nodes
-        (second_copy,) = [
-            arg
-            for arg in second.args[1:]
-            if arg.target == torch.ops.et_copy._h2d_copy.out
-            and arg.args[0] is writer.args[0]
-        ]
-        nodes = list(gm.graph.nodes)
-        self.assertLess(nodes.index(first), nodes.index(writer))
-        self.assertLess(nodes.index(writer), nodes.index(second_copy))
+        for write_view in (False, True):
+            gm = _lower_with_host_writer(
+                Model(write_view),
+                (torch.randn(2, 2), torch.randn(2, 2)),
+                ["cuda:0", "cuda:0"],
+                _host_writer_config(write_view),
+            )
+            with self.subTest(write_view=write_view):
+                writer = _host_writer(gm)
+                first, second = _collect_device_copy_nodes(gm).delegate_nodes
+                (second_copy,) = [
+                    arg
+                    for arg in second.args[1:]
+                    if arg.target == torch.ops.et_copy._h2d_copy.out
+                    and arg.args[0] is _host_written(gm)
+                ]
+                nodes = list(gm.graph.nodes)
+                self.assertLess(nodes.index(first), nodes.index(writer))
+                self.assertLess(nodes.index(writer), nodes.index(second_copy))
 
     def test_skip_h2d_for_input_shared_by_device_delegates(self):
         """With skip_h2d_for_method_inputs=True, a method input read by two

From 0160993e3f4953558d6c0ad1fb2b492b08bad75c Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 24 Oct 2026 00:39:09 -0700
Subject: [PATCH 10/26] Test tensors shared by delegates on different devices

The device pass tests covered a delegate output read by delegates on two
devices and a constant read on two devices. A method input read by
delegates on different devices was only covered with
skip_h2d_for_method_inputs, where it is refused.

The new test has delegates on cuda:0, cuda:1 and cuda:0 that read the
same method input. The second and third each read the output of the
delegate before them, which is on the other device. The input gets one
copy for each device, and both delegates on cuda:0 read the same one.
Each output crosses to the other device through the host. The test fails
if copy reuse ignores the device index.
---
 exir/tests/test_propagate_device_pass.py | 41 ++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index f00819a9eab..c0d26cb00c2 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -1182,6 +1182,47 @@ def forward(self, a, b):
                     f"first delegate's output directly",
                 )
 
+    def test_tensors_shared_by_delegates_on_different_devices(self):
+        """Delegates on cuda:0 and cuda:1 that read the same method input and
+        each other's outputs. The input gets one copy for each device, read by
+        every delegate there, and each output crosses to the other device
+        through the host."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                a = x + y  # delegate 0 (cuda:0)
+                b = x + a  # delegate 1 (cuda:1)
+                return a, b, x + b  # delegate 2 (cuda:0)
+
+        d2h_out = torch.ops.et_copy._d2h_copy.out
+        for pipeline, gm in _lower_model_to_executorch(
+            Model(),
+            (torch.randn(2, 2), torch.randn(2, 2)),
+            PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1", "cuda:0"]),
+            ExecutorchBackendConfig(enable_non_cpu_memory_planning=True),
+        ):
+            with self.subTest(pipeline=pipeline):
+                nodes = _collect_device_copy_nodes(gm)
+                first, second, third = nodes.delegate_nodes
+                x_copies = [n for n in nodes.h2d_nodes if n.args[0].name == "x"]
+                first_x, second_x, third_x = (
+                    next(arg for arg in delegate.args[1:] if arg in x_copies)
+                    for delegate in (first, second, third)
+                )
+                self.assertEqual(len(x_copies), 2)
+                self.assertIs(first_x, third_x)
+                self.assertEqual(first_x.meta["spec"].device_index, 0)
+                self.assertEqual(second_x.meta["spec"].device_index, 1)
+                crossings = sorted(
+                    (
+                        h2d.args[0].args[0].meta["spec"].device_index,
+                        h2d.meta["spec"].device_index,
+                    )
+                    for h2d in nodes.h2d_nodes
+                    if h2d.args[0].target == d2h_out
+                )
+                self.assertEqual(crossings, [(0, 1), (1, 0)])
+
         # ---- Integration tests: device consistency after to_executorch ----
 
     def test_device_consistency_cuda_1(self):

From b10cc1e46ef19fab81b5b656942d6abc9b9180a3 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 24 Oct 2026 00:39:19 -0700
Subject: [PATCH 11/26] Test device placement end to end with CUDA delegates

The device pass fixes were only tested on the graph and on memory
planning. This adds tests that lower small models to real CUDA delegates,
run them, and compare the results with eager PyTorch. Each add becomes its
own delegate, and parameters and buffers stay outside the delegates, so
the program copies them to the device.

- A transposed parameter. Without the parameter fix, to_executorch fails,
  because the copy would move the contiguous bytes the program stores to a
  delegate compiled for the transposed strides, and the pass refuses that.
- A buffer read by two delegates.

The tests need a CUDA GPU and are skipped otherwise.
---
 backends/cuda/tests/test_device_placement.py | 113 +++++++++++++++++++
 1 file changed, 113 insertions(+)
 create mode 100644 backends/cuda/tests/test_device_placement.py

diff --git a/backends/cuda/tests/test_device_placement.py b/backends/cuda/tests/test_device_placement.py
new file mode 100644
index 00000000000..9bedbf6abff
--- /dev/null
+++ b/backends/cuda/tests/test_device_placement.py
@@ -0,0 +1,113 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+from executorch.backends.cuda.cuda_backend import CudaBackend
+from executorch.exir import (
+    EdgeCompileConfig,
+    ExecutorchBackendConfig,
+    to_edge_transform_and_lower,
+)
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.passes.propagate_device_pass import TARGET_DEVICE_COMPILE_SPEC_KEY
+
+
+class EachAddOnCuda(Partitioner):
+    """Lowers each add to its own CUDA delegate. Parameters stay outside, so
+    they cross the delegate boundary as in a partial lowering."""
+
+    def partition(self, exported_program) -> PartitionResult:
+        spec = DelegationSpec(
+            CudaBackend.__name__,
+            [
+                CudaBackend.generate_method_name_compile_spec("forward"),
+                CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, b"cuda:0"),
+            ],
+        )
+        tags = {}
+        adds = [
+            node
+            for node in exported_program.graph_module.graph.nodes
+            if node.op == "call_function"
+            and node.target == exir_ops.edge.aten.add.Tensor
+        ]
+        for i, node in enumerate(adds):
+            node.meta["delegation_tag"] = f"add{i}"
+            tags[f"add{i}"] = spec
+        return PartitionResult(
+            tagged_exported_program=exported_program, partition_tags=tags
+        )
+
+
+def run(model, inputs) -> list[torch.Tensor]:
+    """Lower model with EachAddOnCuda, run it in the runtime, return its outputs."""
+    from executorch.runtime import Runtime
+
+    edge = to_edge_transform_and_lower(
+        torch.export.export(model, inputs, strict=True),
+        partitioner=[EachAddOnCuda()],
+        compile_config=EdgeCompileConfig(
+            _check_ir_validity=False, _skip_dim_order=True
+        ),
+    )
+    with tempfile.TemporaryDirectory() as directory:
+        path = Path(directory) / "model.pte"
+        path.write_bytes(edge.to_executorch(ExecutorchBackendConfig()).buffer)
+        method = Runtime.get().load_program(path).load_method("forward")
+        return [output.cpu() for output in method.execute(list(inputs))]
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+@unittest.skipIf(
+    torch.version.hip is not None,
+    "loads through the Python runtime, which the ROCm build does not include",
+)
+class CudaDevicePlacementTest(unittest.TestCase):
+    def test_transposed_parameter_read_by_a_delegate(self):
+        """The program stores a transposed parameter contiguous and copies
+        those bytes to the device. to_edge makes the parameter contiguous
+        before the delegate is compiled, so the delegate expects that layout."""
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                weight = torch.arange(6.0).reshape(3, 2).t()
+                self.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+            def forward(self, x):
+                return torch.add(x, self.weight)
+
+        model = Model()
+        x = torch.arange(6.0).reshape(2, 3) * 10
+        (output,) = run(model, (x,))
+        torch.testing.assert_close(output, model(x))
+
+    def test_constant_read_by_two_delegates(self):
+        """Two delegates read a constant stored in the program from outside
+        them, so the program copies it to the device for them."""
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("table", torch.arange(6.0).reshape(2, 3))
+
+            def forward(self, x):
+                return torch.add(x, self.table), torch.add(x * 2, self.table)
+
+        model = Model()
+        x = torch.arange(6.0).reshape(2, 3) * 10
+        for output, expected in zip(run(model, (x,)), model(x), strict=True):
+            torch.testing.assert_close(output, expected)

From e9d8c0e38f0b73c118488c66508b3ce3953474c4 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 24 Oct 2026 00:40:15 -0700
Subject: [PATCH 12/26] Say when skip_d2h_for_method_outputs still copies a
 returned output

With skip_h2d_for_method_inputs and skip_d2h_for_method_outputs, users
expect no copies at the method boundary. Two cases still copy on every
call:

- a constant that feeds a device delegate from outside it, because its
  data is stored on the host. The pass already logs a warning with its
  name and size.
- with skip_d2h_for_method_outputs, a returned delegate output that a
  custom host op with a mutable argument, such as a KV cache update, writes
  in place. The method returns a device copy of the written host
  copy, so the output goes to the host and back on every call. This was
  silent.

The second case now logs a warning too. It gives the output's size, the
device and the op that writes it, and suggests moving that op into the
delegate. To name the op, the write check now has a variant,
_in_place_writer, that returns the writer. The config comment for
skip_d2h_for_method_outputs and the CUDA backend docs mention the copies.

The test for the second case now checks the warning, and fails without it.
---
 docs/source/backends/cuda/cuda-overview.md |  6 ++++
 exir/passes/propagate_device_config.py     |  8 ++++--
 exir/passes/propagate_device_pass.py       | 32 +++++++++++++++++-----
 exir/tests/test_propagate_device_pass.py   | 28 +++++++++++--------
 4 files changed, 53 insertions(+), 21 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index ba681b84dea..a104f730b7e 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -185,6 +185,12 @@ on every call, one for each device that reads the constant, and logs a warning.
 device memory until the last delegate on that device has read it. Keep constants inside the
 delegate to avoid that cost.
 
+**With `skip_d2h_for_method_outputs`, a returned output that a custom operator outside the
+delegate writes in place is still copied.** This is a custom operator with a mutable argument, such
+as a KV cache update. It writes the host copy, so `to_executorch` returns a device copy of the
+written tensor. That copies the output to the host and back on every call and logs a warning. Put
+that operator in the delegate to avoid it.
+
 **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the
 program reserves its own input buffer and the runtime copies the caller's memory into it, so the
 copy comes back at run time. On device memory that copy is a host memcpy into a device pointer,
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index ab2cfe1d8f0..037039977d2 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -62,9 +62,11 @@ class PropagateDeviceConfig:
     skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False
 
     # When True, device delegate outputs that the method returns stay on
-    # device: the method output never reads a _d2h_copy of them. When a host op
-    # writes the host copy of such an output in place, the method returns a
-    # device copy of the written tensor, which costs a copy on every call.
+    # device: the method output never reads a _d2h_copy of them. When a custom
+    # host op with a mutable argument, such as a KV cache update, writes the
+    # host copy of such an output in place, the method returns a device copy of
+    # the written tensor. That costs a copy to the host and one back on every
+    # call, and the pass logs a warning about it.
     # Useful for cross-method GPU pipelines where the next method consumes GPU
     # tensors directly.
     #
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index e0aea8f3a63..4ece3e6ccb8 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -172,9 +172,9 @@ def _schema_aliases_input(
     )
 
 
-def _is_written_in_place(node: torch.fx.Node) -> bool:
-    """Whether an op writes *node* in place, directly or through a view of it,
-    per the op's schema.
+def _in_place_writer(node: torch.fx.Node) -> Optional[torch.fx.Node]:
+    """An op that writes *node* in place, directly or through a view of it,
+    per the op's schema, or None.
 
     A mutable op left on the host does this once to_executorch has removed
     auto_functionalized, so a copy of *node* can miss the write. If the op
@@ -197,11 +197,16 @@ def _is_written_in_place(node: torch.fx.Node) -> bool:
             if schema is None:
                 continue
             if _schema_writes_input(user, schema, alias):
-                return True
+                return user
             if user not in aliases and _schema_aliases_input(user, schema, alias):
                 aliases.add(user)
                 pending.append(user)
-    return False
+    return None
+
+
+def _is_written_in_place(node: torch.fx.Node) -> bool:
+    """Whether an op writes *node* in place; see _in_place_writer."""
+    return _in_place_writer(node) is not None
 
 
 def _replace_spec(
@@ -487,7 +492,8 @@ def _keep_method_outputs_on_device(
 
         The copy stays for its other users, such as a host op. When one of them
         writes the copy in place, the device tensor never sees that write, so the
-        output reads a new device copy of the written host copy instead.
+        output reads a new device copy of the written host copy instead, and a
+        warning says so.
         """
         d2h = torch.ops.et_copy._d2h_copy.default
         for node in list(graph_module.graph.nodes):
@@ -498,7 +504,8 @@ def _keep_method_outputs_on_device(
             ):
                 continue
             returned = node.args[0]
-            if _is_written_in_place(node):
+            writer = _in_place_writer(node)
+            if writer is not None:
                 device_spec = returned.meta["spec"]
                 with graph_module.graph.inserting_before(
                     graph_module.graph.output_node()
@@ -512,6 +519,17 @@ def _keep_method_outputs_on_device(
                 returned.meta["val"] = node.meta.get("val")
                 if "tensor_meta" in node.meta:
                     returned.meta["tensor_meta"] = node.meta["tensor_meta"]
+                logger.warning(
+                    "A method output (%d bytes) comes from a delegate on %s:%d "
+                    "and '%s' writes it in place on the host, so it is copied to "
+                    "the host and back to the device on every call, even with "
+                    "skip_d2h_for_method_outputs. Have the partitioner take that "
+                    "op into the delegate to avoid those copies.",
+                    device_spec.nbytes(),
+                    device_spec.device.name.lower(),
+                    device_spec.device_index,
+                    writer.target,
+                )
             _replace_spec(
                 [graph_module.graph.output_node()],
                 node.meta["spec"],
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index c0d26cb00c2..a5dc1472571 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -918,7 +918,8 @@ def test_skip_d2h_returns_device_copy_of_output_written_on_host(self):
         """With skip_d2h_for_method_outputs, a returned delegate output whose
         host copy a host op writes in place, directly or through a view, is
         returned as a device copy of the written host copy, because the device
-        tensor never sees the write."""
+        tensor never sees the write. A warning names the op that writes it,
+        since these copies run on every call."""
 
         class Model(torch.nn.Module):
             def __init__(self, write_view):
@@ -934,18 +935,23 @@ def forward(self, x, y):
                 return a
 
         for write_view in (False, True):
-            gm = _lower_with_host_writer(
-                Model(write_view),
-                (torch.randn(2, 2), torch.randn(2, 2)),
-                ["cuda:0"],
-                _host_writer_config(
-                    write_view,
-                    propagate_device_config=PropagateDeviceConfig(
-                        skip_d2h_for_method_outputs=True
+            with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs:
+                gm = _lower_with_host_writer(
+                    Model(write_view),
+                    (torch.randn(2, 2), torch.randn(2, 2)),
+                    ["cuda:0"],
+                    _host_writer_config(
+                        write_view,
+                        propagate_device_config=PropagateDeviceConfig(
+                            skip_d2h_for_method_outputs=True
+                        ),
                     ),
-                ),
-            )
+                )
             with self.subTest(write_view=write_view):
+                (warning,) = logs.output
+                self.assertIn(
+                    "'propagate_device_test.accumulate.default' writes it", warning
+                )
                 (output,) = gm.graph.output_node().args[0]
                 self.assertEqual(output.target, torch.ops.et_copy._h2d_copy.out)
                 self.assertEqual(output.meta["spec"].device, DeviceType.CUDA)

From 2ff25a8ee4cf0b95b0907837460b64f647410ac0 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 24 Oct 2026 04:04:54 -0700
Subject: [PATCH 13/26] Ignore host writes that come after a device copy is
 read

PropagateDevicePass keeps a device -> host -> device round-trip between
two delegates on one device, and does not reuse a device copy of a
tensor, when a host op writes the tensor in place. It looked for such a
write anywhere in the graph. But a write that runs after the delegate
has read the copy cannot change what the delegate reads. So some
programs paid a copy on every call that main did not make, for example
when reinplace_pass turns a host op into an in-place op, which it only
does on the last use of a tensor, or when a mutable buffer is written
back after its readers.

The round-trip now stays only when the host copy is written before it is
copied back to the device. A copy is reused only when no write runs
between it and the delegate that would reuse it, and the newest copy on
that device is the one considered. The graph order is the order the
method runs in, so it decides what comes before.

The new test writes two tensors on the host after the second of two
delegates on one device has read them. The round-trip is removed and
the tensor both delegates read gets one copy. The test fails without
this change, and the tests that write between the delegates still pass.
---
 exir/passes/propagate_device_pass.py     | 87 +++++++++++++++---------
 exir/tests/test_propagate_device_pass.py | 39 +++++++++++
 2 files changed, 95 insertions(+), 31 deletions(-)

diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 4ece3e6ccb8..11c34ea97a7 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -172,9 +172,9 @@ def _schema_aliases_input(
     )
 
 
-def _in_place_writer(node: torch.fx.Node) -> Optional[torch.fx.Node]:
-    """An op that writes *node* in place, directly or through a view of it,
-    per the op's schema, or None.
+def _in_place_writers(node: torch.fx.Node) -> list[torch.fx.Node]:
+    """The ops that write *node* in place, directly or through a view of it,
+    per the op's schema.
 
     A mutable op left on the host does this once to_executorch has removed
     auto_functionalized, so a copy of *node* can miss the write. If the op
@@ -188,6 +188,7 @@ def _in_place_writer(node: torch.fx.Node) -> Optional[torch.fx.Node]:
     such as a KV cache update. Users without a schema, such as delegate calls
     and getitem, do not write their inputs.
     """
+    writers: dict[torch.fx.Node, None] = {}
     aliases = {node}
     pending = [node]
     while pending:
@@ -197,16 +198,54 @@ def _in_place_writer(node: torch.fx.Node) -> Optional[torch.fx.Node]:
             if schema is None:
                 continue
             if _schema_writes_input(user, schema, alias):
-                return user
+                writers[user] = None
             if user not in aliases and _schema_aliases_input(user, schema, alias):
                 aliases.add(user)
                 pending.append(user)
-    return None
+    return list(writers)
 
 
-def _is_written_in_place(node: torch.fx.Node) -> bool:
-    """Whether an op writes *node* in place; see _in_place_writer."""
-    return _in_place_writer(node) is not None
+def _written_between(
+    node: torch.fx.Node,
+    after: Optional[torch.fx.Node],
+    before: torch.fx.Node,
+) -> bool:
+    """Whether an op writes *node* in place after *after*, or anywhere when it
+    is None, and before *before*, in graph order, which is the order the
+    method runs in. A write outside that range cannot change what a reader at
+    *before* of a copy made at *after* sees."""
+    writers = _in_place_writers(node)
+    if not writers:
+        return False
+    position = {n: i for i, n in enumerate(before.graph.nodes)}
+    start = position[after] if after is not None else -1
+    return any(start < position[writer] < position[before] for writer in writers)
+
+
+def _reusable_copy(
+    tensor: torch.fx.Node,
+    delegate: torch.fx.Node,
+    device_type: schema.DeviceType,
+    device_index: int,
+) -> Optional[torch.fx.Node]:
+    """The _h2d_copy of *tensor* to this device that *delegate* can read too.
+
+    Delegates are visited in graph order, so a copy made for an earlier
+    delegate on this device already comes before *delegate*, and the last such
+    copy is the newest. It is stale when an op writes *tensor* in place between
+    that copy and *delegate*, so then there is none.
+    """
+    copy = None
+    for user in tensor.users:
+        if (
+            user.target == torch.ops.et_copy._h2d_copy.default
+            and user.meta["spec"].device == device_type
+            and user.meta["spec"].device_index == device_index
+        ):
+            copy = user
+    if copy is not None and _written_between(tensor, copy, delegate):
+        return None
+    return copy
 
 
 def _replace_spec(
@@ -238,8 +277,8 @@ class PropagateDevicePass(PassBase):
     When a delegate has a target_device CompileSpec (e.g., "cuda:0"):
     - For each delegate input: insert et_copy._h2d_copy before the delegate call,
       or reuse the one an earlier delegate on the same device reads when no op
-      writes the input in place. A reused copy stays in device memory until its
-      last reader.
+      writes the input in place in between. A reused copy stays in device
+      memory until its last reader.
       The original input node keeps its device; the h2d_copy output is tagged as
       device.
     - For each delegate output: insert et_copy._d2h_copy after each getitem,
@@ -382,21 +421,7 @@ def _insert_h2d_copies(
                     f"it is created, for example in the pass that adds it."
                 )
 
-            # Delegates are visited in graph order, so a copy made for an earlier
-            # delegate on this device already comes before this one. It can be
-            # stale when an op writes the tensor in place.
-            h2d_node = None
-            if not _is_written_in_place(arg):
-                h2d_node = next(
-                    (
-                        user
-                        for user in arg.users
-                        if user.target == torch.ops.et_copy._h2d_copy.default
-                        and user.meta["spec"].device == target_device_type
-                        and user.meta["spec"].device_index == device_index
-                    ),
-                    None,
-                )
+            h2d_node = _reusable_copy(arg, node, target_device_type, device_index)
             if h2d_node is None:
                 with graph_module.graph.inserting_before(node):
                     h2d_node = graph_module.graph.call_function(
@@ -504,8 +529,8 @@ def _keep_method_outputs_on_device(
             ):
                 continue
             returned = node.args[0]
-            writer = _in_place_writer(node)
-            if writer is not None:
+            writers = _in_place_writers(node)
+            if writers:
                 device_spec = returned.meta["spec"]
                 with graph_module.graph.inserting_before(
                     graph_module.graph.output_node()
@@ -528,7 +553,7 @@ def _keep_method_outputs_on_device(
                     device_spec.nbytes(),
                     device_spec.device.name.lower(),
                     device_spec.device_index,
-                    writer.target,
+                    writers[0].target,
                 )
             _replace_spec(
                 [graph_module.graph.output_node()],
@@ -559,8 +584,8 @@ def _elide_redundant_same_device_copies(
         rewires the consumer straight to that device tensor and drops both copies.
         The device comparison is the correctness guard: a genuine cross-device
         move (e.g. ``cuda:0 -> host -> cuda:1``) has differing devices and is left
-        intact. So is a round-trip whose host copy any op writes in place,
-        wherever that write is, because the device tensor never sees it. A D2H
+        intact. So is a round-trip whose host copy an op writes in place before
+        the copy back, because the device tensor never sees that write. A D2H
         feeding other (host) consumers keeps those users and is removed only once
         it becomes dead.
         """
@@ -590,7 +615,7 @@ def _elide_redundant_same_device_copies(
             if (
                 x_spec.device != h2d_spec.device
                 or x_spec.device_index != h2d_spec.device_index
-                or _is_written_in_place(src)
+                or _written_between(src, None, node)
             ):
                 continue
 
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index a5dc1472571..71b7da86e31 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -625,6 +625,45 @@ def forward(self, x, y):
                 y_copies = [n for n in nodes.h2d_nodes if n.args[0].name == "y"]
                 self.assertEqual(len(y_copies), 1)
 
+    def test_host_writes_after_the_device_readers_keep_no_copies(self):
+        """Host ops that write tensors in place only after the delegates on one
+        device have read them cannot change what those delegates read. So the
+        round-trip between the delegates is still removed, and the tensor both
+        delegates read gets one copy."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                t = x * 2
+                a = t + y
+                b = a + t
+                # Both writes read b, so they run after the second delegate.
+                torch.ops.propagate_device_test.accumulate(a, b)
+                torch.ops.propagate_device_test.accumulate(t, b)
+                return a, b, t
+
+        gm = _lower_with_host_writer(
+            Model(),
+            (torch.randn(2, 2), torch.randn(2, 2)),
+            ["cuda:0", "cuda:0"],
+            _host_writer_config(False),
+        )
+        nodes = _collect_device_copy_nodes(gm)
+        first, second = nodes.delegate_nodes
+        order = list(gm.graph.nodes)
+        writers = [
+            n
+            for n in order
+            if n.target == torch.ops.propagate_device_test.accumulate.out
+        ]
+        self.assertEqual(len(writers), 2)
+        for writer in writers:
+            self.assertLess(order.index(second), order.index(writer))
+        # Only t and y are copied, and the second delegate reads the first
+        # delegate's output directly.
+        self.assertEqual(len(nodes.h2d_nodes), 2)
+        (shared,) = set(first.args[1:]) & set(second.args[1:])
+        self.assertIn(shared, nodes.h2d_nodes)
+
     def test_single_tensor_delegate_output_read_on_host(self):
         """A host op that reads the output of a lowered module returning one
         tensor reads a host copy of it."""

From a22b90409b5236efbfd8306234f9a851940c3b3b Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 24 Oct 2026 04:40:47 -0700
Subject: [PATCH 14/26] Share the device copy of a written output with a later
 delegate

With skip_d2h_for_method_outputs, when a host op writes the host copy of
a returned delegate output in place, the method returns a device copy of
the written tensor. If a later delegate on the same device also reads
that tensor, it already reads a device copy of the written host copy, so
the method made two identical copies on every call.

The method output now reads the copy that the later delegate reads, when
no op writes the tensor in place between that copy and the end of the
method. Otherwise it still gets its own copy.

The new test returns a delegate output that a host op writes and that a
second delegate on the same device reads. There is one device copy, read
by both. The test fails without this change.
---
 exir/passes/propagate_device_pass.py     | 27 ++++++++++++--------
 exir/tests/test_propagate_device_pass.py | 32 ++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 11c34ea97a7..8bb858bff50 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -532,18 +532,23 @@ def _keep_method_outputs_on_device(
             writers = _in_place_writers(node)
             if writers:
                 device_spec = returned.meta["spec"]
-                with graph_module.graph.inserting_before(
-                    graph_module.graph.output_node()
-                ):
-                    returned = graph_module.graph.call_function(
-                        torch.ops.et_copy._h2d_copy.default, (node,)
-                    )
-                returned.meta["spec"] = _copy_output_spec(
-                    node.meta["spec"], device_spec.device, device_spec.device_index
+                output_node = graph_module.graph.output_node()
+                # A later delegate on this device may already read a device copy
+                # of the written host copy.
+                returned = _reusable_copy(
+                    node, output_node, device_spec.device, device_spec.device_index
                 )
-                returned.meta["val"] = node.meta.get("val")
-                if "tensor_meta" in node.meta:
-                    returned.meta["tensor_meta"] = node.meta["tensor_meta"]
+                if returned is None:
+                    with graph_module.graph.inserting_before(output_node):
+                        returned = graph_module.graph.call_function(
+                            torch.ops.et_copy._h2d_copy.default, (node,)
+                        )
+                    returned.meta["spec"] = _copy_output_spec(
+                        node.meta["spec"], device_spec.device, device_spec.device_index
+                    )
+                    returned.meta["val"] = node.meta.get("val")
+                    if "tensor_meta" in node.meta:
+                        returned.meta["tensor_meta"] = node.meta["tensor_meta"]
                 logger.warning(
                     "A method output (%d bytes) comes from a delegate on %s:%d "
                     "and '%s' writes it in place on the host, so it is copied to "
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 71b7da86e31..7de68a1e154 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -998,6 +998,38 @@ def forward(self, x, y):
                 nodes = list(gm.graph.nodes)
                 self.assertLess(nodes.index(_host_writer(gm)), nodes.index(output))
 
+    def test_skip_d2h_shares_the_device_copy_of_a_written_output(self):
+        """With skip_d2h_for_method_outputs, a returned output that a host op
+        writes before a later delegate on the same device reads it needs a
+        device copy for both. The method output reads the same copy as that
+        delegate."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                a = x + y
+                torch.ops.propagate_device_test.accumulate(a, y)
+                return a, a + y
+
+        gm = _lower_with_host_writer(
+            Model(),
+            (torch.randn(2, 2), torch.randn(2, 2)),
+            ["cuda:0", "cuda:0"],
+            _host_writer_config(
+                False,
+                propagate_device_config=PropagateDeviceConfig(
+                    skip_d2h_for_method_outputs=True
+                ),
+            ),
+        )
+        (copy,) = [
+            n
+            for n in _collect_device_copy_nodes(gm).h2d_nodes
+            if n.args[0] is _host_written(gm)
+        ]
+        _, second = _collect_device_copy_nodes(gm).delegate_nodes
+        self.assertIn(copy, second.args[1:])
+        self.assertIs(gm.graph.output_node().args[0][0], copy)
+
     def test_cross_device_copies_preserved(self):
         """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two
         directly-connected delegates must be preserved -- the round-trip is a

From e92206b6cb2d9758c85ced6dc75053e90764e252 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 24 Oct 2026 04:41:52 -0700
Subject: [PATCH 15/26] Check the placeholder, not the input spec, in
 convert_constant_dim_order_pass

_update_placeholder_meta looks up the input spec of a lifted tensor and
then its placeholder. The second check tested the input spec again, which
was already known to exist, so a missing placeholder failed later with an
AttributeError on None instead of the intended "Missing placeholder"
error. Parameters now go through this function too.

The new test points a parameter's input spec at a name with no
placeholder and expects the error. It fails without this change.
---
 .../passes/convert_constant_dim_order_pass.py |  2 +-
 exir/tests/test_passes.py                     | 20 +++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/exir/passes/convert_constant_dim_order_pass.py b/exir/passes/convert_constant_dim_order_pass.py
index 9cc00a66498..f93fb4c74da 100644
--- a/exir/passes/convert_constant_dim_order_pass.py
+++ b/exir/passes/convert_constant_dim_order_pass.py
@@ -42,7 +42,7 @@ def _update_placeholder_meta(
         ),
         None,
     )
-    if input_spec is None:
+    if placeholder_node is None:
         raise RuntimeError(f"Missing placeholder for {input_spec.arg.name}.")
 
     placeholder_node.meta["val"] = placeholder_node.meta["val"].contiguous()
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 1a7ba243950..2f775b14c15 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -3185,6 +3185,26 @@ def forward(self, x):
                 self.assertTrue(placeholder.meta["val"].is_contiguous())
                 torch.testing.assert_close(ep.module()(*inputs), model(*inputs))
 
+    def test_convert_constant_dim_order_missing_placeholder(self):
+        """
+        Verify that the convert_constant_dim_order_pass names the placeholder it
+        cannot find for a lifted tensor it converts.
+        """
+
+        class AddParameter(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.randn(5, 4).t())
+
+            def forward(self, x):
+                return x + self.weight
+
+        ep = torch.export.export(AddParameter(), (torch.randn(4, 5),))
+        spec = next(s for s in ep.graph_signature.input_specs if s.target == "weight")
+        spec.arg = TensorArgument(name="missing")
+        with self.assertRaisesRegex(RuntimeError, "Missing placeholder for missing"):
+            convert_constant_dim_order_pass.convert_constant_dim_order_pass(ep)
+
 
 class TestMemoryFormatOpsPassPreserveFormat(unittest.TestCase):
     """

From dda74d38b41fde91298a5ef647871f5e4097197e Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 24 Oct 2026 05:14:53 -0700
Subject: [PATCH 16/26] Run the CUDA unit tests when the device placement
 passes change

backends/cuda/tests/test_device_placement.py runs programs whose
host/device copies come from PropagateDevicePass and whose parameter
layout comes from convert_constant_dim_order_pass. The CUDA workflow did
not start for changes to those passes, so a change to them alone could
merge without running the test.

The workflow's pull_request paths and the unittest-cuda job now include
those two passes, the device pass config, and the write-back pass whose
helpers the device pass uses to find in-place writes, the same way they
already include the pybindings paths for the pybindings device test. The
other CUDA jobs keep their own path checks, so they do not start for
these files.

The test skips without a GPU, and pytest exits 0 on a skip, so the system
shard now runs it next to the pybindings device tests and treats a skip
as a failure, as it already does for those.
---
 .github/workflows/cuda.yml | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml
index 41dd61ebcda..07653ae8bd4 100644
--- a/.github/workflows/cuda.yml
+++ b/.github/workflows/cuda.yml
@@ -43,6 +43,10 @@ on:
       - examples/models/muse-glimmer/**
       - extension/pybindings/**
       - runtime/__init__.py
+      - exir/passes/convert_constant_dim_order_pass.py
+      - exir/passes/insert_write_back_for_buffers_pass.py
+      - exir/passes/propagate_device_config.py
+      - exir/passes/propagate_device_pass.py
   workflow_dispatch:
     inputs:
       models:
@@ -469,8 +473,9 @@ jobs:
   unittest-cuda:
     name: unittest-cuda (${{ matrix.shard }})
     needs: [changed-files, run-decision]
-    # This job runs the pybindings device test, so it also has to fire on the
-    # two pybindings paths the workflow filter lists.
+    # This job runs the pybindings device test and the device placement test,
+    # so it also has to fire on the pybindings and device pass paths the
+    # workflow filter lists.
     if: |
       contains(needs.changed-files.outputs.changed-files, 'backends/cuda') ||
       contains(needs.changed-files.outputs.changed-files, 'backends/aoti') ||
@@ -481,6 +486,10 @@ jobs:
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') ||
       contains(needs.changed-files.outputs.changed-files, 'extension/pybindings') ||
       contains(needs.changed-files.outputs.changed-files, 'runtime/__init__.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/convert_constant_dim_order_pass.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/insert_write_back_for_buffers_pass.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/propagate_device_config.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/propagate_device_pass.py') ||
       needs.run-decision.outputs.is-full-run == 'true'
     # kernel: Triton kernels and their autotuning. system: lowering passes, op
     # dispatch, export/runtime, and model tests. Together they do not fit the
@@ -567,6 +576,21 @@ jobs:
             echo "::error::a pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
             exit 1
           fi
+
+          # Same for the device placement test, which skips without a GPU.
+          set +e
+          python -m pytest backends/cuda/tests/test_device_placement.py -v -rs -o "addopts=" \
+            2>&1 | tee /tmp/device_placement_test.log
+          device_placement_status=${PIPESTATUS[0]}
+          set -e
+          if [ "${device_placement_status}" -ne 0 ]; then
+            echo "::error::a device placement test failed"
+            exit "${device_placement_status}"
+          fi
+          if grep -qE "[0-9]+ skipped" /tmp/device_placement_test.log; then
+            echo "::error::a device placement test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
+            exit 1
+          fi
         fi
 
         # Build ExecuTorch with CUDA support
@@ -589,7 +613,7 @@ jobs:
           python -m pytest "${KERNEL_TESTS[@]}" -v -o "addopts=" --durations=20
         else
           IGNORES=()
-          for t in "${KERNEL_TESTS[@]}" "${QUANTIZED_GEMM_TESTS[@]}"; do
+          for t in "${KERNEL_TESTS[@]}" "${QUANTIZED_GEMM_TESTS[@]}" backends/cuda/tests/test_device_placement.py; do
             IGNORES+=("--ignore=${t}")
           done
           python -m pytest backends/cuda/tests backends/cuda/passes/tests -v -o "addopts=" --durations=20 \

From 5f40202cacf8f2c992448e15f556f27c8824f858 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 25 Oct 2026 16:39:18 -0700
Subject: [PATCH 17/26] Warn when to_edge changes the layout of a parameter,
 buffer or constant

to_edge stores a parameter, buffer or constant that is neither contiguous
nor channels-last, such as a transposed weight, as contiguous. Partitioners
and delegates then see the contiguous layout, not the strides in the model.
This happened without any message, for every backend.

convert_constant_dim_order_pass now logs a warning with the tensor's kind,
name and original strides each time it makes one contiguous. to_edge runs
the pass twice, but only the first run changes a tensor, so each one is
reported once.

The updated tests expect the warning for a transposed parameter, from the
pass and from to_edge with and without _skip_dim_order, and no warning for
a channels-last parameter. They fail without this change.
---
 .../passes/convert_constant_dim_order_pass.py | 14 ++++++++
 exir/tests/test_passes.py                     | 34 ++++++++++++++-----
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/exir/passes/convert_constant_dim_order_pass.py b/exir/passes/convert_constant_dim_order_pass.py
index f93fb4c74da..36037961f62 100644
--- a/exir/passes/convert_constant_dim_order_pass.py
+++ b/exir/passes/convert_constant_dim_order_pass.py
@@ -3,10 +3,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
+
 import torch
 from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 def _should_transform(tensor: torch.Tensor) -> bool:
     """
@@ -45,7 +49,17 @@ def _update_placeholder_meta(
     if placeholder_node is None:
         raise RuntimeError(f"Missing placeholder for {input_spec.arg.name}.")
 
+    strides = placeholder_node.meta["val"].stride()
     placeholder_node.meta["val"] = placeholder_node.meta["val"].contiguous()
+    logger.warning(
+        "%s '%s' has strides %s, which are neither contiguous nor channels-last. "
+        "to_edge stores it contiguous, so partitioners and delegates see the "
+        "contiguous layout, not the strides in the model. Make it contiguous "
+        "where it is created to avoid this change.",
+        kind.name.replace("_", " ").capitalize(),
+        target,
+        tuple(strides),
+    )
 
 
 def convert_constant_dim_order_pass(
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 2f775b14c15..b338dc2d3ea 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -3141,11 +3141,21 @@ def forward(self, x):
             with self.subTest(description=description):
                 inputs = (torch.randn(weight.shape),)
                 ep = torch.export.export(AddParameter(weight), inputs)
-                modified_ep = (
-                    convert_constant_dim_order_pass.convert_constant_dim_order_pass(
-                        copy.deepcopy(ep)
-                    )
+                logs = (
+                    self.assertLogs(convert_constant_dim_order_pass.logger, "WARNING")
+                    if converted
+                    else self.assertNoLogs(convert_constant_dim_order_pass.logger)
                 )
+                with logs as captured:
+                    modified_ep = (
+                        convert_constant_dim_order_pass.convert_constant_dim_order_pass(
+                            copy.deepcopy(ep)
+                        )
+                    )
+                if converted:
+                    self.assertIn(
+                        "Parameter 'weight' has strides (1, 4)", captured.output[0]
+                    )
 
                 torch.testing.assert_close(
                     ep.module()(*inputs), modified_ep.module()(*inputs)
@@ -3176,10 +3186,18 @@ def forward(self, x):
         inputs = (torch.randn(4, 5),)
         for skip_dim_order in (False, True):
             with self.subTest(skip_dim_order=skip_dim_order):
-                ep = to_edge(
-                    torch.export.export(model, inputs),
-                    compile_config=EdgeCompileConfig(_skip_dim_order=skip_dim_order),
-                ).exported_program()
+                with self.assertLogs(
+                    convert_constant_dim_order_pass.logger, "WARNING"
+                ) as captured:
+                    ep = to_edge(
+                        torch.export.export(model, inputs),
+                        compile_config=EdgeCompileConfig(
+                            _skip_dim_order=skip_dim_order
+                        ),
+                    ).exported_program()
+                # to_edge runs the pass twice, but only the first run converts.
+                self.assertEqual(len(captured.output), 1)
+                self.assertIn("Parameter 'weight'", captured.output[0])
                 self.assertTrue(ep.state_dict["weight"].is_contiguous())
                 placeholder = next(n for n in ep.graph.nodes if n.name == "p_weight")
                 self.assertTrue(placeholder.meta["val"].is_contiguous())

From 825f18748173cad8e72fad73fe63e4caf6b86b81 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 25 Oct 2026 18:02:22 -0700
Subject: [PATCH 18/26] Give a delegate whose target device is the CPU no
 device copies

A partitioner can set the target_device compile spec to "cpu". The pass
treated that delegate like one on a device: it copied each input with
_h2d_copy and each output with _d2h_copy, both to and from the CPU. The
copy kernels refuse a CPU destination for _h2d_copy and a CPU source for
_d2h_copy, so such a program failed at run time.

A delegate on the CPU reads and writes host memory directly, so the pass
now leaves it alone, the same as a delegate without a target device.

The new test lowers an add to a delegate with target device "cpu" and
expects no copies. It fails without this change.
---
 exir/passes/propagate_device_pass.py     |  9 ++++++---
 exir/tests/test_propagate_device_pass.py | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 8bb858bff50..b29ad98b8e9 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -274,7 +274,8 @@ class PropagateDevicePass(PassBase):
     After to_backend, walk the graph and insert H2D/D2H copy ops at delegate
     boundaries based on partitioner-assigned device info.
 
-    When a delegate has a target_device CompileSpec (e.g., "cuda:0"):
+    When a delegate has a target_device CompileSpec (e.g., "cuda:0") other
+    than the CPU:
     - For each delegate input: insert et_copy._h2d_copy before the delegate call,
       or reuse the one an earlier delegate on the same device reads when no op
       writes the input in place in between. A reused copy stays in device
@@ -639,7 +640,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         #            H2D copy nodes before delegate inputs, or reuse the copy an
         #            earlier delegate on the same device reads, and tag the delegate
         #            output specs with the target device.  Delegates without a
-        #            target_device are left untouched (no copies, specs stay CPU).
+        #            target_device, or with the CPU as their target_device, are
+        #            left untouched (no copies, specs stay CPU).
         #   Pass 2 – For each getitem that extracts from a device-tagged delegate
         #            (tracked in device_delegates), propagate the device onto the
         #            getitem spec and insert a D2H copy after it so downstream
@@ -664,7 +666,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                     )
 
                 result = _get_target_device_from_compile_specs(lowered_module)
-                if result is None:
+                # A delegate on the host reads and writes host memory directly.
+                if result is None or result[0] == schema.DeviceType.CPU:
                     continue
 
                 target_device_type, device_index = result
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 7de68a1e154..3b8b77bf9e9 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -541,6 +541,29 @@ def forward(self, a, b):
                 self.assertEqual(len(device_copy_nodes.h2d_nodes), 0)
                 self.assertEqual(len(device_copy_nodes.d2h_nodes), 0)
 
+    def test_no_copy_nodes_for_cpu_target_device(self):
+        """A delegate whose target device is the CPU reads and writes host
+        memory, so it gets no copies, like a delegate without a target device."""
+
+        class Model(torch.nn.Module):
+            def forward(self, a, b):
+                return torch.add(a, b)
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False, enable_non_cpu_memory_planning=True
+        )
+        for pipeline, gm in _lower_model_to_executorch(
+            Model(),
+            (torch.randn(2, 2), torch.randn(2, 2)),
+            PerAddDeviceDelegatePartitioner(["cpu"]),
+            et_config,
+        ):
+            with self.subTest(pipeline=pipeline):
+                nodes = _collect_device_copy_nodes(gm)
+                self.assertEqual(len(nodes.delegate_nodes), 1)
+                self.assertEqual(nodes.h2d_nodes, [])
+                self.assertEqual(nodes.d2h_nodes, [])
+
     def test_redundant_same_device_copies_elided(self):
         """Two directly-connected delegates on the *same* device must not leave a
         redundant device->host->device round-trip between them. The

From d3d397851d4eba0c56328f2894a5c8203431aaf2 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 25 Oct 2026 16:43:26 -0700
Subject: [PATCH 19/26] Copy a device input to the host for readers off the
 device

With skip_h2d_for_method_inputs, the caller passes an input that a device
delegate reads in device memory. Until now, to_executorch refused such an
input when anything else read it too: a host op, a delegate on the host,
or a delegate on another device.

The pass now keeps the input on the device for the delegates there and
gives every other reader one _d2h_copy of it, the same way it already
handles a device delegate output read on the host. A delegate on another
device then reads a copy of that host copy. The input stays on the device
of the first delegate that reads it, in graph order, and later delegates
on that device find it already placed. An op that reads only the input's
shape, such as sym_size, reads it on the device, since sizes stay in host
memory. The copy runs on every call, so the pass logs a warning with the
input's size and a reader. With
skip_d2h_for_method_outputs, a method that returns the input returns it
on the device, with no copy. An input that an op writes in place on the
host is still refused, now with its own message, because the write would
only reach the host copy.

The refusal test becomes a test that the input stays on the device and
that the other readers read one D2H copy, with the warning, for a host
op, two delegates on the device plus a host op, a delegate without a
device, a delegate on the CPU, a delegate on another device and a buffer
write-back. It fails without this change. New tests return an input
with both skip flags, and read only the shape of an input on the host,
and expect no copy and no warning. A new refusal
case covers a host op that writes the input in place, and a new CUDA test
runs a program whose input on the device is also read on the host and
compares with eager PyTorch.
---
 backends/cuda/tests/test_device_placement.py |  46 ++++-
 docs/source/backends/cuda/cuda-overview.md   |  18 +-
 exir/passes/propagate_device_config.py       |  17 +-
 exir/passes/propagate_device_pass.py         | 151 ++++++++++----
 exir/tests/test_propagate_device_pass.py     | 200 ++++++++++++++-----
 5 files changed, 317 insertions(+), 115 deletions(-)

diff --git a/backends/cuda/tests/test_device_placement.py b/backends/cuda/tests/test_device_placement.py
index 9bedbf6abff..b54242802a1 100644
--- a/backends/cuda/tests/test_device_placement.py
+++ b/backends/cuda/tests/test_device_placement.py
@@ -22,7 +22,11 @@
     PartitionResult,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.passes.propagate_device_pass import TARGET_DEVICE_COMPILE_SPEC_KEY
+from executorch.exir.passes import MemoryPlanningPass
+from executorch.exir.passes.propagate_device_pass import (
+    PropagateDeviceConfig,
+    TARGET_DEVICE_COMPILE_SPEC_KEY,
+)
 
 
 class EachAddOnCuda(Partitioner):
@@ -52,8 +56,9 @@ def partition(self, exported_program) -> PartitionResult:
         )
 
 
-def run(model, inputs) -> list[torch.Tensor]:
-    """Lower model with EachAddOnCuda, run it in the runtime, return its outputs."""
+def load(model, inputs, directory, config=None):
+    """Lower model with EachAddOnCuda and load its forward method from a file
+    in directory."""
     from executorch.runtime import Runtime
 
     edge = to_edge_transform_and_lower(
@@ -63,10 +68,15 @@ def run(model, inputs) -> list[torch.Tensor]:
             _check_ir_validity=False, _skip_dim_order=True
         ),
     )
+    path = Path(directory) / "model.pte"
+    path.write_bytes(edge.to_executorch(config or ExecutorchBackendConfig()).buffer)
+    return Runtime.get().load_program(path).load_method("forward")
+
+
+def run(model, inputs) -> list[torch.Tensor]:
+    """Lower model with EachAddOnCuda, run it in the runtime, return its outputs."""
     with tempfile.TemporaryDirectory() as directory:
-        path = Path(directory) / "model.pte"
-        path.write_bytes(edge.to_executorch(ExecutorchBackendConfig()).buffer)
-        method = Runtime.get().load_program(path).load_method("forward")
+        method = load(model, inputs, directory)
         return [output.cpu() for output in method.execute(list(inputs))]
 
 
@@ -111,3 +121,27 @@ def forward(self, x):
         x = torch.arange(6.0).reshape(2, 3) * 10
         for output, expected in zip(run(model, (x,)), model(x), strict=True):
             torch.testing.assert_close(output, expected)
+
+    def test_input_on_device_also_read_on_host(self):
+        """With skip_h2d_for_method_inputs, the caller passes the input on the
+        device. A delegate reads it there, and a host op reads a copy of it in
+        host memory."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.add(x, y), torch.sin(x)
+
+        config = ExecutorchBackendConfig(
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
+        )
+        model = Model()
+        inputs = (torch.randn(2, 3), torch.randn(2, 3))
+        with tempfile.TemporaryDirectory() as directory:
+            method = load(model, inputs, directory, config)
+            outputs = method.execute([t.cuda() for t in inputs])
+        for output, expected in zip(outputs, model(*inputs), strict=True):
+            torch.testing.assert_close(output.cpu(), expected)
diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index a104f730b7e..69e8724507f 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -169,14 +169,16 @@ out works, and the snippet above sets it explicitly only to make the requirement
 it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens
 during device-aware memory planning.
 
-**With `skip_h2d_for_method_inputs`, an input that a device delegate reads may only be read by
-delegates on that device.** If anything else reads it, even only its shape, such as an operator
-outside those delegates, a delegate on another device, or the method output, `to_executorch` raises
-an error instead of adding a copy. Turn the flag off, or have the partitioner put those operators in
-a delegate on that device. `to_executorch` also refuses an input that the method changes, because
-the change is written back on the host. For a mutable buffer, have the partitioner put the buffer
-and every operator that reads or writes it in the delegate. `tag_mutated_buffer` tags the buffer
-once those operators are there, and the CUDA partitioner already calls it.
+**With `skip_h2d_for_method_inputs`, an input that a device delegate reads stays on that device.**
+If delegates on more than one device read it, that is the device of the first one in the graph.
+An operator that reads only its shape reads it there too. Anything else that reads it, such as an
+operator outside the delegates or a delegate on another device, reads a copy of it in host memory,
+made on every call, and `to_executorch` logs a warning. Have the partitioner put those operators in
+a delegate on that device to avoid that copy. `to_executorch` refuses an input that the method
+changes, because the change is written back on the host, and an input that an operator outside the
+delegates writes in place. For a mutable buffer, have the partitioner put the buffer and every
+operator that reads or writes it in the delegate. `tag_mutated_buffer` tags the buffer once those
+operators are there, and the CUDA partitioner already calls it.
 
 **A constant that a delegate reads from outside is still copied.** Constants, such as weights and
 buffers the method does not change, are not passed by the caller. Their data stays in host memory,
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index 037039977d2..a688e74b141 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -48,13 +48,16 @@ class PropagateDeviceConfig:
 
     # When True, method-level input tensors that feed directly into a device
     # delegate are NOT wrapped with _h2d_copy. The user must provide tensors
-    # already on the target device. Every reader of such an input must then be
-    # a delegate on that device, or to_executorch raises. Useful for pipelines
-    # where inputs are pre-staged on GPU. Constants are not passed by the
-    # caller, so they are still copied: their data stays on the host, in the
-    # program or in its external data file. That copy runs on every call, once
-    # for each device that reads the constant, and stays in device memory until
-    # the last delegate there has read it.
+    # already on the target device: the device of the first delegate that
+    # reads the input, in graph order. Any other reader of such an input, such
+    # as a host op or a delegate on another device, reads a copy of it in host
+    # memory, made on every call, and the pass logs a warning. An input that an
+    # op writes in place on the host makes to_executorch raise. Useful for
+    # pipelines where inputs are pre-staged on GPU. Constants are not passed
+    # by the caller, so they are still copied: their data stays on the host,
+    # in the program or in its external data file. That copy runs on every
+    # call, once for each device that reads the constant, and stays in device
+    # memory until the last delegate there has read it.
     #
     # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer
     # for the input and copies the caller's memory into it, which is what this exists to avoid.
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index b29ad98b8e9..c041be2a6f1 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -42,6 +42,16 @@
 # with this key and a value encoding the device string (e.g., b"cuda:0").
 TARGET_DEVICE_COMPILE_SPEC_KEY = "target_device"
 
+# Ops that read only the sizes, strides or offset of a tensor. Those stay in
+# host memory when the data is on a device, so these ops need no copy of it.
+_METADATA_READERS = (
+    torch.ops.aten.sym_size.default,
+    torch.ops.aten.sym_size.int,
+    torch.ops.aten.sym_stride.int,
+    torch.ops.aten.sym_numel.default,
+    torch.ops.aten.sym_storage_offset.default,
+)
+
 
 def _parse_device_spec_value(value: bytes) -> tuple[schema.DeviceType, int]:
     """
@@ -290,10 +300,10 @@ class PropagateDevicePass(PassBase):
     Skip-copy optimizations:
     - skip_h2d_for_method_inputs: If the input is a graph-level placeholder
       feeding a delegate, don't insert H2D; tag the placeholder as device
-      instead (user provides GPU tensor at runtime). Every reader of the
-      placeholder must then be a delegate on that device, or the pass raises.
-      Constants are placeholders too, but their data is stored on the host,
-      so they still get an H2D copy.
+      instead (user provides GPU tensor at runtime). Every other reader of the
+      placeholder reads one D2H copy of it. A placeholder that an op writes in
+      place on the host is refused. Constants are placeholders too, but their
+      data is stored on the host, so they still get an H2D copy.
     - skip_d2h_for_method_outputs: The graph output reads a delegate output
       on device, not its D2H copy; the copy stays only for other readers. When
       one of them writes the copy in place, the graph output reads a device
@@ -365,45 +375,9 @@ def _insert_h2d_copies(
                 and self._is_placeholder(arg)
                 and not arg_spec.const
             ):
-                # The spec is shared by every reader, so all of them must be
-                # delegates on this same device.
-                other_readers = [
-                    user
-                    for user in arg.users
-                    if self._delegate_device(graph_module, user)
-                    != (target_device_type, device_index)
-                ]
-                if other_readers:
-                    write_back = next(
-                        (
-                            user
-                            for user in other_readers
-                            if user.target == torch.ops.aten.copy_.default
-                            and user.args[0] is arg
-                        ),
-                        None,
-                    )
-                    refusal = (
-                        f"skip_h2d_for_method_inputs=True cannot place placeholder "
-                        f"'{arg.name}' on {target_device_type.name.lower()}:"
-                        f"{device_index}, because "
-                        f"'{(write_back or other_readers[0]).name}'"
-                    )
-                    if write_back is not None:
-                        raise RuntimeError(
-                            f"{refusal} writes it back on the host (the method "
-                            f"mutates it). Either disable "
-                            f"skip_h2d_for_method_inputs, or, for a mutable buffer, "
-                            f"have the partitioner put the buffer and every op that "
-                            f"reads or writes it in the delegate; tag_mutated_buffer "
-                            f"tags the buffer once those ops are there."
-                        )
-                    raise RuntimeError(
-                        f"{refusal} also reads it and is not a delegate on that "
-                        f"device. Either disable skip_h2d_for_method_inputs, or use "
-                        f"the placeholder only from delegates on one device."
-                    )
-                _set_device_on_spec(arg_spec, target_device_type, device_index)
+                self._place_input_on_device(
+                    graph_module, arg, target_device_type, device_index
+                )
                 changed = True
                 continue
 
@@ -452,6 +426,97 @@ def _insert_h2d_copies(
         node.args = tuple(new_args)
         return changed
 
+    def _place_input_on_device(
+        self,
+        graph_module: torch.fx.GraphModule,
+        placeholder: torch.fx.Node,
+        device_type: schema.DeviceType,
+        device_index: int,
+    ) -> None:
+        """Move a method input to the device for skip_h2d_for_method_inputs.
+
+        Delegates on that device read it there, and so do ops that read only
+        its shape. Every other reader, such as a host op, a delegate on the
+        host, a delegate on another device or, without
+        skip_d2h_for_method_outputs, the method output, reads one D2H copy of it
+        instead, made on every call. An
+        input that an op writes in place on the host is refused: the write would
+        go to the copy, and neither the caller nor the delegates would see it.
+        Delegates are visited in graph order, so the first one that reads the
+        input decides its device, and later ones on that device find it placed.
+        """
+        spec = placeholder.meta["spec"]
+        if (spec.device, spec.device_index) == (device_type, device_index):
+            return
+        device = f"{device_type.name.lower()}:{device_index}"
+        writers = _in_place_writers(placeholder)
+        if writers:
+            write_back = next(
+                (
+                    writer
+                    for writer in writers
+                    if writer.target == torch.ops.aten.copy_.default
+                    and writer.args[0] is placeholder
+                ),
+                None,
+            )
+            refusal = (
+                f"skip_h2d_for_method_inputs=True cannot place placeholder "
+                f"'{placeholder.name}' on {device}, because "
+                f"'{(write_back or writers[0]).name}' writes it"
+            )
+            if write_back is not None:
+                raise RuntimeError(
+                    f"{refusal} back on the host (the method mutates it). Either "
+                    f"disable skip_h2d_for_method_inputs, or, for a mutable buffer, "
+                    f"have the partitioner put the buffer and every op that reads "
+                    f"or writes it in the delegate; tag_mutated_buffer tags the "
+                    f"buffer once those ops are there."
+                )
+            raise RuntimeError(
+                f"{refusal} in place on the host. Either disable "
+                f"skip_h2d_for_method_inputs, or have the partitioner put that op "
+                f"in a delegate on {device}."
+            )
+
+        _set_device_on_spec(spec, device_type, device_index)
+        # With skip_d2h_for_method_outputs, the method returns the input where
+        # the caller put it.
+        other_readers = [
+            user
+            for user in placeholder.users
+            if self._delegate_device(graph_module, user) != (device_type, device_index)
+            and user.target not in _METADATA_READERS
+            and not (user.op == "output" and self.skip_d2h_for_method_outputs)
+        ]
+        if not other_readers:
+            return
+        position = {n: i for i, n in enumerate(graph_module.graph.nodes)}
+        first_reader = min(other_readers, key=position.__getitem__)
+        logger.warning(
+            "Method input '%s' (%s bytes) stays on %s for the delegates there, "
+            "but '%s' also reads it, so it is copied to the host on every call. "
+            "Have the partitioner put every reader in a delegate on %s to avoid "
+            "that copy.",
+            placeholder.name,
+            spec.nbytes(),
+            device,
+            first_reader.name,
+            device,
+        )
+        with graph_module.graph.inserting_before(first_reader):
+            d2h_node = graph_module.graph.call_function(
+                torch.ops.et_copy._d2h_copy.default, (placeholder,)
+            )
+        d2h_spec = _copy_output_spec(spec, schema.DeviceType.CPU, 0)
+        d2h_node.meta["spec"] = d2h_spec
+        d2h_node.meta["val"] = placeholder.meta.get("val")
+        if "tensor_meta" in placeholder.meta:
+            d2h_node.meta["tensor_meta"] = placeholder.meta["tensor_meta"]
+        for reader in other_readers:
+            reader.replace_input_with(placeholder, d2h_node)
+        _replace_spec(other_readers, spec, d2h_spec)
+
     def _insert_d2h_for_getitem(
         self,
         graph_module: torch.fx.GraphModule,
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 3b8b77bf9e9..f8915cd6c76 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -134,10 +134,11 @@ class PerAddDeviceDelegatePartitioner(Partitioner):
     directly-connected adds become two directly-connected delegates -- the shape
     that leaves a redundant ``device -> host -> device`` round-trip between them.
     ``devices[i]`` is the target device for the i-th add in graph order (the last
-    entry is reused if there are more adds than devices).
+    entry is reused if there are more adds than devices). ``None`` gives that add
+    a delegate without a target device.
     """
 
-    def __init__(self, devices: List[str]) -> None:
+    def __init__(self, devices: List[Optional[str]]) -> None:
         super().__init__()
         self._devices = list(devices)
 
@@ -152,12 +153,13 @@ def partition(self, exported_program) -> PartitionResult:
             device = self._devices[min(i, len(self._devices) - 1)]
             tag = f"tag{i}"
             node.meta["delegation_tag"] = tag
+            compile_specs = [CompileSpec("max_value", bytes([4]))]
+            if device is not None:
+                compile_specs.append(
+                    CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8"))
+                )
             partition_tags[tag] = DelegationSpec(
-                BackendWithCompilerDemo.__name__,
-                [
-                    CompileSpec("max_value", bytes([4])),
-                    CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")),
-                ],
+                BackendWithCompilerDemo.__name__, compile_specs
             )
         return PartitionResult(
             tagged_exported_program=exported_program,
@@ -2121,9 +2123,12 @@ def forward(self, a, b):
                     cuda_ph, program, DeviceType.CUDA, pipeline, "Placeholder"
                 )
 
-    def test_skip_h2d_refuses_input_read_off_device(self):
-        """A method input cannot move to a device when a host op, or a delegate
-        on another device, also reads it."""
+    def test_skip_h2d_copies_input_for_readers_off_device(self):
+        """With skip_h2d_for_method_inputs=True, a method input that a device
+        delegate reads stays on that device. A host op, a delegate on the host,
+        or a delegate on another device that also reads it reads one _d2h_copy
+        of it, also when more than one delegate on the device reads it, and a
+        warning says so."""
 
         class HostOp(torch.nn.Module):
             def forward(self, a, b):
@@ -2133,6 +2138,10 @@ class TwoDelegates(torch.nn.Module):
             def forward(self, a, b, c):
                 return torch.add(a, b), torch.add(a, c)
 
+        class TwoDelegatesAndHostOp(torch.nn.Module):
+            def forward(self, a, b, c):
+                return torch.add(a, b), torch.add(a, c), torch.sin(a)
+
         class CopiedIntoBuffer(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -2150,56 +2159,128 @@ def forward(self, a, b):
             enable_non_cpu_memory_planning=True,
         )
         cases = {
-            "host op": (
-                HostOp(),
-                2,
-                DeviceAwarePartitioner("cuda:0"),
-                "aten_sin_default",
+            "host op": (HostOp(), 2, DeviceAwarePartitioner("cuda:0")),
+            "two delegates on the device and a host op": (
+                TwoDelegatesAndHostOp(),
+                3,
+                PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:0"]),
             ),
-            "other device index": (
+            "delegate without a device": (
                 TwoDelegates(),
                 3,
-                PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1"]),
-                "executorch_call_delegate_1",
+                PerAddDeviceDelegatePartitioner(["cuda:0", None]),
             ),
-            "other device type": (
+            "delegate on the cpu": (
                 TwoDelegates(),
                 3,
                 PerAddDeviceDelegatePartitioner(["cuda:0", "cpu"]),
-                "executorch_call_delegate_1",
+            ),
+            "other device index": (
+                TwoDelegates(),
+                3,
+                PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1"]),
             ),
             # The write-back of the buffer only reads the input.
             "buffer write-back": (
                 CopiedIntoBuffer(),
                 2,
                 DeviceAwarePartitioner("cuda:0"),
-                "copy_",
             ),
         }
-        for case, (model, num_inputs, partitioner, reader) in cases.items():
-            ep = export(model, tuple(torch.randn(2, 2) for _ in range(num_inputs)))
-            # Each pipeline gets its own check, so the first refusal cannot hide
-            # the second pipeline.
-            for pipeline in ("to_edge+to_backend", "to_edge_transform_and_lower"):
-                if pipeline == "to_edge+to_backend":
-                    edge = to_edge(
-                        deepcopy(ep),
-                        compile_config=EdgeCompileConfig(_check_ir_validity=False),
-                    ).to_backend(partitioner)
-                else:
-                    edge = to_edge_transform_and_lower(
-                        deepcopy(ep), partitioner=[partitioner]
-                    )
+        for case, (model, num_inputs, partitioner) in cases.items():
+            inputs = tuple(torch.randn(2, 2) for _ in range(num_inputs))
+            with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs:
+                results = self._get_executorch_program(
+                    model, inputs, partitioner, et_config
+                )
+            self.assertIn("Method input 'a' (16 bytes) stays on cuda:0", logs.output[0])
+            for pipeline, program, gm in results:
                 with self.subTest(case=case, pipeline=pipeline):
-                    with self.assertRaisesRegex(
-                        RuntimeError,
-                        f"placeholder 'a' on cuda:0, because '{reader}' also reads it",
-                    ):
-                        edge.to_executorch(deepcopy(et_config))
+                    a = next(n for n in gm.graph.nodes if n.name == "a")
+                    self.assertEqual(a.meta["spec"].device, DeviceType.CUDA)
+                    self.assertEqual(a.meta["spec"].device_index, 0)
+                    nodes = _collect_device_copy_nodes(gm)
+                    self.assertIs(nodes.delegate_nodes[0].args[1], a)
+                    # The copy is the only other reader of the input, and is
+                    # not copied again.
+                    (copy,) = [n for n in a.users if n not in nodes.delegate_nodes]
+                    self.assertIn(copy, nodes.d2h_nodes)
+                    self.assertFalse(set(copy.users) & set(nodes.d2h_nodes))
+                    self._assert_nodes_buffer_device(
+                        [copy], program, DeviceType.CPU, pipeline, "D2H copy"
+                    )
+
+    def test_skip_h2d_shape_reader_needs_no_copy(self):
+        """With skip_h2d_for_method_inputs=True, an op that reads only the shape
+        of an input on the device reads the input there, since its sizes stay
+        in host memory, so the input gets no copy and no warning."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.add(x, y), torch.zeros(x.shape[0])
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        batch = torch.export.Dim("batch", min=2, max=8)
+        ep = export(
+            Model(),
+            (torch.randn(4, 2), torch.randn(4, 2)),
+            dynamic_shapes={"x": {0: batch}, "y": {0: batch}},
+        )
+        edge = to_edge(
+            ep, compile_config=EdgeCompileConfig(_check_ir_validity=False)
+        ).to_backend(DeviceAwarePartitioner("cuda:0"))
+        with self.assertNoLogs(propagate_device_pass.logger, "WARNING"):
+            gm = edge.to_executorch(et_config).exported_program().graph_module
+        (size,) = [n for n in gm.graph.nodes if n.target == torch.ops.aten.sym_size.int]
+        read = size.args[0]
+        self.assertEqual(read.op, "placeholder")
+        self.assertEqual(read.meta["spec"].device, DeviceType.CUDA)
+        self.assertFalse(
+            set(read.users) & set(_collect_device_copy_nodes(gm).d2h_nodes)
+        )
+
+    def test_skip_both_returns_input_on_device_without_copy(self):
+        """With both skip flags, a method that returns an input a device
+        delegate also reads returns it on the device, with no copy and no
+        warning."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.add(x, y), x
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True, skip_d2h_for_method_outputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        with self.assertNoLogs(propagate_device_pass.logger, "WARNING"):
+            results = self._get_executorch_program(
+                Model(),
+                (torch.randn(2, 2), torch.randn(2, 2)),
+                DeviceAwarePartitioner("cuda:0"),
+                et_config,
+            )
+        for pipeline, _program, gm in results:
+            with self.subTest(pipeline=pipeline):
+                nodes = _collect_device_copy_nodes(gm)
+                self.assertEqual(nodes.h2d_nodes + nodes.d2h_nodes, [])
+                x = next(n for n in gm.graph.nodes if n.name == "x")
+                self.assertIn(x, gm.graph.output_node().args[0])
+                self.assertEqual(x.meta["spec"].device, DeviceType.CUDA)
 
     def test_skip_h2d_refuses_placeholder_written_on_host(self):
         """A mutated input or mutable buffer is written back on the host, so the
-        refusal says so, also when a host op reads it too."""
+        refusal says so, also when a host op reads it too. An input that a host
+        op writes in place is refused as well, since the write would only reach
+        its host copy."""
 
         class MutableBuffer(torch.nn.Module):
             def __init__(self):
@@ -2228,6 +2309,12 @@ def forward(self, a, x):
                 x.add_(a)
                 return out
 
+        class InputWrittenByHostOp(torch.nn.Module):
+            def forward(self, a, x):
+                out = torch.add(a, x)
+                torch.ops.propagate_device_test.accumulate(x, a)
+                return out
+
         et_config = ExecutorchBackendConfig(
             emit_stacktrace=False,
             propagate_device_config=PropagateDeviceConfig(
@@ -2235,21 +2322,32 @@ def forward(self, a, x):
             ),
             enable_non_cpu_memory_planning=True,
         )
+        write_back = r"'copy_' writes it back on the host \(the method mutates it\)"
         cases = {
-            "mutable buffer": (MutableBuffer(), 1, "b_state"),
-            "mutable buffer read on host": (MutableBufferReadOnHost(), 1, "b_state"),
-            "mutated input": (MutatedInput(), 2, "x"),
+            "mutable buffer": (MutableBuffer(), 1, "b_state", write_back),
+            "mutable buffer read on host": (
+                MutableBufferReadOnHost(),
+                1,
+                "b_state",
+                write_back,
+            ),
+            "mutated input": (MutatedInput(), 2, "x", write_back),
+            "input written by a host op": (
+                InputWrittenByHostOp(),
+                2,
+                "x",
+                "'[a-z_]*accumulate[a-z_]*' writes it in place on the host",
+            ),
         }
-        for case, (model, num_inputs, name) in cases.items():
+        for case, (model, num_inputs, name, reason) in cases.items():
             inputs = tuple(torch.randn(2, 2) for _ in range(num_inputs))
-            edge = to_edge(export(model, inputs)).to_backend(
-                DeviceAwarePartitioner("cuda:0")
-            )
+            edge = to_edge(
+                export(model, inputs),
+                compile_config=EdgeCompileConfig(_check_ir_validity=False),
+            ).to_backend(DeviceAwarePartitioner("cuda:0"))
             with self.subTest(case=case):
                 with self.assertRaisesRegex(
-                    RuntimeError,
-                    f"placeholder '{name}' on cuda:0, because 'copy_' writes it "
-                    r"back on the host \(the method mutates it\)",
+                    RuntimeError, f"placeholder '{name}' on cuda:0, because {reason}"
                 ):
                     edge.to_executorch(deepcopy(et_config))
 

From ad9e14ce34c7b8b474237c80ce8f7c12f62303c3 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 25 Oct 2026 16:46:25 -0700
Subject: [PATCH 20/26] Add error_on_per_call_copies to make copies the skip
 flags keep an error

With skip_h2d_for_method_inputs and skip_d2h_for_method_outputs, users
expect no copies at the method boundary. Three copies still stay and run
on every call, and to_executorch only logs a warning for them: a constant
that feeds a device delegate from outside it, a method input on the device
that a reader off that device also reads, and a returned output that a
host op writes in place, which is copied to the host and back. A warning
is easy to miss in a long export log.

The new opt-in PropagateDeviceConfig option error_on_per_call_copies makes
to_executorch raise with the same message instead. It applies to a
constant copy only with skip_h2d_for_method_inputs, since without it every
input is copied anyway. The input copy exists only with that flag, and the
written output copy only with skip_d2h_for_method_outputs. The messages
now format sizes with an f-string, as the error and the warning share them.

The new test expects the error for all three cases, and only the warning
for a constant copy without the skip flag. It fails without this change.
---
 docs/source/backends/cuda/cuda-overview.md |  4 ++
 exir/passes/propagate_device_config.py     | 11 ++++
 exir/passes/propagate_device_pass.py       | 71 ++++++++++++----------
 exir/program/_program.py                   |  1 +
 exir/tests/test_propagate_device_pass.py   | 57 +++++++++++++++++
 5 files changed, 113 insertions(+), 31 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index 69e8724507f..2696ec2945e 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -193,6 +193,10 @@ as a KV cache update. It writes the host copy, so `to_executorch` returns a devi
 written tensor. That copies the output to the host and back on every call and logs a warning. Put
 that operator in the delegate to avoid it.
 
+**Set `error_on_per_call_copies=True` in `PropagateDeviceConfig` to turn these warnings into
+errors** when the matching skip flag is set, so `to_executorch` fails instead of adding a copy that
+runs on every call.
+
 **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the
 program reserves its own input buffer and the runtime copies the caller's memory into it, so the
 copy comes back at run time. On device memory that copy is a host memcpy into a device pointer,
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index a688e74b141..9e6fe92408e 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -77,11 +77,21 @@ class PropagateDeviceConfig:
     # A dict can be used to set per-method values, keyed by method name.
     skip_d2h_for_method_outputs: Union[bool, Dict[str, bool]] = False
 
+    # When True, the copies that the skip flags above cannot remove make
+    # to_executorch raise instead of logging a warning: with
+    # skip_h2d_for_method_inputs, a constant copied to the device and a method
+    # input copied to the host for a reader off its device, and with
+    # skip_d2h_for_method_outputs, a returned output copied to the host and
+    # back. Each runs on every call. Use it to make sure a program that skips
+    # its boundary copies has none left.
+    error_on_per_call_copies: bool = False
+
     def __hash__(self) -> int:
         return hash(
             (
                 str(self.skip_h2d_for_method_inputs),
                 str(self.skip_d2h_for_method_outputs),
+                self.error_on_per_call_copies,
             )
         )
 
@@ -91,4 +101,5 @@ def __eq__(self, other: object) -> bool:
         return (
             self.skip_h2d_for_method_inputs == other.skip_h2d_for_method_inputs
             and self.skip_d2h_for_method_outputs == other.skip_d2h_for_method_outputs
+            and self.error_on_per_call_copies == other.error_on_per_call_copies
         )
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index c041be2a6f1..bca0dcf5946 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -308,6 +308,10 @@ class PropagateDevicePass(PassBase):
       on device, not its D2H copy; the copy stays only for other readers. When
       one of them writes the copy in place, the graph output reads a device
       copy of the written tensor instead.
+    - error_on_per_call_copies: Raise instead of logging a warning when a
+      copy that a skip flag asked to avoid has to stay: a constant copied to a
+      device, a method input copied for a reader off its device, or a
+      returned output written on the host.
     """
 
     def __init__(
@@ -315,11 +319,13 @@ def __init__(
         skip_h2d_for_method_inputs: bool = False,
         skip_d2h_for_method_outputs: bool = False,
         enable_non_cpu_memory_planning: bool = False,
+        error_on_per_call_copies: bool = False,
     ) -> None:
         super().__init__()
         self.skip_h2d_for_method_inputs = skip_h2d_for_method_inputs
         self.skip_d2h_for_method_outputs = skip_d2h_for_method_outputs
         self.enable_non_cpu_memory_planning = enable_non_cpu_memory_planning
+        self.error_on_per_call_copies = error_on_per_call_copies
 
         if (
             skip_h2d_for_method_inputs or skip_d2h_for_method_outputs
@@ -347,6 +353,15 @@ def _delegate_device(
             return None
         return _get_target_device_from_compile_specs(lowered_module)
 
+    def _report_per_call_copy(self, message: str, skip_requested: bool) -> None:
+        """Warn about a copy that runs on every call. When a skip flag asked
+        for no such copy and error_on_per_call_copies is set, raise instead."""
+        if skip_requested and self.error_on_per_call_copies:
+            raise RuntimeError(
+                f"{message} Raised because error_on_per_call_copies=True."
+            )
+        logger.warning("%s", message)
+
     def _insert_h2d_copies(
         self,
         graph_module: torch.fx.GraphModule,
@@ -398,6 +413,15 @@ def _insert_h2d_copies(
 
             h2d_node = _reusable_copy(arg, node, target_device_type, device_index)
             if h2d_node is None:
+                if arg_spec.const:
+                    self._report_per_call_copy(
+                        f"Constant '{arg.name}' ({arg_spec.nbytes()} bytes) feeds a "
+                        f"delegate on {target_device_type.name.lower()}:"
+                        f"{device_index} from outside it, so it is copied to the "
+                        f"device on every call. Have the partitioner take it into "
+                        f"the delegate to avoid that copy.",
+                        self.skip_h2d_for_method_inputs,
+                    )
                 with graph_module.graph.inserting_before(node):
                     h2d_node = graph_module.graph.call_function(
                         torch.ops.et_copy._h2d_copy.default,
@@ -409,17 +433,6 @@ def _insert_h2d_copies(
                     h2d_node.meta["val"] = arg.meta.get("val")
                     if "tensor_meta" in arg.meta:
                         h2d_node.meta["tensor_meta"] = arg.meta["tensor_meta"]
-                if arg_spec.const:
-                    logger.warning(
-                        "Constant '%s' (%d bytes) feeds a delegate on %s:%d from "
-                        "outside it, so it is copied to the device on every call. "
-                        "Have the partitioner take it into the delegate to avoid "
-                        "that copy.",
-                        arg.name,
-                        arg_spec.nbytes(),
-                        target_device_type.name.lower(),
-                        device_index,
-                    )
             new_args[i] = h2d_node
             changed = True
 
@@ -493,16 +506,13 @@ def _place_input_on_device(
             return
         position = {n: i for i, n in enumerate(graph_module.graph.nodes)}
         first_reader = min(other_readers, key=position.__getitem__)
-        logger.warning(
-            "Method input '%s' (%s bytes) stays on %s for the delegates there, "
-            "but '%s' also reads it, so it is copied to the host on every call. "
-            "Have the partitioner put every reader in a delegate on %s to avoid "
-            "that copy.",
-            placeholder.name,
-            spec.nbytes(),
-            device,
-            first_reader.name,
-            device,
+        self._report_per_call_copy(
+            f"Method input '{placeholder.name}' ({spec.nbytes()} bytes) stays on "
+            f"{device} for the delegates there, but '{first_reader.name}' also "
+            f"reads it, so it is copied to the host on every call. Have the "
+            f"partitioner put every reader in a delegate on {device} to avoid "
+            f"that copy.",
+            True,
         )
         with graph_module.graph.inserting_before(first_reader):
             d2h_node = graph_module.graph.call_function(
@@ -615,16 +625,15 @@ def _keep_method_outputs_on_device(
                     returned.meta["val"] = node.meta.get("val")
                     if "tensor_meta" in node.meta:
                         returned.meta["tensor_meta"] = node.meta["tensor_meta"]
-                logger.warning(
-                    "A method output (%d bytes) comes from a delegate on %s:%d "
-                    "and '%s' writes it in place on the host, so it is copied to "
-                    "the host and back to the device on every call, even with "
-                    "skip_d2h_for_method_outputs. Have the partitioner take that "
-                    "op into the delegate to avoid those copies.",
-                    device_spec.nbytes(),
-                    device_spec.device.name.lower(),
-                    device_spec.device_index,
-                    writers[0].target,
+                self._report_per_call_copy(
+                    f"A method output ({device_spec.nbytes()} bytes) comes from a "
+                    f"delegate on {device_spec.device.name.lower()}:"
+                    f"{device_spec.device_index} and '{writers[0].target}' writes "
+                    f"it in place on the host, so it is copied to the host and back "
+                    f"to the device on every call, even with "
+                    f"skip_d2h_for_method_outputs. Have the partitioner take that "
+                    f"op into the delegate to avoid those copies.",
+                    True,
                 )
             _replace_spec(
                 [graph_module.graph.output_node()],
diff --git a/exir/program/_program.py b/exir/program/_program.py
index f2084235bda..b1d896281c0 100644
--- a/exir/program/_program.py
+++ b/exir/program/_program.py
@@ -793,6 +793,7 @@ def edge_to_executorch_passes(
             skip_h2d_for_method_inputs=device_cfg.skip_h2d_for_method_inputs,
             skip_d2h_for_method_outputs=device_cfg.skip_d2h_for_method_outputs,
             enable_non_cpu_memory_planning=config.enable_non_cpu_memory_planning,
+            error_on_per_call_copies=device_cfg.error_on_per_call_copies,
         ),
         EdgeToBackendOpsPass(),
         RemoveGraphAssertsPass(),
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index f8915cd6c76..b29884bf9b7 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -1910,6 +1910,63 @@ def forward(self, a):
                         nodes.h2d_nodes, program, DeviceType.CUDA, pipeline, "H2D copy"
                     )
 
+    def test_error_on_per_call_copies(self):
+        """error_on_per_call_copies turns the warning about a copy that a skip
+        flag cannot remove into an error: a constant copied to the device and
+        an input copied for a host reader under skip_h2d_for_method_inputs, and
+        a returned output written on the host under skip_d2h_for_method_outputs.
+        Without the skip flag, a constant copy only logs the warning."""
+
+        class ReadsConstant(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("offset", torch.randn(2, 2))
+
+            def forward(self, a):
+                return torch.add(a, self.offset), self.offset * 2
+
+        class WritesOutputOnHost(torch.nn.Module):
+            def forward(self, x, y):
+                a = x + y
+                torch.ops.propagate_device_test.accumulate(a, y)
+                return a
+
+        class ReadsInputOnHost(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y, torch.sin(x)
+
+        def lower(model, num_inputs, **flags):
+            return _lower_with_host_writer(
+                model,
+                tuple(torch.randn(2, 2) for _ in range(num_inputs)),
+                ["cuda:0"],
+                _host_writer_config(
+                    False,
+                    propagate_device_config=PropagateDeviceConfig(
+                        error_on_per_call_copies=True, **flags
+                    ),
+                ),
+            )
+
+        cases = {
+            "constant": (ReadsConstant(), 1, "skip_h2d_for_method_inputs"),
+            "input read on host": (ReadsInputOnHost(), 2, "skip_h2d_for_method_inputs"),
+            "output written on host": (
+                WritesOutputOnHost(),
+                2,
+                "skip_d2h_for_method_outputs",
+            ),
+        }
+        for case, (model, num_inputs, flag) in cases.items():
+            with self.subTest(case=case):
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    "on every call.*Raised because error_on_per_call_copies=True",
+                ):
+                    lower(model, num_inputs, **{flag: True})
+        with self.assertLogs(propagate_device_pass.logger, "WARNING"):
+            lower(ReadsConstant(), 1)
+
     def test_constant_copy_keeps_stored_layout(self):
         """The copy of a constant moves its bytes as they are, so it must use the
         layout the constant is stored in: contiguous for a transposed parameter,

From 2a8bff8bd93196611efd5277e34ae3fe6a922fb8 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 26 Oct 2026 08:48:45 -0700
Subject: [PATCH 21/26] Shorten the device copy notes in the CUDA docs

The notes this change adds to the CUDA backend page, about the copies
that stay with skip_h2d_for_method_inputs and skip_d2h_for_method_outputs,
were long paragraphs that explained the internals. They are now a short
list of the copies that stay, how to remove them, how to make them an
error, and which inputs the skip flag refuses. They also move after the
other notes on the Python export options.
---
 docs/source/backends/cuda/cuda-overview.md | 48 +++++++++-------------
 1 file changed, 20 insertions(+), 28 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index 2696ec2945e..40d809ed07e 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -169,34 +169,6 @@ out works, and the snippet above sets it explicitly only to make the requirement
 it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens
 during device-aware memory planning.
 
-**With `skip_h2d_for_method_inputs`, an input that a device delegate reads stays on that device.**
-If delegates on more than one device read it, that is the device of the first one in the graph.
-An operator that reads only its shape reads it there too. Anything else that reads it, such as an
-operator outside the delegates or a delegate on another device, reads a copy of it in host memory,
-made on every call, and `to_executorch` logs a warning. Have the partitioner put those operators in
-a delegate on that device to avoid that copy. `to_executorch` refuses an input that the method
-changes, because the change is written back on the host, and an input that an operator outside the
-delegates writes in place. For a mutable buffer, have the partitioner put the buffer and every
-operator that reads or writes it in the delegate. `tag_mutated_buffer` tags the buffer once those
-operators are there, and the CUDA partitioner already calls it.
-
-**A constant that a delegate reads from outside is still copied.** Constants, such as weights and
-buffers the method does not change, are not passed by the caller. Their data stays in host memory,
-in the program or in its external data file, so `to_executorch` adds a copy to the device that runs
-on every call, one for each device that reads the constant, and logs a warning. Each copy stays in
-device memory until the last delegate on that device has read it. Keep constants inside the
-delegate to avoid that cost.
-
-**With `skip_d2h_for_method_outputs`, a returned output that a custom operator outside the
-delegate writes in place is still copied.** This is a custom operator with a mutable argument, such
-as a KV cache update. It writes the host copy, so `to_executorch` returns a device copy of the
-written tensor. That copies the output to the host and back on every call and logs a warning. Put
-that operator in the delegate to avoid it.
-
-**Set `error_on_per_call_copies=True` in `PropagateDeviceConfig` to turn these warnings into
-errors** when the matching skip flag is set, so `to_executorch` fails instead of adding a copy that
-runs on every call.
-
 **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the
 program reserves its own input buffer and the runtime copies the caller's memory into it, so the
 copy comes back at run time. On device memory that copy is a host memcpy into a device pointer,
@@ -236,6 +208,26 @@ Do not pass a dict to `skip_h2d_for_method_inputs` itself. The pass reads that f
 truthiness, so any non-empty dict enables the skip for every method regardless of the values
 inside it.
 
+**Some copies stay even with the skip flags.** In the cases below, `to_executorch` keeps a copy
+that runs on every call and logs a warning. To remove the copy, have the partitioner put more of
+the graph in the delegate.
+
+- An input stays on the device of the first delegate that reads it. Other readers, such as an
+  operator outside the delegate or a delegate on another device, read a copy in host memory.
+- A constant, such as a weight, that a delegate reads from outside is copied to each device that
+  reads it.
+- A returned output that a custom operator outside the delegate writes in place, such as a KV
+  cache update, is copied to the host and back.
+
+Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail
+instead of warning. It fails only on a copy that one of the skip flags you set was meant to remove.
+
+**`skip_h2d_for_method_inputs` refuses inputs that the method changes.** `to_executorch` fails if
+the method changes an input, even when the delegate makes the change. This includes a mutable
+buffer, such as a KV cache, even though the caller does not pass it. For a buffer, put it and every
+operator that uses it in the delegate. The CUDA partitioner already moves the buffer in, with
+`tag_mutated_buffer`, once those operators are there.
+
 **The choice is baked into the `.pte`.** A program exported with copies expects host tensors; one
 exported without them expects device tensors. Passing the wrong kind is a caller error, not
 something the runtime corrects, so the two are not interchangeable at run time.

From c367248ac6de74586c4889b1d40b629a2e1a0198 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 26 Oct 2026 11:16:52 -0700
Subject: [PATCH 22/26] Return an output written on the host from the host

With skip_d2h_for_method_outputs, the method returns delegate outputs on
the device. When a custom host op with a mutable argument, such as a KV
cache update, writes the host copy of such an output in place, the device
tensor never sees the write. The pass used to copy the written host
tensor back to the device and return that copy, so every call paid for a
copy to the host and another one back.

The written host tensor is already the right result, so the method now
returns it as it is, and the method output says it is in host memory.
The copy to the host stays, because the host op needs it. The warning
about it stays too, and error_on_per_call_copies still turns it into an
error.
---
 docs/source/backends/cuda/cuda-overview.md |  3 +-
 exir/passes/propagate_device_config.py     | 16 +++++------
 exir/passes/propagate_device_pass.py       | 32 ++++++---------------
 exir/tests/test_propagate_device_pass.py   | 33 +++++++++++++---------
 4 files changed, 37 insertions(+), 47 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index 40d809ed07e..3a40e71a0a4 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -217,7 +217,8 @@ the graph in the delegate.
 - A constant, such as a weight, that a delegate reads from outside is copied to each device that
   reads it.
 - A returned output that a custom operator outside the delegate writes in place, such as a KV
-  cache update, is copied to the host and back.
+  cache update, is copied to the host and returned there. The method metadata gives the device
+  of each output.
 
 Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail
 instead of warning. It fails only on a copy that one of the skip flags you set was meant to remove.
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index 9e6fe92408e..9a5c0dbfd53 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -65,11 +65,11 @@ class PropagateDeviceConfig:
     skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False
 
     # When True, device delegate outputs that the method returns stay on
-    # device: the method output never reads a _d2h_copy of them. When a custom
-    # host op with a mutable argument, such as a KV cache update, writes the
-    # host copy of such an output in place, the method returns a device copy of
-    # the written tensor. That costs a copy to the host and one back on every
-    # call, and the pass logs a warning about it.
+    # device, with one exception: when a custom host op with a mutable
+    # argument, such as a KV cache update, writes the host copy of such an
+    # output in place, the method returns that host copy, since the device
+    # tensor never sees the write. The copy to the host runs on every call,
+    # and the pass logs a warning about it.
     # Useful for cross-method GPU pipelines where the next method consumes GPU
     # tensors directly.
     #
@@ -81,9 +81,9 @@ class PropagateDeviceConfig:
     # to_executorch raise instead of logging a warning: with
     # skip_h2d_for_method_inputs, a constant copied to the device and a method
     # input copied to the host for a reader off its device, and with
-    # skip_d2h_for_method_outputs, a returned output copied to the host and
-    # back. Each runs on every call. Use it to make sure a program that skips
-    # its boundary copies has none left.
+    # skip_d2h_for_method_outputs, a returned output copied to the host for a
+    # host op that writes it. Each runs on every call. Use it to make sure a
+    # program that skips its boundary copies has none left.
     error_on_per_call_copies: bool = False
 
     def __hash__(self) -> int:
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index bca0dcf5946..5e821300b3d 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -306,8 +306,8 @@ class PropagateDevicePass(PassBase):
       data is stored on the host, so they still get an H2D copy.
     - skip_d2h_for_method_outputs: The graph output reads a delegate output
       on device, not its D2H copy; the copy stays only for other readers. When
-      one of them writes the copy in place, the graph output reads a device
-      copy of the written tensor instead.
+      one of them writes the copy in place, the graph output reads the written
+      host copy instead.
     - error_on_per_call_copies: Raise instead of logging a warning when a
       copy that a skip flag asked to avoid has to stay: a constant copied to a
       device, a method input copied for a reader off its device, or a
@@ -593,8 +593,8 @@ def _keep_method_outputs_on_device(
 
         The copy stays for its other users, such as a host op. When one of them
         writes the copy in place, the device tensor never sees that write, so the
-        output reads a new device copy of the written host copy instead, and a
-        warning says so.
+        output keeps reading the written host copy, and a warning says so.
+        Copying it back to the device would add a second copy on every call.
         """
         d2h = torch.ops.et_copy._d2h_copy.default
         for node in list(graph_module.graph.nodes):
@@ -608,33 +608,17 @@ def _keep_method_outputs_on_device(
             writers = _in_place_writers(node)
             if writers:
                 device_spec = returned.meta["spec"]
-                output_node = graph_module.graph.output_node()
-                # A later delegate on this device may already read a device copy
-                # of the written host copy.
-                returned = _reusable_copy(
-                    node, output_node, device_spec.device, device_spec.device_index
-                )
-                if returned is None:
-                    with graph_module.graph.inserting_before(output_node):
-                        returned = graph_module.graph.call_function(
-                            torch.ops.et_copy._h2d_copy.default, (node,)
-                        )
-                    returned.meta["spec"] = _copy_output_spec(
-                        node.meta["spec"], device_spec.device, device_spec.device_index
-                    )
-                    returned.meta["val"] = node.meta.get("val")
-                    if "tensor_meta" in node.meta:
-                        returned.meta["tensor_meta"] = node.meta["tensor_meta"]
                 self._report_per_call_copy(
                     f"A method output ({device_spec.nbytes()} bytes) comes from a "
                     f"delegate on {device_spec.device.name.lower()}:"
                     f"{device_spec.device_index} and '{writers[0].target}' writes "
-                    f"it in place on the host, so it is copied to the host and back "
-                    f"to the device on every call, even with "
+                    f"it in place on the host, so it is copied to the host on every "
+                    f"call and returned there, even with "
                     f"skip_d2h_for_method_outputs. Have the partitioner take that "
-                    f"op into the delegate to avoid those copies.",
+                    f"op into the delegate to avoid that copy.",
                     True,
                 )
+                continue
             _replace_spec(
                 [graph_module.graph.output_node()],
                 node.meta["spec"],
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index b29884bf9b7..ec154140abf 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -978,12 +978,12 @@ def forward(self, x, y):
         self.assertIsNone(output.meta["spec"].mem_id)
         self.assertIsNotNone(d2h.meta["spec"].mem_id)
 
-    def test_skip_d2h_returns_device_copy_of_output_written_on_host(self):
+    def test_skip_d2h_returns_output_written_on_host_from_host(self):
         """With skip_d2h_for_method_outputs, a returned delegate output whose
         host copy a host op writes in place, directly or through a view, is
-        returned as a device copy of the written host copy, because the device
-        tensor never sees the write. A warning names the op that writes it,
-        since these copies run on every call."""
+        returned as that written host copy, because the device tensor never
+        sees the write. No copy goes back to the device. A warning names the
+        op that writes it, since the copy to the host runs on every call."""
 
         class Model(torch.nn.Module):
             def __init__(self, write_view):
@@ -1017,17 +1017,22 @@ def forward(self, x, y):
                     "'propagate_device_test.accumulate.default' writes it", warning
                 )
                 (output,) = gm.graph.output_node().args[0]
-                self.assertEqual(output.target, torch.ops.et_copy._h2d_copy.out)
-                self.assertEqual(output.meta["spec"].device, DeviceType.CUDA)
-                self.assertIs(output.args[0], _host_written(gm))
-                nodes = list(gm.graph.nodes)
-                self.assertLess(nodes.index(_host_writer(gm)), nodes.index(output))
+                self.assertIs(output, _host_written(gm))
+                self.assertEqual(output.target, torch.ops.et_copy._d2h_copy.out)
+                (output_spec,) = gm.graph.output_node().meta["spec"]
+                self.assertEqual(output_spec.device, DeviceType.CPU)
+                self.assertFalse(
+                    any(
+                        copy.args[0] is output
+                        for copy in _collect_device_copy_nodes(gm).h2d_nodes
+                    )
+                )
 
-    def test_skip_d2h_shares_the_device_copy_of_a_written_output(self):
+    def test_skip_d2h_copies_a_written_output_only_for_a_later_delegate(self):
         """With skip_d2h_for_method_outputs, a returned output that a host op
-        writes before a later delegate on the same device reads it needs a
-        device copy for both. The method output reads the same copy as that
-        delegate."""
+        writes before a later delegate on the same device reads it is copied
+        back to the device for that delegate only. The method returns the
+        written host copy."""
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -1053,7 +1058,7 @@ def forward(self, x, y):
         ]
         _, second = _collect_device_copy_nodes(gm).delegate_nodes
         self.assertIn(copy, second.args[1:])
-        self.assertIs(gm.graph.output_node().args[0][0], copy)
+        self.assertIs(gm.graph.output_node().args[0][0], _host_written(gm))
 
     def test_cross_device_copies_preserved(self):
         """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two

From 4041857a92f9855b5fcd31b97bccf34939c2e061 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 26 Oct 2026 12:14:49 -0700
Subject: [PATCH 23/26] Describe every per-call copy left by the skip flags
 correctly

The docs and the config comment said that a returned output goes back
to the host only when a custom op writes it in place. The pass does
this for any op outside the delegate that writes it, aten in-place ops
included, so the text now says that.

The warning for a method input that stays on the device was also off
in two cases. When the method returns the input, the advice to put
every reader in a delegate cannot work, and skip_d2h_for_method_outputs
is what removes the copy. When a delegate on another device reads the
input, it gets a second copy, from the host to its device, which the
warning did not mention. The warning, the config comment and the docs
now cover both cases.
---
 docs/source/backends/cuda/cuda-overview.md |  8 +++++---
 exir/passes/propagate_device_config.py     | 20 +++++++++++---------
 exir/passes/propagate_device_pass.py       |  7 ++++---
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index 3a40e71a0a4..45feff2684b 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -213,11 +213,13 @@ that runs on every call and logs a warning. To remove the copy, have the partiti
 the graph in the delegate.
 
 - An input stays on the device of the first delegate that reads it. Other readers, such as an
-  operator outside the delegate or a delegate on another device, read a copy in host memory.
+  operator outside the delegate, a delegate on another device, or the method output without
+  `skip_d2h_for_method_outputs`, read a copy in host memory. A delegate on another device then
+  gets a second copy, from the host.
 - A constant, such as a weight, that a delegate reads from outside is copied to each device that
   reads it.
-- A returned output that a custom operator outside the delegate writes in place, such as a KV
-  cache update, is copied to the host and returned there. The method metadata gives the device
+- A returned output that an operator outside the delegate writes in place, such as a KV cache
+  update, is copied to the host and returned there. The method metadata gives the device
   of each output.
 
 Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index 9a5c0dbfd53..a6e87b9de40 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -50,10 +50,12 @@ class PropagateDeviceConfig:
     # delegate are NOT wrapped with _h2d_copy. The user must provide tensors
     # already on the target device: the device of the first delegate that
     # reads the input, in graph order. Any other reader of such an input, such
-    # as a host op or a delegate on another device, reads a copy of it in host
-    # memory, made on every call, and the pass logs a warning. An input that an
-    # op writes in place on the host makes to_executorch raise. Useful for
-    # pipelines where inputs are pre-staged on GPU. Constants are not passed
+    # as a host op, a delegate on another device, or the method output without
+    # skip_d2h_for_method_outputs, reads a copy of it in host memory, made on
+    # every call, and the pass logs a warning. A delegate on another device
+    # then gets a second copy, from the host. An input that an op writes in
+    # place on the host makes to_executorch raise. Useful for pipelines where
+    # inputs are pre-staged on GPU. Constants are not passed
     # by the caller, so they are still copied: their data stays on the host,
     # in the program or in its external data file. That copy runs on every
     # call, once for each device that reads the constant, and stays in device
@@ -65,11 +67,11 @@ class PropagateDeviceConfig:
     skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False
 
     # When True, device delegate outputs that the method returns stay on
-    # device, with one exception: when a custom host op with a mutable
-    # argument, such as a KV cache update, writes the host copy of such an
-    # output in place, the method returns that host copy, since the device
-    # tensor never sees the write. The copy to the host runs on every call,
-    # and the pass logs a warning about it.
+    # device, with one exception: when an op outside the delegate, such as a
+    # KV cache update, writes the host copy of such an output in place, the
+    # method returns that host copy, since the device tensor never sees the
+    # write. The copy to the host runs on every call, and the pass logs a
+    # warning about it.
     # Useful for cross-method GPU pipelines where the next method consumes GPU
     # tensors directly.
     #
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index 5e821300b3d..f9444cc620e 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -509,9 +509,10 @@ def _place_input_on_device(
         self._report_per_call_copy(
             f"Method input '{placeholder.name}' ({spec.nbytes()} bytes) stays on "
             f"{device} for the delegates there, but '{first_reader.name}' also "
-            f"reads it, so it is copied to the host on every call. Have the "
-            f"partitioner put every reader in a delegate on {device} to avoid "
-            f"that copy.",
+            f"reads it, so it is copied to the host on every call, and from there "
+            f"to any other device that reads it. To avoid that, have the "
+            f"partitioner put every op that reads it in a delegate on {device}, "
+            f"and set skip_d2h_for_method_outputs if the method returns it.",
             True,
         )
         with graph_module.graph.inserting_before(first_reader):

From d517e6783f211fbd4b623f67f0e3465dcb9c4716 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 26 Oct 2026 12:14:49 -0700
Subject: [PATCH 24/26] Fix a test comment about how often to_edge converts a
 constant

With _skip_dim_order, to_edge runs convert_constant_dim_order_pass once,
not twice, so the comment now covers both modes.
---
 exir/tests/test_passes.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index b338dc2d3ea..fe3bc675ff4 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -3195,7 +3195,8 @@ def forward(self, x):
                             _skip_dim_order=skip_dim_order
                         ),
                     ).exported_program()
-                # to_edge runs the pass twice, but only the first run converts.
+                # Only one run of the pass converts: the first, or with
+                # _skip_dim_order the only one.
                 self.assertEqual(len(captured.output), 1)
                 self.assertIn("Parameter 'weight'", captured.output[0])
                 self.assertTrue(ep.state_dict["weight"].is_contiguous())

From 4c56ae2d0a1b7d37c5abcc36aa46e41c778b9df8 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 26 Oct 2026 13:23:28 -0700
Subject: [PATCH 25/26] Pin auto_functionalized v2 in the host writer tests

Some PyTorch builds turn off auto_functionalized v2 for export. There, a
custom op that writes through a view writes into a view_copy, not into an
as_strided of the tensor, so the tests that write through a view see a
different graph and fail. Under v1 that write is lost even without these
passes, so there is nothing correct to check there. Functionalize with v2 in
these tests so they check the same graph in every build.
---
 exir/tests/test_propagate_device_pass.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index ec154140abf..893a9d2cefa 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -197,10 +197,16 @@ def _lower_with_host_writer(
     et_config: ExecutorchBackendConfig,
 ) -> torch.fx.GraphModule:
     """Delegate each add of a model that also calls accumulate on the host."""
-    edge = to_edge(
-        export(model, inputs),
-        compile_config=EdgeCompileConfig(_check_ir_validity=False),
-    )
+    # torch._export.config turns off auto_functionalized v2 for export in some
+    # PyTorch builds (see is_fbcode() there). Under v1, a custom op that writes
+    # through a view writes a view_copy and the write is lost, a limitation
+    # that predates these passes. Pin v2 so these tests check the same graph in
+    # every build.
+    with torch._export.config.patch(enable_auto_functionalized_v2_for_export=True):
+        edge = to_edge(
+            export(model, inputs),
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        )
     et = edge.to_backend(PerAddDeviceDelegatePartitioner(devices)).to_executorch(
         et_config
     )

From 964a6111bafc21fce0c225d141b119fcf7e5e1d1 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 26 Oct 2026 13:23:44 -0700
Subject: [PATCH 26/26] Name the output in the warning for an output returned
 from the host

The warning for a returned output that an op writes on the host now says
which output it is and that a caller passing its own output buffers must
pass a host buffer for it. The docs and the config comment now give a
mutable custom op left on the host as the example, since a KV cache
update writes a buffer, which is a different case.
---
 docs/source/backends/cuda/cuda-overview.md |  6 +++---
 exir/passes/propagate_device_config.py     |  6 +++---
 exir/passes/propagate_device_pass.py       | 10 ++++++----
 exir/tests/test_propagate_device_pass.py   |  1 +
 4 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md
index 45feff2684b..e93186b6941 100644
--- a/docs/source/backends/cuda/cuda-overview.md
+++ b/docs/source/backends/cuda/cuda-overview.md
@@ -218,9 +218,9 @@ the graph in the delegate.
   gets a second copy, from the host.
 - A constant, such as a weight, that a delegate reads from outside is copied to each device that
   reads it.
-- A returned output that an operator outside the delegate writes in place, such as a KV cache
-  update, is copied to the host and returned there. The method metadata gives the device
-  of each output.
+- A returned output that an operator outside the delegate writes in place, such as a mutable
+  custom operator left on the host, is copied to the host and returned there. The method
+  metadata gives the device of each output.
 
 Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail
 instead of warning. It fails only on a copy that one of the skip flags you set was meant to remove.
diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py
index a6e87b9de40..49b1ffd20a1 100644
--- a/exir/passes/propagate_device_config.py
+++ b/exir/passes/propagate_device_config.py
@@ -68,9 +68,9 @@ class PropagateDeviceConfig:
 
     # When True, device delegate outputs that the method returns stay on
     # device, with one exception: when an op outside the delegate, such as a
-    # KV cache update, writes the host copy of such an output in place, the
-    # method returns that host copy, since the device tensor never sees the
-    # write. The copy to the host runs on every call, and the pass logs a
+    # mutable custom op left on the host, writes the host copy of such an
+    # output in place, the method returns that host copy, since the device
+    # tensor never sees the write. The copy to the host runs on every call, and the pass logs a
     # warning about it.
     # Useful for cross-method GPU pipelines where the next method consumes GPU
     # tensors directly.
diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py
index f9444cc620e..4e3ed5035bc 100644
--- a/exir/passes/propagate_device_pass.py
+++ b/exir/passes/propagate_device_pass.py
@@ -609,14 +609,16 @@ def _keep_method_outputs_on_device(
             writers = _in_place_writers(node)
             if writers:
                 device_spec = returned.meta["spec"]
+                index = list(graph_module.graph.output_node().args[0]).index(node)
                 self._report_per_call_copy(
-                    f"A method output ({device_spec.nbytes()} bytes) comes from a "
-                    f"delegate on {device_spec.device.name.lower()}:"
+                    f"Method output {index} ({device_spec.nbytes()} bytes) comes "
+                    f"from a delegate on {device_spec.device.name.lower()}:"
                     f"{device_spec.device_index} and '{writers[0].target}' writes "
                     f"it in place on the host, so it is copied to the host on every "
                     f"call and returned there, even with "
-                    f"skip_d2h_for_method_outputs. Have the partitioner take that "
-                    f"op into the delegate to avoid that copy.",
+                    f"skip_d2h_for_method_outputs. A caller that passes its own "
+                    f"output buffers must pass a host buffer for it. Have the "
+                    f"partitioner take that op into the delegate to avoid that copy.",
                     True,
                 )
                 continue
diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py
index 893a9d2cefa..612e8a06290 100644
--- a/exir/tests/test_propagate_device_pass.py
+++ b/exir/tests/test_propagate_device_pass.py
@@ -1019,6 +1019,7 @@ def forward(self, x, y):
                 )
             with self.subTest(write_view=write_view):
                 (warning,) = logs.output
+                self.assertIn("Method output 0 (16 bytes)", warning)
                 self.assertIn(
                     "'propagate_device_test.accumulate.default' writes it", warning
                 )