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.

For example, a constant read by delegates on cuda:0, cuda:1 and cuda:0 now
gets two copies instead of three.

Author: shoumikhin

docs/source/backends/cuda/cuda-overview.md

@@ -158,8 +158,8 @@ outputs = method([torch.randn(4, 64, device="cuda")])
 
 Constants that the partitioner leaves outside the delegate, such as weights or buffers the method
 does not change, are not passed by the caller. Their data stays in the program in host memory, so
-they are still copied to the device on every call, once for each delegate that reads them. Keep
-constants inside the delegate to avoid that cost.
+they are still copied to the device on every call. Keep constants inside the delegate to avoid
+that cost.
 
 Outputs also stay on the device, which is what you want when the next stage of your pipeline is
 also on the GPU.

exir/passes/propagate_device_config.py

@@ -51,7 +51,7 @@ class PropagateDeviceConfig:
     # 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 in the program on the host. That copy runs
-    # on every call, once for each delegate that reads the constant.
+    # on every call, once for each device 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.

exir/passes/propagate_device_pass.py

@@ -175,7 +175,8 @@ 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.
+    - 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.
       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,
       or after the call itself when it returns a single tensor.
@@ -304,18 +305,32 @@ def _insert_h2d_copies(
                     f"it is created, for example in the pass that adds it."
                 )
 
-            with graph_module.graph.inserting_before(node):
-                h2d_node = graph_module.graph.call_function(
-                    torch.ops.et_copy._h2d_copy.default,
-                    (arg,),
-                )
-                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
+            # Delegates are visited in graph order, so a copy made for an earlier
+            # delegate on this device already comes before this one.
+            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,
+            )
+            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"]
+            new_args[i] = h2d_node
+            changed = True
 
         node.args = tuple(new_args)
         return changed

exir/tests/test_propagate_device_pass.py

@@ -580,9 +580,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)
                 )
@@ -779,14 +779,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(
@@ -1609,6 +1609,55 @@ 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():
+            et = to_edge(export(Model(devices), inputs[:1])).to_executorch(
+                deepcopy(et_config)
+            )
+            gm = et.exported_program().graph_module
+            with self.subTest(case=case):
+                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_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