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 the reader is the
write-back of a mutated input or mutable buffer, the error says so, and
for a mutable buffer it points to tag_mutated_buffer, which lets the
partitioner keep the buffer inside the delegate.

Author: shoumikhin

exir/passes/propagate_device_pass.py

@@ -183,10 +183,10 @@ 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.
+      read only by delegates on one device, 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.
     - skip_d2h_for_method_outputs: If a delegate output feeds directly to
       graph output, don't insert D2H; the output stays on device.
     """
@@ -217,6 +217,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 _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)
@@ -243,17 +254,38 @@ 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_reader = next(
+                    (
+                        user
+                        for user in arg.users
+                        if self._delegate_device(graph_module, user)
+                        != (target_device_type, device_index)
+                    ),
+                    None,
+                )
+                if other_reader is not 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 '{other_reader.name}'"
+                    )
+                    if (
+                        other_reader.target == torch.ops.aten.copy_.default
+                        and other_reader.args[0] is arg
+                    ):
+                        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 delegate it, for example with "
+                            f"tag_mutated_buffer."
+                        )
                     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

exir/tests/test_propagate_device_pass.py

@@ -515,7 +515,8 @@ def forward(self, a, b, c):
     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")
@@ -528,26 +529,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]
@@ -560,18 +570,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):
@@ -1521,6 +1523,142 @@ 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)
+
+        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",
+            ),
+        }
+        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}'"
+                    ):
+                        edge.to_executorch(deepcopy(et_config))
+
+    def test_skip_h2d_refuses_placeholder_written_on_host(self):
+        """A mutated input or mutable buffer that only a device delegate reads is
+        still written back on the host, so the refusal says so."""
+
+        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 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"),
+            "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)