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.

Author: shoumikhin

exir/passes/propagate_device_pass.py

@@ -143,6 +143,8 @@ def _clone_spec_with_device(
     """Create a copy of a TensorSpec with a different device."""
     new_spec = copy.copy(spec)
     new_spec.init_mem_planning_fields()
+    # A copy's output is planned memory, even when it copies a constant.
+    new_spec.const = False
     _set_device_on_spec(new_spec, device_type, device_index)
     return new_spec
 
@@ -214,7 +216,13 @@ def _insert_h2d_copies(
             if not isinstance(arg_spec, TensorSpec):
                 continue
 
-            if self.skip_h2d_for_method_inputs and self._is_placeholder(arg):
+            # A constant is a placeholder too, but its data is in the program on
+            # the host, so it still needs a copy like any other host tensor.
+            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:

exir/tests/test_propagate_device_pass.py

@@ -1265,6 +1265,43 @@ 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, cpu_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.assertTrue(all(spec.mem_id is not None for spec in copies))
+
     def test_tensorspec_repr_includes_device(self):
         spec = TensorSpec(dtype=torch.float32, shape=torch.Size([2, 3]))
         repr_str = repr(spec)