Add skip_h2d_for_constants to keep constants in device memory

backends/cuda/runtime/test/test_cuda_allocator.cpp

@@ -10,6 +10,8 @@
 
 #include <executorch/extension/cuda/runtime_api.h>
 
+#include <atomic>
+#include <chrono>
 #include <cstdint>
 #include <cstdlib>
 #include <limits>
@@ -178,6 +180,58 @@ TEST_F(CudaAllocatorTest, CopyRoundtripWithCallerStream) {
   ASSERT_EQ(cudaStreamDestroy(s), cudaSuccess);
 }
 
+// The compatibility header has no HIP names for non-blocking streams or host
+// callbacks.
+#if !defined(EXECUTORCH_USE_HIP)
+// A method can be loaded on one stream and run on another that is not ordered
+// after it, so synchronize() has to wait for the stream the copies used: the
+// caller stream when one is set, otherwise the default stream.
+TEST_F(CudaAllocatorTest, SynchronizeWaitsForTheCopyStream) {
+  const cudaHostFn_t hold = [](void* flag) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(200));
+    static_cast<std::atomic<bool>*>(flag)->store(true);
+  };
+  cudaStream_t caller;
+  ASSERT_EQ(
+      cudaStreamCreateWithFlags(&caller, cudaStreamNonBlocking), cudaSuccess);
+  for (cudaStream_t s : {caller, cudaStream_t{nullptr}}) {
+    std::optional<executorch::extension::cuda::CallerStreamGuard> g;
+    if (s != nullptr) {
+      g.emplace(s);
+    }
+    std::atomic<bool> done{false};
+    ASSERT_EQ(cudaLaunchHostFunc(s, hold, &done), cudaSuccess);
+    EXPECT_EQ(CudaAllocator::instance().synchronize(0), Error::Ok);
+    EXPECT_TRUE(done.load())
+        << (s != nullptr ? "caller" : "default") << " stream";
+    // If synchronize() did not wait, the callback must still finish before
+    // `done` goes out of scope.
+    ASSERT_EQ(cudaStreamSynchronize(s), cudaSuccess);
+  }
+  ASSERT_EQ(cudaStreamDestroy(caller), cudaSuccess);
+
+  // A caller stream of nullptr names the default stream of the current
+  // device, so with another device current, synchronize(0) still has to wait
+  // for device 0.
+  int device_count = 0;
+  ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess);
+  if (device_count < 2) {
+    return;
+  }
+  std::atomic<bool> done{false};
+  ASSERT_EQ(cudaSetDevice(0), cudaSuccess);
+  ASSERT_EQ(cudaLaunchHostFunc(nullptr, hold, &done), cudaSuccess);
+  ASSERT_EQ(cudaSetDevice(1), cudaSuccess);
+  {
+    executorch::extension::cuda::CallerStreamGuard g(nullptr);
+    EXPECT_EQ(CudaAllocator::instance().synchronize(0), Error::Ok);
+  }
+  EXPECT_TRUE(done.load()) << "default stream of device 0, device 1 current";
+  ASSERT_EQ(cudaSetDevice(0), cudaSuccess);
+  ASSERT_EQ(cudaStreamSynchronize(nullptr), cudaSuccess);
+}
+#endif // !EXECUTORCH_USE_HIP
+
 TEST_F(CudaAllocatorTest, CopyHostToDeviceNullDstReturnsInvalidArgument) {
   CudaAllocator& a = CudaAllocator::instance();
   // null dst should fail gracefully not CHECK abort

backends/cuda/tests/test_device_placement.py

@@ -145,3 +145,37 @@ def forward(self, x, y):
             outputs = method.execute([t.cuda() for t in inputs])
         for output, expected in zip(outputs, model(*inputs), strict=True):
             torch.testing.assert_close(output.cpu(), expected)
+
+    def test_constant_kept_in_device_memory(self):
+        """With skip_h2d_for_constants, a constant that only delegates read is
+        loaded into device memory once, when the method loads, and later
+        delegates must not reuse that memory, so every call reads it.
+        error_on_per_call_copies makes the export fail if the constant is
+        copied on every call instead."""
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("table", torch.arange(6.0).reshape(2, 3))
+
+            def forward(self, x):
+                a = torch.add(x, self.table)
+                b = torch.add(a, a)
+                return torch.add(b, b)
+
+        config = ExecutorchBackendConfig(
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True,
+                skip_h2d_for_constants=True,
+                error_on_per_call_copies=True,
+            ),
+            enable_non_cpu_memory_planning=True,
+            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
+        )
+        model = Model()
+        x = torch.arange(6.0).reshape(2, 3)
+        with tempfile.TemporaryDirectory() as directory:
+            method = load(model, (x,), directory, config)
+            for scale in (1.0, 10.0):
+                (output,) = method.execute([(x * scale).cuda()])
+                torch.testing.assert_close(output.cpu(), model(x * scale))

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

@@ -164,13 +164,14 @@ overwrite, unless you pass `clone_outputs=False`.
 
 ### Things worth knowing
 
-**The two flags are independent.** Skip only the input copies if you produce data on the GPU but
-want results back on the host, or only the output copies for the reverse.
+**`skip_h2d_for_method_inputs` and `skip_d2h_for_method_outputs` are independent.** Skip only the
+input copies if you produce data on the GPU but want results back on the host, or only the output
+copies for the reverse.
 
-**Both flags need `enable_non_cpu_memory_planning=True`.** It defaults to `True`, so leaving it
-out works, and the snippet above sets it explicitly only to make the requirement visible. Setting
-it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens
-during device-aware memory planning.
+**Every skip flag needs `enable_non_cpu_memory_planning=True`.** It defaults to `True`, so leaving
+it out works, and the snippet above sets it explicitly only to make the requirement visible.
+Setting it to `False` while asking for any skip raises a `ValueError`, because copy insertion
+happens during device-aware memory planning.
 
 **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the
 program reserves its own input buffer on the device, and the runtime refuses to copy the caller's
@@ -213,14 +214,27 @@ the graph in the delegate.
   `skip_d2h_for_method_outputs`, read a copy in host memory. A delegate on another device then
   gets a second copy, from the host.
 - A constant, such as a weight, that a delegate reads from outside is copied to each device that
-  reads it.
+  reads it. `skip_h2d_for_constants` (below) can remove this copy.
 - A returned output that an operator outside the delegate writes in place, such as a mutable
   custom operator left on the host, is copied to the host and returned there. The method
   metadata gives the device of each output.
 
 Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail
 instead of warning. It fails only on a copy that one of the skip flags you set was meant to remove.
 
+**Set `skip_h2d_for_constants=True` to keep constants in device memory.** If only delegates on
+one device read a constant, it is loaded into that device's memory once, when the method loads,
+and not copied on every call. The cost: it takes device memory for as long as the method is
+loaded, and each method that reads it has its own copy. Its data is stored as the initial state
+of that memory, so if the same data is also stored as an ordinary constant, for example for a
+method that reads it on the host, the `.pte` holds it twice. A constant that a host operator or a
+second device also reads is still copied on every call. So is a constant that is changed on the
+host before the delegate reads it, such as by a view, a transpose or a clone, and
+`error_on_per_call_copies` does not catch that copy. Have the partitioner take that change into
+the delegate, so that the delegate reads the constant itself. `share_mutable_buffers` needs every
+planned tensor in host memory, so it does not work with this flag. This flag works on its own: it
+needs neither `skip_h2d_for_method_inputs` nor unplanned inputs.
+
 **`skip_h2d_for_method_inputs` refuses inputs that the method changes.** `to_executorch` fails if
 the method changes an input, even when the delegate makes the change. This includes a mutable
 buffer, such as a KV cache, even though the caller does not pass it. For a buffer, put it and every

exir/emit/_emitter.py

@@ -2057,11 +2057,13 @@ def _is_buffer(node: Node, graph_signature: ExportGraphSignature) -> bool:
             # overall be considered a mutable buffer if it has already been memory planned. This would
             # suggest that the same abstract buffer is mutable in another entry point so we should
             # compel it to be considered mutable in all entry points at emission just as the user did with
-            # memory planning.
+            # memory planning. A device resident constant is planned too, but
+            # it is not mutable.
             is_mutable_buffer |= (
                 _is_buffer(self.node, self.exported_program.graph_signature)
                 and spec.mem_id is not None
                 and spec.mem_offset is not None
+                and not self.node.meta.get("device_resident_constant", False)
             )
 
             # If the placeholder has a constant_tag, it is external to the PTE file

exir/emit/test/test_emit.py

@@ -297,6 +297,26 @@ def forward(self, x):
             torch.ones(10, dtype=torch.int64),
         )
 
+        # An empty buffer has no initial data, so the program has no mutable
+        # segment, and loading the method must not look for one.
+        class EmptyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("cache_pos", torch.zeros(0, 2))
+
+            def forward(self, x):
+                self.cache_pos.add_(x)
+                return self.cache_pos + 1
+
+        x = torch.zeros(0, 2)
+        et_program_empty = to_edge(
+            torch.export.export(EmptyModule(), (x,), strict=True)
+        ).to_executorch(config=et_config)
+        method_empty = runtime.load_program(et_program_empty.buffer).load_method(
+            "forward"
+        )
+        self.assertEqual(method_empty.execute((x,))[0].shape, (0, 2))
+
     def test_int_list_input(self):
         class M(torch.nn.Module):
             def forward(self, x, y, z):

exir/memory_planning.py

@@ -351,8 +351,11 @@ def update_tensor_lifetime(
     else:
         start = node_idx if start is None or start > node_idx else start
 
-    if node.op == "placeholder" and _is_mutable_buffer(node, gs):
-        # mutable buffers are never freed
+    if node.op == "placeholder" and (
+        _is_mutable_buffer(node, gs) or _is_device_resident_constant(node)
+    ):
+        # Mutable buffers and device resident constants keep their data across
+        # calls, so they are never freed.
         end = max_node_idx
     else:
         end = node_idx if end is None or end < node_idx else end
@@ -386,6 +389,14 @@ def _is_mutable_buffer(
     return False
 
 
+def _is_device_resident_constant(node: Node) -> bool:
+    """Whether node is a constant that PropagateDevicePass placed in device
+    memory. The runtime copies its data from the program or its data file
+    into its planned memory once, at load, so it is planned like a mutable
+    buffer instead of being used in place."""
+    return node.op == "placeholder" and node.meta.get("device_resident_constant", False)
+
+
 def _get_mutable_buffer_specs(
     nodes: Iterable[Node], graph_signature: Optional[ExportGraphSignature]
 ) -> Set[TensorSpec]:
@@ -425,7 +436,11 @@ def get_graph_input_tensors(
 ) -> Set[TensorSpec]:
     graph_input_tensors = set()
     for node in nodes:
-        if node.op == "placeholder" and not _is_mutable_buffer(node, graph_signature):
+        if (
+            node.op == "placeholder"
+            and not _is_mutable_buffer(node, graph_signature)
+            and not _is_device_resident_constant(node)
+        ):
             for spec in get_node_tensor_specs(node):
                 graph_input_tensors.add(spec)
 
@@ -535,6 +550,7 @@ def collect_specs_from_nodes(  # noqa: C901
                 ignore_const
                 and spec.const
                 and not node.meta.get("weight_has_gradient", False)
+                and not _is_device_resident_constant(node)
             ):
                 continue
             if dedup:

exir/passes/propagate_device_config.py

@@ -25,7 +25,7 @@
 @compatibility(is_backward_compatible=False)
 @dataclass
 class PropagateDeviceConfig:
-    """Controls whether the runtime copies method inputs and outputs across the device boundary.
+    """Controls whether the runtime copies method inputs, outputs and constants across the device boundary.
 
     Memory planning reserves buffers for graph inputs and outputs by default. A program that skips
     the input copy must use ``alloc_graph_input=False``. Unplanning the outputs as well, as below,
@@ -46,6 +46,8 @@ class PropagateDeviceConfig:
     already owns and refuses to copy into it, so every call fails with ``NotSupported``. Planned
     outputs work, but they live in the method's memory, which the next run overwrites, and the
     Python ``execute()`` copies them unless it gets ``clone_outputs=False``.
+
+    ``skip_h2d_for_constants`` does not need ``alloc_graph_input=False``: it works with planned inputs.
     """
 
     # When True, method-level input tensors that feed directly into a device
@@ -61,7 +63,8 @@ class PropagateDeviceConfig:
     # by the caller, so they are still copied: their data stays on the host,
     # in the program or in its external data file. That copy runs on every
     # call, once for each device that reads the constant, and stays in device
-    # memory until the last delegate there has read it.
+    # memory until the last delegate there has read it. Set
+    # skip_h2d_for_constants to keep constants on the device instead.
     #
     # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer
     # for the input and refuses the caller's tensor.
@@ -85,21 +88,37 @@ class PropagateDeviceConfig:
     # propagate_device_config of ExecutorchBackendConfig.
     skip_d2h_for_method_outputs: Union[bool, Dict[str, bool]] = False
 
-    # When True, the copies that the skip flags above cannot remove make
+    # When True, the copies that the skip flags cannot remove make
     # to_executorch raise instead of logging a warning: with
-    # skip_h2d_for_method_inputs, a constant copied to the device and a method
-    # input copied to the host for a reader off its device, and with
+    # skip_h2d_for_method_inputs or skip_h2d_for_constants, a constant copied
+    # to the device; with skip_h2d_for_method_inputs, a method input copied to
+    # the host for a reader off its device; and with
     # skip_d2h_for_method_outputs, a returned output copied to the host for a
     # host op that writes it. Each runs on every call. Use it to make sure a
     # program that skips its boundary copies has none left.
     error_on_per_call_copies: bool = False
 
+    # When True, a constant that only delegates on one device read is kept in
+    # that device's memory instead of being copied there on every call. The
+    # runtime loads its data there once, when it loads the method, through
+    # host memory. The constant then uses device memory for as long as the
+    # method is loaded, and each method that reads it keeps its own copy. A
+    # constant that anything else reads is still copied on every call. Only a
+    # constant that a delegate reads directly is kept: one changed on the host
+    # first, such as by a view, a transpose or a clone, is copied on every
+    # call, and error_on_per_call_copies does not catch that copy. To keep
+    # it, have the partitioner take that change into the delegate too.
+    # MemoryPlanningPass(share_mutable_buffers=True) needs every planned tensor
+    # in host memory, so it does not work with this option.
+    skip_h2d_for_constants: bool = False
+
     def __hash__(self) -> int:
         return hash(
             (
                 str(self.skip_h2d_for_method_inputs),
                 str(self.skip_d2h_for_method_outputs),
                 self.error_on_per_call_copies,
+                self.skip_h2d_for_constants,
             )
         )
 
@@ -110,4 +129,5 @@ def __eq__(self, other: object) -> bool:
             self.skip_h2d_for_method_inputs == other.skip_h2d_for_method_inputs
             and self.skip_d2h_for_method_outputs == other.skip_d2h_for_method_outputs
             and self.error_on_per_call_copies == other.error_on_per_call_copies
+            and self.skip_h2d_for_constants == other.skip_h2d_for_constants
         )