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.

Author: shoumikhin

exir/passes/propagate_device_pass.py

@@ -195,6 +195,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 for any other node."""
+        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)
@@ -223,17 +234,22 @@ 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.
+                if not all(
+                    self._delegate_device(graph_module, user)
+                    == (target_device_type, device_index)
+                    for user in arg.users
+                ):
                     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=True requires every user of "
+                        f"placeholder '{arg.name}' to be a delegate on the same "
+                        f"device, but it has {len(arg.users)} users. The placeholder "
+                        f"is shared by consumers that need it elsewhere, so its "
+                        f"TensorSpec cannot be safely mutated in-place to the "
+                        f"delegate's device. Either disable "
                         f"skip_h2d_for_method_inputs, or ensure the placeholder is "
-                        f"used exclusively by this delegate."
+                        f"used only by delegates on one device."
                     )
                 _set_device_on_spec(arg_spec, target_device_type, device_index)
                 changed = True

exir/tests/test_propagate_device_pass.py

@@ -1302,6 +1302,61 @@ def forward(self, a):
                     self.assertTrue(all(not spec.const for spec in copies))
                     self.assertTrue(all(spec.mem_id is not None for spec in copies))
 
+    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, as when a program calls one lowered module several times."""
+
+        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, cpu_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_shared_with_host_op(self):
+        """A method input that a host op also reads cannot move to the device."""
+
+        class Model(torch.nn.Module):
+            def forward(self, a, b):
+                return torch.add(a, b), torch.sin(a)
+
+        et_config = ExecutorchBackendConfig(
+            emit_stacktrace=False,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+        )
+        with self.assertRaisesRegex(Exception, "delegate on the same"):
+            self._get_executorch_program(
+                Model(),
+                (torch.randn(2, 2), torch.randn(2, 2)),
+                DeviceAwarePartitioner("cuda:0"),
+                et_config,
+            )
+
     def test_tensorspec_repr_includes_device(self):
         spec = TensorSpec(dtype=torch.float32, shape=torch.Size([2, 3]))
         repr_str = repr(spec)