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.

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.

Author: shoumikhin

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

@@ -156,6 +156,11 @@ exec_program = et_program.to_executorch(
 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.
+
 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/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",
         ":propagate_device_config",
         "//caffe2:torch",

exir/passes/propagate_device_config.py

@@ -49,7 +49,9 @@ 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 in the program on the host. 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.

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
 
 # Re-exported for backward compatibility; the dataclass lives in a lightweight
 # module so that ExecutorchBackendConfig can reference it without importing the
@@ -135,14 +136,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
 
@@ -179,7 +185,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.
     """
@@ -227,7 +235,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:
@@ -244,14 +256,27 @@ 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."
+                )
+
             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:
@@ -308,7 +333,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:

exir/tests/test_propagate_device_pass.py

@@ -30,6 +30,7 @@
 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.pass_base import ExportedProgramPassBase, ExportedProgramPassResult
 from executorch.exir.passes import MemoryPlanningPass
 from executorch.exir.passes.propagate_device_pass import (
     _get_target_device_from_compile_specs,
@@ -1478,6 +1479,134 @@ 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,
+            )
+            for pipeline, program, gm in self._get_executorch_program(
+                model, inputs, DeviceAwarePartitioner("cuda:0"), et_config
+            ):
+                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, for a transposed parameter as well as
+        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)