Fix device placement of constants, shared inputs and delegate outputs

.github/workflows/cuda.yml

@@ -43,6 +43,10 @@ on:
       - examples/models/muse-glimmer/**
       - extension/pybindings/**
       - runtime/__init__.py
+      - exir/passes/convert_constant_dim_order_pass.py
+      - exir/passes/insert_write_back_for_buffers_pass.py
+      - exir/passes/propagate_device_config.py
+      - exir/passes/propagate_device_pass.py
   workflow_dispatch:
     inputs:
       models:
@@ -469,8 +473,9 @@ jobs:
   unittest-cuda:
     name: unittest-cuda (${{ matrix.shard }})
     needs: [changed-files, run-decision]
-    # This job runs the pybindings device test, so it also has to fire on the
-    # two pybindings paths the workflow filter lists.
+    # This job runs the pybindings device test and the device placement test,
+    # so it also has to fire on the pybindings and device pass paths the
+    # workflow filter lists.
     if: |
       contains(needs.changed-files.outputs.changed-files, 'backends/cuda') ||
       contains(needs.changed-files.outputs.changed-files, 'backends/aoti') ||
@@ -481,6 +486,10 @@ jobs:
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') ||
       contains(needs.changed-files.outputs.changed-files, 'extension/pybindings') ||
       contains(needs.changed-files.outputs.changed-files, 'runtime/__init__.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/convert_constant_dim_order_pass.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/insert_write_back_for_buffers_pass.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/propagate_device_config.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/passes/propagate_device_pass.py') ||
       needs.run-decision.outputs.is-full-run == 'true'
     # kernel: Triton kernels and their autotuning. system: lowering passes, op
     # dispatch, export/runtime, and model tests. Together they do not fit the
@@ -567,6 +576,21 @@ jobs:
             echo "::error::a pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
             exit 1
           fi
+
+          # Same for the device placement test, which skips without a GPU.
+          set +e
+          python -m pytest backends/cuda/tests/test_device_placement.py -v -rs -o "addopts=" \
+            2>&1 | tee /tmp/device_placement_test.log
+          device_placement_status=${PIPESTATUS[0]}
+          set -e
+          if [ "${device_placement_status}" -ne 0 ]; then
+            echo "::error::a device placement test failed"
+            exit "${device_placement_status}"
+          fi
+          if grep -qE "[0-9]+ skipped" /tmp/device_placement_test.log; then
+            echo "::error::a device placement test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
+            exit 1
+          fi
         fi
 
         # Build ExecuTorch with CUDA support
@@ -589,7 +613,7 @@ jobs:
           python -m pytest "${KERNEL_TESTS[@]}" -v -o "addopts=" --durations=20
         else
           IGNORES=()
-          for t in "${KERNEL_TESTS[@]}" "${QUANTIZED_GEMM_TESTS[@]}"; do
+          for t in "${KERNEL_TESTS[@]}" "${QUANTIZED_GEMM_TESTS[@]}" backends/cuda/tests/test_device_placement.py; do
             IGNORES+=("--ignore=${t}")
           done
           python -m pytest backends/cuda/tests backends/cuda/passes/tests -v -o "addopts=" --durations=20 \

backends/cuda/tests/test_device_placement.py

@@ -0,0 +1,147 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+from executorch.backends.cuda.cuda_backend import CudaBackend
+from executorch.exir import (
+    EdgeCompileConfig,
+    ExecutorchBackendConfig,
+    to_edge_transform_and_lower,
+)
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.passes import MemoryPlanningPass
+from executorch.exir.passes.propagate_device_pass import (
+    PropagateDeviceConfig,
+    TARGET_DEVICE_COMPILE_SPEC_KEY,
+)
+
+
+class EachAddOnCuda(Partitioner):
+    """Lowers each add to its own CUDA delegate. Parameters stay outside, so
+    they cross the delegate boundary as in a partial lowering."""
+
+    def partition(self, exported_program) -> PartitionResult:
+        spec = DelegationSpec(
+            CudaBackend.__name__,
+            [
+                CudaBackend.generate_method_name_compile_spec("forward"),
+                CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, b"cuda:0"),
+            ],
+        )
+        tags = {}
+        adds = [
+            node
+            for node in exported_program.graph_module.graph.nodes
+            if node.op == "call_function"
+            and node.target == exir_ops.edge.aten.add.Tensor
+        ]
+        for i, node in enumerate(adds):
+            node.meta["delegation_tag"] = f"add{i}"
+            tags[f"add{i}"] = spec
+        return PartitionResult(
+            tagged_exported_program=exported_program, partition_tags=tags
+        )
+
+
+def load(model, inputs, directory, config=None):
+    """Lower model with EachAddOnCuda and load its forward method from a file
+    in directory."""
+    from executorch.runtime import Runtime
+
+    edge = to_edge_transform_and_lower(
+        torch.export.export(model, inputs, strict=True),
+        partitioner=[EachAddOnCuda()],
+        compile_config=EdgeCompileConfig(
+            _check_ir_validity=False, _skip_dim_order=True
+        ),
+    )
+    path = Path(directory) / "model.pte"
+    path.write_bytes(edge.to_executorch(config or ExecutorchBackendConfig()).buffer)
+    return Runtime.get().load_program(path).load_method("forward")
+
+
+def run(model, inputs) -> list[torch.Tensor]:
+    """Lower model with EachAddOnCuda, run it in the runtime, return its outputs."""
+    with tempfile.TemporaryDirectory() as directory:
+        method = load(model, inputs, directory)
+        return [output.cpu() for output in method.execute(list(inputs))]
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+@unittest.skipIf(
+    torch.version.hip is not None,
+    "loads through the Python runtime, which the ROCm build does not include",
+)
+class CudaDevicePlacementTest(unittest.TestCase):
+    def test_transposed_parameter_read_by_a_delegate(self):
+        """The program stores a transposed parameter contiguous and copies
+        those bytes to the device. to_edge makes the parameter contiguous
+        before the delegate is compiled, so the delegate expects that layout."""
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                weight = torch.arange(6.0).reshape(3, 2).t()
+                self.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+            def forward(self, x):
+                return torch.add(x, self.weight)
+
+        model = Model()
+        x = torch.arange(6.0).reshape(2, 3) * 10
+        (output,) = run(model, (x,))
+        torch.testing.assert_close(output, model(x))
+
+    def test_constant_read_by_two_delegates(self):
+        """Two delegates read a constant stored in the program from outside
+        them, so the program copies it to the device for them."""
+
+        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):
+                return torch.add(x, self.table), torch.add(x * 2, self.table)
+
+        model = Model()
+        x = torch.arange(6.0).reshape(2, 3) * 10
+        for output, expected in zip(run(model, (x,)), model(x), strict=True):
+            torch.testing.assert_close(output, expected)
+
+    def test_input_on_device_also_read_on_host(self):
+        """With skip_h2d_for_method_inputs, the caller passes the input on the
+        device. A delegate reads it there, and a host op reads a copy of it in
+        host memory."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.add(x, y), torch.sin(x)
+
+        config = ExecutorchBackendConfig(
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            enable_non_cpu_memory_planning=True,
+            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
+        )
+        model = Model()
+        inputs = (torch.randn(2, 3), torch.randn(2, 3))
+        with tempfile.TemporaryDirectory() as directory:
+            method = load(model, inputs, directory, config)
+            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)

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

@@ -208,6 +208,29 @@ Do not pass a dict to `skip_h2d_for_method_inputs` itself. The pass reads that f
 truthiness, so any non-empty dict enables the skip for every method regardless of the values
 inside it.
 
+**Some copies stay even with the skip flags.** In the cases below, `to_executorch` keeps a copy
+that runs on every call and logs a warning. To remove the copy, have the partitioner put more of
+the graph in the delegate.
+
+- An input stays on the device of the first delegate that reads it. Other readers, such as an
+  operator outside the delegate, a delegate on another device, or the method output without
+  `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.
+- 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.
+
+**`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
+operator that uses it in the delegate. The CUDA partitioner already moves the buffer in, with
+`tag_mutated_buffer`, once those operators are there.
+
 **The choice is baked into the `.pte`.** A program exported with copies expects host tensors; one
 exported without them expects device tensors. Passing the wrong kind is a caller error, not
 something the runtime corrects, so the two are not interchangeable at run time.

exir/passes/BUCK

@@ -523,7 +523,9 @@ fbcode_target(_kind = runtime.python_library,
         "propagate_device_pass.py",
     ],
     deps = [
+        ":convert_constant_dim_order_pass",
         ":device_copy_ops_registry",
+        ":insert_write_back_for_buffers_pass",
         ":propagate_device_config",
         "//caffe2:torch",
         "//executorch/exir:delegate",

exir/passes/convert_constant_dim_order_pass.py

@@ -3,10 +3,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
+
 import torch
 from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 def _should_transform(tensor: torch.Tensor) -> bool:
     """
@@ -42,10 +46,20 @@ def _update_placeholder_meta(
         ),
         None,
     )
-    if input_spec is None:
+    if placeholder_node is None:
         raise RuntimeError(f"Missing placeholder for {input_spec.arg.name}.")
 
+    strides = placeholder_node.meta["val"].stride()
     placeholder_node.meta["val"] = placeholder_node.meta["val"].contiguous()
+    logger.warning(
+        "%s '%s' has strides %s, which are neither contiguous nor channels-last. "
+        "to_edge stores it contiguous, so partitioners and delegates see the "
+        "contiguous layout, not the strides in the model. Make it contiguous "
+        "where it is created to avoid this change.",
+        kind.name.replace("_", " ").capitalize(),
+        target,
+        tuple(strides),
+    )
 
 
 def convert_constant_dim_order_pass(
@@ -54,9 +68,9 @@ def convert_constant_dim_order_pass(
     """
     Normalize lifted tensors to default or channels-last dim order.
 
-    Tensor constants and buffers with other dim orders or striding are converted
-    to contiguous tensors. This pass acts in-place on the unlifted exported
-    program.
+    Tensor constants, buffers and parameters with other dim orders or striding
+    are converted to contiguous tensors. This pass acts in-place on the unlifted
+    exported program.
 
     Args:
         exported_program: The ExportedProgram to transform.
@@ -84,6 +98,13 @@ def convert_constant_dim_order_pass(
             )
             _update_placeholder_meta(exported_program, key, kind)
 
+    for key, param in exported_program.named_parameters():
+        if _should_transform(param):
+            exported_program.state_dict[key] = torch.nn.Parameter(
+                param.contiguous(), requires_grad=param.requires_grad
+            )
+            _update_placeholder_meta(exported_program, key, InputKind.PARAMETER)
+
     # Convert buffers.
     for key, buffer in exported_program.named_buffers():
         if (

exir/passes/propagate_device_config.py

@@ -48,28 +48,52 @@ 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.
+    # already on the target device: the device of the first delegate that
+    # reads the input, in graph order. Any other reader of such an input, such
+    # as a host op, a delegate on another device, or the method output without
+    # skip_d2h_for_method_outputs, reads a copy of it in host memory, made on
+    # every call, and the pass logs a warning. A delegate on another device
+    # then gets a second copy, from the host. An input that an op writes in
+    # place on the host makes to_executorch raise. 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 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.
     #
     # 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.
     # A dict can be used to set per-method values, keyed by method name.
     skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False
 
-    # When True, device delegate outputs that are directly method outputs
-    # are NOT wrapped with _d2h_copy. The method outputs stay on device.
-    # Useful for cross-method GPU pipelines where the next method consumes
-    # GPU tensors directly.
+    # When True, device delegate outputs that the method returns stay on
+    # device, with one exception: when an op outside the delegate, such as a
+    # mutable custom op left on the host, writes the host copy of such an
+    # output in place, the method returns that host copy, since the device
+    # tensor never sees the write. The copy to the host runs on every call, and the pass logs a
+    # warning about it.
+    # Useful for cross-method GPU pipelines where the next method consumes GPU
+    # tensors directly.
     #
     # Pair with MemoryPlanningPass(alloc_graph_output=False) for the same reason as the input flag.
     # A dict can be used to set per-method values, keyed by method name.
     skip_d2h_for_method_outputs: Union[bool, Dict[str, bool]] = False
 
+    # When True, the copies that the skip flags above 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_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
+
     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,
             )
         )
 
@@ -79,4 +103,5 @@ def __eq__(self, other: object) -> bool:
         return (
             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
         )