From 3b90a492044ded5f7683c907f50902bded21dc67 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 09:55:46 -0700 Subject: [PATCH] Fix device placement of constants, shared inputs and delegate outputs (#24088) ## What is wrong today When a delegate runs on a GPU, `PropagateDevicePass` copies its inputs to the device (`_h2d_copy`) and its outputs back (`_d2h_copy`). `skip_h2d_for_method_inputs` and `skip_d2h_for_method_outputs` remove those copies at the method's edges. 1. A transposed parameter is stored contiguous, but a delegate reading it is compiled for its original strides. `to_edge` does this silently. 2. When a host op writes a delegate output in place, the next delegate on that device reads the old value. Such a `relu_` makes memory planning fail. 3. A lowered module returning one tensor, called from another module, gets no `_d2h_copy`. 4. With the output skip flag, a returned output that something else reads comes back on the host. 5. A constant that feeds a device delegate makes `to_executorch` fail. With the input skip flag, it arrives as host memory marked as device memory. 6. With the input skip flag, an input with more than one reader is refused. 7. A delegate targeting the CPU gets copies that fail at run time. ## What this change does Numbered as above: 1. `to_edge` makes such a parameter contiguous before partitioning, and warns for each tensor whose layout it changes. 2. The host op writes the host copy of the output. The copies stay when it writes between two delegates, also through a view. 3. The output gets its `_d2h_copy`. 4. Returned outputs stay on the device. One that a host op writes in place is returned on the host, with a warning. 5. The constant gets a device copy, with a warning. One neither contiguous nor channels-last is refused. 6. The input stays on the device of its first delegate. Shape-only ops read it there. Other readers share one `_d2h_copy` of it, with a warning. An input written in place on the host is refused. 7. The delegate gets no copies. Also: - A tensor that several delegates on one device read is copied there once, unless written in between. - New option `PropagateDeviceConfig(error_on_per_call_copies=True)` turns the warnings in 4 to 6 into errors when the matching skip flag is set. ## What was tested - Device pass tests and new `exir/tests/test_passes.py` tests pass. All new tests but one fail without their fix. - `backends/cuda/tests/test_device_placement.py` matches eager PyTorch with real CUDA delegates on an H100. The `unittest-cuda` job now runs when these passes change, and fails if that test skips. (cherry picked from commit cc4b514d2577e22ca9ca20bad60e5a2f80c4a4f7) --- .github/workflows/cuda.yml | 30 +- backends/cuda/tests/test_device_placement.py | 147 ++ docs/source/backends/cuda/cuda-overview.md | 23 + exir/passes/BUCK | 2 + .../passes/convert_constant_dim_order_pass.py | 29 +- exir/passes/propagate_device_config.py | 37 +- exir/passes/propagate_device_pass.py | 480 ++++++- exir/program/_program.py | 12 +- exir/tests/test_passes.py | 109 ++ exir/tests/test_propagate_device_pass.py | 1183 ++++++++++++++++- 10 files changed, 1955 insertions(+), 97 deletions(-) create mode 100644 backends/cuda/tests/test_device_placement.py diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 41dd61ebcda..07653ae8bd4 100644 --- a/.github/workflows/cuda.yml +++ b/.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 \ diff --git a/backends/cuda/tests/test_device_placement.py b/backends/cuda/tests/test_device_placement.py new file mode 100644 index 00000000000..b54242802a1 --- /dev/null +++ b/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) diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index db7d0adfb75..32effc7e520 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -204,6 +204,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. diff --git a/exir/passes/BUCK b/exir/passes/BUCK index 8338a58725e..6eeed352f34 100644 --- a/exir/passes/BUCK +++ b/exir/passes/BUCK @@ -524,7 +524,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", diff --git a/exir/passes/convert_constant_dim_order_pass.py b/exir/passes/convert_constant_dim_order_pass.py index dd98fc9ece6..36037961f62 100644 --- a/exir/passes/convert_constant_dim_order_pass.py +++ b/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 ( diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py index 119231f563a..3d6dfeeefda 100644 --- a/exir/passes/propagate_device_config.py +++ b/exir/passes/propagate_device_config.py @@ -50,8 +50,18 @@ 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 refuses the caller's tensor. @@ -59,10 +69,14 @@ class PropagateDeviceConfig: # propagate_device_config of ExecutorchBackendConfig. 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. # # Planned outputs work, but they live in the method's memory, which the next run overwrites, # and the Python runtime copies them. MemoryPlanningPass(alloc_graph_output=False) avoids @@ -71,11 +85,21 @@ 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 + # 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, ) ) @@ -85,4 +109,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 ) diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py index aac2de75593..4e3ed5035bc 100644 --- a/exir/passes/propagate_device_pass.py +++ b/exir/passes/propagate_device_pass.py @@ -18,6 +18,13 @@ 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 +from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + _alias_sets, + _contains_node, + _schema_paired_args, + _schema_writes_input, +) # Re-exported for backward compatibility; the dataclass lives in a lightweight # module so that ExecutorchBackendConfig can reference it without importing the @@ -35,6 +42,16 @@ # with this key and a value encoding the device string (e.g., b"cuda:0"). TARGET_DEVICE_COMPILE_SPEC_KEY = "target_device" +# Ops that read only the sizes, strides or offset of a tensor. Those stay in +# host memory when the data is on a device, so these ops need no copy of it. +_METADATA_READERS = ( + torch.ops.aten.sym_size.default, + torch.ops.aten.sym_size.int, + torch.ops.aten.sym_stride.int, + torch.ops.aten.sym_numel.default, + torch.ops.aten.sym_storage_offset.default, +) + def _parse_device_spec_value(value: bytes) -> tuple[schema.DeviceType, int]: """ @@ -135,36 +152,166 @@ 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 +def _schema_aliases_input( + node: torch.fx.Node, schema: torch.FunctionSchema, input_node: torch.fx.Node +) -> bool: + """Whether schema marks the result of node as an alias of input_node, as for + a view or an in-place op.""" + returned = set().union(*(_alias_sets(ret.alias_info) for ret in schema.returns)) + return any( + _alias_sets(schema_arg.alias_info) & returned + for arg, schema_arg in _schema_paired_args(node, schema) + if _contains_node(arg, input_node) + ) + + +def _in_place_writers(node: torch.fx.Node) -> list[torch.fx.Node]: + """The ops that write *node* in place, directly or through a view of it, + per the op's schema. + + A mutable op left on the host does this once to_executorch has removed + auto_functionalized, so a copy of *node* can miss the write. If the op + wrote a view of *node* in the model, that removal makes the view again, + for example with as_strided. The result of an in-place op is a view of + its input too. A view_copy is not: it still makes its own tensor here, + and ReplaceViewCopyWithViewPass turns it into a view later only where no + write can tell the difference. Custom op schemas are trusted, as + functionalization trusts them. The write-back pass trusts only aten + schemas, but here that would keep a copy for every custom op on the host, + such as a KV cache update. Users without a schema, such as delegate calls + and getitem, do not write their inputs. + """ + writers: dict[torch.fx.Node, None] = {} + aliases = {node} + pending = [node] + while pending: + alias = pending.pop() + for user in alias.users: + schema = getattr(user.target, "_schema", None) + if schema is None: + continue + if _schema_writes_input(user, schema, alias): + writers[user] = None + if user not in aliases and _schema_aliases_input(user, schema, alias): + aliases.add(user) + pending.append(user) + return list(writers) + + +def _written_between( + node: torch.fx.Node, + after: Optional[torch.fx.Node], + before: torch.fx.Node, +) -> bool: + """Whether an op writes *node* in place after *after*, or anywhere when it + is None, and before *before*, in graph order, which is the order the + method runs in. A write outside that range cannot change what a reader at + *before* of a copy made at *after* sees.""" + writers = _in_place_writers(node) + if not writers: + return False + position = {n: i for i, n in enumerate(before.graph.nodes)} + start = position[after] if after is not None else -1 + return any(start < position[writer] < position[before] for writer in writers) + + +def _reusable_copy( + tensor: torch.fx.Node, + delegate: torch.fx.Node, + device_type: schema.DeviceType, + device_index: int, +) -> Optional[torch.fx.Node]: + """The _h2d_copy of *tensor* to this device that *delegate* can read too. + + Delegates are visited in graph order, so a copy made for an earlier + delegate on this device already comes before *delegate*, and the last such + copy is the newest. It is stale when an op writes *tensor* in place between + that copy and *delegate*, so then there is none. + """ + copy = None + for user in tensor.users: + if ( + user.target == torch.ops.et_copy._h2d_copy.default + and user.meta["spec"].device == device_type + and user.meta["spec"].device_index == device_index + ): + copy = user + if copy is not None and _written_between(tensor, copy, delegate): + return None + return copy + + +def _replace_spec( + readers: list[torch.fx.Node], old_spec: TensorSpec, new_spec: TensorSpec +) -> None: + """Replace *old_spec* with *new_spec* in *readers*, which now read the + tensor with *new_spec*, and in the readers of their results. + + SpecPropPass gives the output node the spec of each tensor it returns, and + an in-place result the spec of the tensor it writes. Memory planning relies + on both, so these specs have to follow the tensor that is read now. + """ + pending = list(readers) + while pending: + node = pending.pop() + spec = node.meta.get("spec") + if any(value is old_spec for value in pytree.tree_leaves(spec)): + node.meta["spec"] = pytree.tree_map( + lambda value: new_spec if value is old_spec else value, spec + ) + pending.extend(node.users) + + class PropagateDevicePass(PassBase): """ After to_backend, walk the graph and insert H2D/D2H copy ops at delegate 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. - 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. - The getitem stays device; the d2h_copy output is tagged as CPU. + When a delegate has a target_device CompileSpec (e.g., "cuda:0") other + than the CPU: + - 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 when no op + writes the input in place in between. A reused copy stays in device + memory until its last reader. + The original input node keeps its device; 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. + The delegate output stays device; the d2h_copy output is tagged as CPU. - Getitem nodes that extract from a delegate call inherit the device. 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). - - skip_d2h_for_method_outputs: If the getitem feeds directly to graph - output, don't insert D2H — the output stays on device. + feeding a delegate, don't insert H2D; tag the placeholder as device + instead (user provides GPU tensor at runtime). Every other reader of the + placeholder reads one D2H copy of it. A placeholder that an op writes in + place on the host is refused. 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. When + one of them writes the copy in place, the graph output reads the written + host copy instead. + - error_on_per_call_copies: Raise instead of logging a warning when a + copy that a skip flag asked to avoid has to stay: a constant copied to a + device, a method input copied for a reader off its device, or a + returned output written on the host. """ def __init__( @@ -172,11 +319,13 @@ def __init__( skip_h2d_for_method_inputs: bool = False, skip_d2h_for_method_outputs: bool = False, enable_non_cpu_memory_planning: bool = False, + error_on_per_call_copies: bool = False, ) -> None: super().__init__() self.skip_h2d_for_method_inputs = skip_h2d_for_method_inputs self.skip_d2h_for_method_outputs = skip_d2h_for_method_outputs self.enable_non_cpu_memory_planning = enable_non_cpu_memory_planning + self.error_on_per_call_copies = error_on_per_call_copies if ( skip_h2d_for_method_inputs or skip_d2h_for_method_outputs @@ -193,9 +342,25 @@ def _is_placeholder(self, node: torch.fx.Node) -> bool: """Check if a node is a graph-level input (placeholder).""" return node.op == "placeholder" - def _feeds_directly_to_output(self, node: torch.fx.Node) -> bool: - """Check if all users of a node are output nodes.""" - return all(user.op == "output" for user in node.users) + def _delegate_device( + self, graph_module: torch.fx.GraphModule, node: torch.fx.Node + ) -> Optional[tuple[schema.DeviceType, int]]: + """The target device of a delegate call node, or None when there is none.""" + if node.op != "call_function" or node.target != executorch_call_delegate: + return None + lowered_module = _get_lowered_module(graph_module, node) + if lowered_module is None: + return None + return _get_target_device_from_compile_specs(lowered_module) + + def _report_per_call_copy(self, message: str, skip_requested: bool) -> None: + """Warn about a copy that runs on every call. When a skip flag asked + for no such copy and error_on_per_call_copies is set, raise instead.""" + if skip_requested and self.error_on_per_call_copies: + raise RuntimeError( + f"{message} Raised because error_on_per_call_copies=True." + ) + logger.warning("%s", message) def _insert_h2d_copies( self, @@ -204,7 +369,13 @@ def _insert_h2d_copies( target_device_type: schema.DeviceType, device_index: int, ) -> bool: - """Insert H2D copy nodes for each tensor input to a delegate call.""" + """Make each tensor input of a delegate call read device memory. + + A method input that is not a constant moves to the device under + skip_h2d_for_method_inputs. + Any other input reads an H2D copy, reused from an earlier delegate on + the same device when possible. + """ changed = False new_args = list(node.args) for i, arg in enumerate(node.args[1:], start=1): @@ -214,41 +385,149 @@ def _insert_h2d_copies( if not isinstance(arg_spec, TensorSpec): continue - if self.skip_h2d_for_method_inputs and self._is_placeholder(arg): - # TODO(gasoonjia): support skip_h2d_for_method_inputs for - # multiple-user placeholder inputs. - if len(arg.users) != 1: - raise RuntimeError( - f"skip_h2d_for_method_inputs=True requires placeholder " - f"'{arg.name}' to have exactly one user, but it has " - f"{len(arg.users)} users. The placeholder is shared by " - f"multiple consumers, so its TensorSpec cannot be safely " - f"mutated in-place to the delegate's device. Either disable " - f"skip_h2d_for_method_inputs, or ensure the placeholder is " - f"used exclusively by this delegate." - ) - _set_device_on_spec(arg_spec, target_device_type, device_index) + if ( + self.skip_h2d_for_method_inputs + and self._is_placeholder(arg) + and not arg_spec.const + ): + self._place_input_on_device( + graph_module, arg, target_device_type, device_index + ) changed = True continue - 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 + 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." ) - 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 + + h2d_node = _reusable_copy(arg, node, target_device_type, device_index) + if h2d_node is None: + if arg_spec.const: + self._report_per_call_copy( + f"Constant '{arg.name}' ({arg_spec.nbytes()} bytes) feeds a " + f"delegate on {target_device_type.name.lower()}:" + f"{device_index} from outside it, so it is copied to the " + f"device on every call. Have the partitioner take it into " + f"the delegate to avoid that copy.", + self.skip_h2d_for_method_inputs, + ) + 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 + def _place_input_on_device( + self, + graph_module: torch.fx.GraphModule, + placeholder: torch.fx.Node, + device_type: schema.DeviceType, + device_index: int, + ) -> None: + """Move a method input to the device for skip_h2d_for_method_inputs. + + Delegates on that device read it there, and so do ops that read only + its shape. Every other reader, such as a host op, a delegate on the + host, a delegate on another device or, without + skip_d2h_for_method_outputs, the method output, reads one D2H copy of it + instead, made on every call. An + input that an op writes in place on the host is refused: the write would + go to the copy, and neither the caller nor the delegates would see it. + Delegates are visited in graph order, so the first one that reads the + input decides its device, and later ones on that device find it placed. + """ + spec = placeholder.meta["spec"] + if (spec.device, spec.device_index) == (device_type, device_index): + return + device = f"{device_type.name.lower()}:{device_index}" + writers = _in_place_writers(placeholder) + if writers: + write_back = next( + ( + writer + for writer in writers + if writer.target == torch.ops.aten.copy_.default + and writer.args[0] is placeholder + ), + None, + ) + refusal = ( + f"skip_h2d_for_method_inputs=True cannot place placeholder " + f"'{placeholder.name}' on {device}, because " + f"'{(write_back or writers[0]).name}' writes it" + ) + if write_back is not None: + raise RuntimeError( + f"{refusal} back on the host (the method mutates it). Either " + f"disable skip_h2d_for_method_inputs, or, for a mutable buffer, " + f"have the partitioner put the buffer and every op that reads " + f"or writes it in the delegate; tag_mutated_buffer tags the " + f"buffer once those ops are there." + ) + raise RuntimeError( + f"{refusal} in place on the host. Either disable " + f"skip_h2d_for_method_inputs, or have the partitioner put that op " + f"in a delegate on {device}." + ) + + _set_device_on_spec(spec, device_type, device_index) + # With skip_d2h_for_method_outputs, the method returns the input where + # the caller put it. + other_readers = [ + user + for user in placeholder.users + if self._delegate_device(graph_module, user) != (device_type, device_index) + and user.target not in _METADATA_READERS + and not (user.op == "output" and self.skip_d2h_for_method_outputs) + ] + if not other_readers: + return + position = {n: i for i, n in enumerate(graph_module.graph.nodes)} + first_reader = min(other_readers, key=position.__getitem__) + self._report_per_call_copy( + f"Method input '{placeholder.name}' ({spec.nbytes()} bytes) stays on " + f"{device} for the delegates there, but '{first_reader.name}' also " + f"reads it, so it is copied to the host on every call, and from there " + f"to any other device that reads it. To avoid that, have the " + f"partitioner put every op that reads it in a delegate on {device}, " + f"and set skip_d2h_for_method_outputs if the method returns it.", + True, + ) + with graph_module.graph.inserting_before(first_reader): + d2h_node = graph_module.graph.call_function( + torch.ops.et_copy._d2h_copy.default, (placeholder,) + ) + d2h_spec = _copy_output_spec(spec, schema.DeviceType.CPU, 0) + d2h_node.meta["spec"] = d2h_spec + d2h_node.meta["val"] = placeholder.meta.get("val") + if "tensor_meta" in placeholder.meta: + d2h_node.meta["tensor_meta"] = placeholder.meta["tensor_meta"] + for reader in other_readers: + reader.replace_input_with(placeholder, d2h_node) + _replace_spec(other_readers, spec, d2h_spec) + def _insert_d2h_for_getitem( self, graph_module: torch.fx.GraphModule, @@ -280,16 +559,22 @@ def _insert_d2h_for_getitem( return False _set_device_on_spec(spec, source_spec.device, source_spec.device_index) + self._insert_d2h_after(graph_module, node, spec) + return True - if self.skip_d2h_for_method_outputs and self._feeds_directly_to_output(node): - return True - + def _insert_d2h_after( + self, + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + spec: TensorSpec, + ) -> None: + """Insert a D2H copy after *node*, a device tensor, for all of its users.""" with graph_module.graph.inserting_after(node): d2h_node = graph_module.graph.call_function( 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: @@ -299,15 +584,54 @@ def _insert_d2h_for_getitem( d2h_node, delete_user_cb=lambda user, _d2h=d2h_node: user != _d2h, ) - # Memory planning finds the graph outputs through the output node's - # spec, so it has to name the copy now rather than the device tensor. - for user in d2h_node.users: - if user.op == "output" and "spec" in user.meta: - user.meta["spec"] = pytree.tree_map( - lambda value: d2h_spec if value is spec else value, - user.meta["spec"], - ) - return True + _replace_spec(list(d2h_node.users), spec, d2h_spec) + + def _keep_method_outputs_on_device( + self, + graph_module: torch.fx.GraphModule, + ) -> None: + """Make the graph output read the device tensor instead of its D2H copy. + + The copy stays for its other users, such as a host op. When one of them + writes the copy in place, the device tensor never sees that write, so the + output keeps reading the written host copy, and a warning says so. + Copying it back to the device would add a second copy on every call. + """ + d2h = torch.ops.et_copy._d2h_copy.default + for node in list(graph_module.graph.nodes): + if ( + node.op != "call_function" + or node.target != d2h + or not any(user.op == "output" for user in node.users) + ): + continue + returned = node.args[0] + writers = _in_place_writers(node) + if writers: + device_spec = returned.meta["spec"] + index = list(graph_module.graph.output_node().args[0]).index(node) + self._report_per_call_copy( + f"Method output {index} ({device_spec.nbytes()} bytes) comes " + f"from a delegate on {device_spec.device.name.lower()}:" + f"{device_spec.device_index} and '{writers[0].target}' writes " + f"it in place on the host, so it is copied to the host on every " + f"call and returned there, even with " + f"skip_d2h_for_method_outputs. A caller that passes its own " + f"output buffers must pass a host buffer for it. Have the " + f"partitioner take that op into the delegate to avoid that copy.", + True, + ) + continue + _replace_spec( + [graph_module.graph.output_node()], + node.meta["spec"], + returned.meta["spec"], + ) + node.replace_all_uses_with( + returned, delete_user_cb=lambda user: user.op == "output" + ) + if not node.users: + graph_module.graph.erase_node(node) def _elide_redundant_same_device_copies( self, @@ -327,8 +651,10 @@ def _elide_redundant_same_device_copies( rewires the consumer straight to that device tensor and drops both copies. The device comparison is the correctness guard: a genuine cross-device move (e.g. ``cuda:0 -> host -> cuda:1``) has differing devices and is left - intact. A D2H feeding other (host) consumers keeps those users and is - removed only once it becomes dead. + intact. So is a round-trip whose host copy an op writes in place before + the copy back, because the device tensor never sees that write. A D2H + feeding other (host) consumers keeps those users and is removed only once + it becomes dead. """ h2d = torch.ops.et_copy._h2d_copy.default d2h = torch.ops.et_copy._d2h_copy.default @@ -356,6 +682,7 @@ def _elide_redundant_same_device_copies( if ( x_spec.device != h2d_spec.device or x_spec.device_index != h2d_spec.device_index + or _written_between(src, None, node) ): continue @@ -371,15 +698,20 @@ def _elide_redundant_same_device_copies( def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 # Three-pass approach: # Pass 1 – For each delegate with a target_device CompileSpec, insert - # H2D copy nodes before delegate inputs and tag the delegate + # H2D copy nodes before delegate inputs, or reuse the copy an + # earlier delegate on the same device reads, and tag the delegate # output specs with the target device. Delegates without a - # target_device are left untouched (no copies, specs stay CPU). + # target_device, or with the CPU as their target_device, are + # left untouched (no copies, specs stay CPU). # Pass 2 – For each getitem that extracts from a device-tagged delegate # (tracked in device_delegates), propagate the device onto the # getitem spec and insert a D2H copy after it so downstream - # non-delegated ops receive CPU tensors. + # non-delegated ops receive CPU tensors. A delegate call that + # returns a single tensor gets the D2H copy itself. # Pass 3 – Elide the redundant device->host->device round-trips left # between consecutive same-device delegates by passes 1 and 2. + # With skip_d2h_for_method_outputs, then make the graph output + # read device tensors instead of D2H copies. changed = False device_delegates: set[torch.fx.Node] = set() @@ -395,7 +727,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 ) result = _get_target_device_from_compile_specs(lowered_module) - if result is None: + # A delegate on the host reads and writes host memory directly. + if result is None or result[0] == schema.DeviceType.CPU: continue target_device_type, device_index = result @@ -428,10 +761,20 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 lowered_module.backend_id, ) - # Second pass: propagate device through getitem nodes and insert D2H - # only for delegates that have a target_device. + # Second pass: propagate device through getitem nodes and insert D2H, + # after each getitem or single-tensor call, only for delegates that have + # a target_device. for node in list(graph_module.graph.nodes): - if node.op == "call_function" and node.target == operator.getitem: + # A lowered module that returns one tensor and is called from another + # module has no getitem after its call, so copy its output directly. + if ( + node in device_delegates + and self.enable_non_cpu_memory_planning + and isinstance(node.meta.get("spec"), TensorSpec) + ): + self._insert_d2h_after(graph_module, node, node.meta["spec"]) + changed = True + elif node.op == "call_function" and node.target == operator.getitem: source = node.args[0] if isinstance(source, torch.fx.Node) and source in device_delegates: if self.enable_non_cpu_memory_planning: @@ -457,11 +800,14 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 # Third pass: elide the redundant device->host->device round-trips left # between consecutive same-device delegates by the per-boundary H2D/D2H - # insertion above. Only meaningful when passes 1/2 modified the graph - # (there was a device delegate); when nothing changed there are no copies - # to elide, so skip the scan entirely. + # insertion above, then, with skip_d2h_for_method_outputs, make the graph + # output read device tensors. Only meaningful when passes 1/2 modified the + # graph (there was a device delegate); when nothing changed there are no + # copies to elide, so skip the scan entirely. if changed: self._elide_redundant_same_device_copies(graph_module) + if self.skip_d2h_for_method_outputs: + self._keep_method_outputs_on_device(graph_module) graph_module.recompile() return PassResult(graph_module, changed) diff --git a/exir/program/_program.py b/exir/program/_program.py index 2bdb84aa052..f0f8f5e62be 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -794,6 +794,7 @@ def edge_to_executorch_passes( skip_h2d_for_method_inputs=device_cfg.skip_h2d_for_method_inputs, skip_d2h_for_method_outputs=device_cfg.skip_d2h_for_method_outputs, enable_non_cpu_memory_planning=config.enable_non_cpu_memory_planning, + error_on_per_call_copies=device_cfg.error_on_per_call_copies, ), EdgeToBackendOpsPass(), RemoveGraphAssertsPass(), @@ -824,8 +825,10 @@ def _generate_edge_program( program = remove_unused_parameters_pass(program) if not config._skip_dim_order: - # Normalize constants and buffers before MemoryFormatOpsPass rewrites - # preserve-format copies, which cannot represent arbitrary dim orders. + # Normalize constants, buffers and parameters before partitioning, so a + # delegate sees the layout the program stores, and before + # MemoryFormatOpsPass rewrites preserve-format copies, which cannot + # represent arbitrary dim orders. program = convert_constant_dim_order_pass.convert_constant_dim_order_pass( program ) @@ -878,7 +881,10 @@ def _generate_edge_program( # Lift the tensor constants created in ScalarToTensorPass edge_program = lift_constant_tensor_pass(edge_program) - # Normalize tensor constants introduced by the preceding passes. + # Normalize tensor constants introduced by the preceding passes. With + # _skip_dim_order this is the only call, so it fixes the layout of + # parameters, buffers and constants, but ops computed from them keep the + # old strides. edge_program = convert_constant_dim_order_pass.convert_constant_dim_order_pass( edge_program ) diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 083f446019c..2fa00618679 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -3175,6 +3175,115 @@ def test_convert_constant_dim_order_to_contiguous(self): f"Constant should be contiguous after pass, got strides {modified_const.stride()}", ) + def test_convert_constant_dim_order_parameter(self): + """ + Verify that the convert_constant_dim_order_pass converts a parameter + outside the allowed dim orders to contiguous, keeps it a parameter, and + leaves a channels-last parameter as it is. + """ + + class AddParameter(torch.nn.Module): + def __init__(self, weight: torch.Tensor): + super().__init__() + self.weight = torch.nn.Parameter(weight) + + def forward(self, x): + return x + self.weight + + weights = { + "transposed": (torch.randn(5, 4).t(), True), + "channels_last": ( + torch.randn(2, 3, 4, 5).to(memory_format=torch.channels_last), + False, + ), + } + for description, (weight, converted) in weights.items(): + with self.subTest(description=description): + inputs = (torch.randn(weight.shape),) + ep = torch.export.export(AddParameter(weight), inputs) + logs = ( + self.assertLogs(convert_constant_dim_order_pass.logger, "WARNING") + if converted + else self.assertNoLogs(convert_constant_dim_order_pass.logger) + ) + with logs as captured: + modified_ep = ( + convert_constant_dim_order_pass.convert_constant_dim_order_pass( + copy.deepcopy(ep) + ) + ) + if converted: + self.assertIn( + "Parameter 'weight' has strides (1, 4)", captured.output[0] + ) + + torch.testing.assert_close( + ep.module()(*inputs), modified_ep.module()(*inputs) + ) + param = modified_ep.state_dict["weight"] + self.assertIsInstance(param, torch.nn.Parameter) + self.assertEqual(param.is_contiguous(), converted) + placeholder = next( + n for n in modified_ep.graph.nodes if n.name == "p_weight" + ) + self.assertEqual(placeholder.meta["val"].stride(), param.stride()) + + def test_to_edge_makes_transposed_parameter_contiguous(self): + """ + Verify that to_edge stores a transposed parameter contiguous with and + without _skip_dim_order, so a partitioner sees the stored layout. + """ + + class AddParameter(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(5, 4).t()) + + def forward(self, x): + return self.weight + x + + model = AddParameter() + inputs = (torch.randn(4, 5),) + for skip_dim_order in (False, True): + with self.subTest(skip_dim_order=skip_dim_order): + with self.assertLogs( + convert_constant_dim_order_pass.logger, "WARNING" + ) as captured: + ep = to_edge( + torch.export.export(model, inputs), + compile_config=EdgeCompileConfig( + _skip_dim_order=skip_dim_order + ), + ).exported_program() + # Only one run of the pass converts: the first, or with + # _skip_dim_order the only one. + self.assertEqual(len(captured.output), 1) + self.assertIn("Parameter 'weight'", captured.output[0]) + self.assertTrue(ep.state_dict["weight"].is_contiguous()) + placeholder = next(n for n in ep.graph.nodes if n.name == "p_weight") + self.assertTrue(placeholder.meta["val"].is_contiguous()) + torch.testing.assert_close(ep.module()(*inputs), model(*inputs)) + + def test_convert_constant_dim_order_missing_placeholder(self): + """ + Verify that the convert_constant_dim_order_pass names the placeholder it + cannot find for a lifted tensor it converts. + """ + + class AddParameter(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(5, 4).t()) + + def forward(self, x): + return x + self.weight + + ep = torch.export.export(AddParameter(), (torch.randn(4, 5),)) + spec = next(s for s in ep.graph_signature.input_specs if s.target == "weight") + spec.arg = TensorArgument(name="missing") + with self.assertRaisesRegex(RuntimeError, "Missing placeholder for missing"): + convert_constant_dim_order_pass.convert_constant_dim_order_pass(ep) + class TestMemoryFormatOpsPassPreserveFormat(unittest.TestCase): """ diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py index 41902ddced4..612e8a06290 100644 --- a/exir/tests/test_propagate_device_pass.py +++ b/exir/tests/test_propagate_device_pass.py @@ -13,6 +13,7 @@ import executorch.exir.passes._device_copy_ops_registry # noqa: F401 import torch from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.exir.backend.backend_api import to_backend from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import ( DelegationSpec, @@ -29,7 +30,12 @@ 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.passes import MemoryPlanningPass +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes import ( + MemoryPlanningPass, + propagate_device_pass, + ToOutVarPass, +) from executorch.exir.passes.propagate_device_pass import ( _get_target_device_from_compile_specs, _parse_device_spec_value, @@ -102,6 +108,23 @@ def _collect_device_copy_nodes(gm: torch.fx.GraphModule) -> DeviceCopyNodes: ) +def _lower_add_to_device(inputs: tuple, device: str) -> torch.nn.Module: + """Lower a module that returns ``x + y`` to a delegate on ``device``.""" + + class Add(torch.nn.Module): + def forward(self, x, y): + return torch.add(x, y) + + return to_backend( + BackendWithCompilerDemo.__name__, + to_edge(export(Add(), inputs)).exported_program(), + [ + CompileSpec("max_value", bytes([4])), + CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")), + ], + ) + + class PerAddDeviceDelegatePartitioner(Partitioner): """Delegates each ``aten.add.Tensor`` as its OWN delegate, each with a per-op ``target_device``. @@ -111,10 +134,11 @@ class PerAddDeviceDelegatePartitioner(Partitioner): directly-connected adds become two directly-connected delegates -- the shape that leaves a redundant ``device -> host -> device`` round-trip between them. ``devices[i]`` is the target device for the i-th add in graph order (the last - entry is reused if there are more adds than devices). + entry is reused if there are more adds than devices). ``None`` gives that add + a delegate without a target device. """ - def __init__(self, devices: List[str]) -> None: + def __init__(self, devices: List[Optional[str]]) -> None: super().__init__() self._devices = list(devices) @@ -129,12 +153,13 @@ def partition(self, exported_program) -> PartitionResult: device = self._devices[min(i, len(self._devices) - 1)] tag = f"tag{i}" node.meta["delegation_tag"] = tag + compile_specs = [CompileSpec("max_value", bytes([4]))] + if device is not None: + compile_specs.append( + CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")) + ) partition_tags[tag] = DelegationSpec( - BackendWithCompilerDemo.__name__, - [ - CompileSpec("max_value", bytes([4])), - CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")), - ], + BackendWithCompilerDemo.__name__, compile_specs ) return PartitionResult( tagged_exported_program=exported_program, @@ -142,6 +167,80 @@ def partition(self, exported_program) -> PartitionResult: ) +# A mutable op that stays on the host, like a KV cache update. Once +# to_executorch removes auto_functionalized, it adds other into x in place. +_test_lib = torch.library.Library("propagate_device_test", "DEF") +_test_lib.define("accumulate(Tensor(a!) x, Tensor other) -> Tensor") +_test_lib.define( + "accumulate.out(Tensor(a!) x, Tensor other, *, Tensor(b!) out) -> Tensor(b!)" +) + + +@torch.library.impl(_test_lib, "accumulate", "CompositeExplicitAutograd") +def _accumulate(x: torch.Tensor, other: torch.Tensor) -> torch.Tensor: + x.add_(other) + return x.clone() + + +@torch.library.impl(_test_lib, "accumulate.out", "CompositeExplicitAutograd") +def _accumulate_out( + x: torch.Tensor, other: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + x.add_(other) + return out.copy_(x) + + +def _lower_with_host_writer( + model: torch.nn.Module, + inputs: tuple, + devices: List[str], + et_config: ExecutorchBackendConfig, +) -> torch.fx.GraphModule: + """Delegate each add of a model that also calls accumulate on the host.""" + # torch._export.config turns off auto_functionalized v2 for export in some + # PyTorch builds (see is_fbcode() there). Under v1, a custom op that writes + # through a view writes a view_copy and the write is lost, a limitation + # that predates these passes. Pin v2 so these tests check the same graph in + # every build. + with torch._export.config.patch(enable_auto_functionalized_v2_for_export=True): + edge = to_edge( + export(model, inputs), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + et = edge.to_backend(PerAddDeviceDelegatePartitioner(devices)).to_executorch( + et_config + ) + return et.exported_program().graph_module + + +def _host_writer(gm: torch.fx.GraphModule) -> torch.fx.Node: + (writer,) = [ + n + for n in gm.graph.nodes + if n.target == torch.ops.propagate_device_test.accumulate.out + ] + return writer + + +def _host_written(gm: torch.fx.GraphModule) -> torch.fx.Node: + """The tensor that accumulate writes, looking through the view it may write.""" + written = _host_writer(gm).args[0] + if written.target == torch.ops.aten.as_strided.default: + return written.args[0] + return written + + +def _host_writer_config(write_view: bool, **kwargs) -> ExecutorchBackendConfig: + return ExecutorchBackendConfig( + emit_stacktrace=False, + enable_non_cpu_memory_planning=True, + # A write through a view leaves the aten.as_strided that + # to_executorch makes for it, which has no out variant. + to_out_var_pass=ToOutVarPass(ignore_to_out_var_failure=write_view), + **kwargs, + ) + + class TestPropagateDevicePass(unittest.TestCase): @staticmethod def _collect_tensor_specs(node: torch.fx.Node) -> List[TensorSpec]: @@ -450,6 +549,29 @@ def forward(self, a, b): self.assertEqual(len(device_copy_nodes.h2d_nodes), 0) self.assertEqual(len(device_copy_nodes.d2h_nodes), 0) + def test_no_copy_nodes_for_cpu_target_device(self): + """A delegate whose target device is the CPU reads and writes host + memory, so it gets no copies, like a delegate without a target device.""" + + class Model(torch.nn.Module): + def forward(self, a, b): + return torch.add(a, b) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, enable_non_cpu_memory_planning=True + ) + for pipeline, gm in _lower_model_to_executorch( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + PerAddDeviceDelegatePartitioner(["cpu"]), + et_config, + ): + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(len(nodes.delegate_nodes), 1) + self.assertEqual(nodes.h2d_nodes, []) + self.assertEqual(nodes.d2h_nodes, []) + def test_redundant_same_device_copies_elided(self): """Two directly-connected delegates on the *same* device must not leave a redundant device->host->device round-trip between them. The @@ -493,6 +615,458 @@ def forward(self, a, b, c): f"(1 boundary D2H), got {len(nodes.d2h_nodes)}", ) + def test_same_device_copies_kept_when_host_op_writes_copy(self): + """A host op that writes the host copy of a delegate output in place, + directly or through a view, changes what the next delegate on the same + device reads, so that round-trip stays. The input it only reads still + gets one copy for both delegates.""" + + class Model(torch.nn.Module): + def __init__(self, write_view): + super().__init__() + self.write_view = write_view + + def forward(self, x, y): + a = x + y + if self.write_view: + torch.ops.propagate_device_test.accumulate(a.view(-1), y.view(-1)) + else: + torch.ops.propagate_device_test.accumulate(a, y) + return a + y + + for write_view in (False, True): + gm = _lower_with_host_writer( + Model(write_view), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config(write_view), + ) + with self.subTest(write_view=write_view): + written = _host_written(gm) + self.assertEqual(written.target, torch.ops.et_copy._d2h_copy.out) + nodes = _collect_device_copy_nodes(gm) + self.assertIn( + written, + [ + arg.args[0] + for arg in nodes.delegate_nodes[1].args[1:] + if arg.target == torch.ops.et_copy._h2d_copy.out + ], + ) + y_copies = [n for n in nodes.h2d_nodes if n.args[0].name == "y"] + self.assertEqual(len(y_copies), 1) + + def test_host_writes_after_the_device_readers_keep_no_copies(self): + """Host ops that write tensors in place only after the delegates on one + device have read them cannot change what those delegates read. So the + round-trip between the delegates is still removed, and the tensor both + delegates read gets one copy.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + t = x * 2 + a = t + y + b = a + t + # Both writes read b, so they run after the second delegate. + torch.ops.propagate_device_test.accumulate(a, b) + torch.ops.propagate_device_test.accumulate(t, b) + return a, b, t + + gm = _lower_with_host_writer( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config(False), + ) + nodes = _collect_device_copy_nodes(gm) + first, second = nodes.delegate_nodes + order = list(gm.graph.nodes) + writers = [ + n + for n in order + if n.target == torch.ops.propagate_device_test.accumulate.out + ] + self.assertEqual(len(writers), 2) + for writer in writers: + self.assertLess(order.index(second), order.index(writer)) + # Only t and y are copied, and the second delegate reads the first + # delegate's output directly. + self.assertEqual(len(nodes.h2d_nodes), 2) + (shared,) = set(first.args[1:]) & set(second.args[1:]) + self.assertIn(shared, nodes.h2d_nodes) + + def test_single_tensor_delegate_output_read_on_host(self): + """A host op that reads the output of a lowered module returning one + tensor reads a host copy of it.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return torch.sin(self.lowered(x, y)) + + et = to_edge(export(Model(), inputs)).to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, enable_non_cpu_memory_planning=True + ) + ) + gm = et.exported_program().graph_module + nodes = _collect_device_copy_nodes(gm) + (d2h,) = nodes.d2h_nodes + self.assertIs(d2h.args[0], nodes.delegate_nodes[0]) + (sin,) = d2h.users + self.assertEqual(sin.target, torch.ops.aten.sin.out) + self._assert_nodes_buffer_device( + [d2h], et._emitter_output.program, DeviceType.CPU, "host op", "D2H copy" + ) + + def test_single_tensor_delegate_output_copied_to_host(self): + """A lowered module that returns one tensor, called from another module, + has no getitem after its call. Its device output still gets a _d2h_copy, + and the next call on the same device reads it without a round-trip. With + skip_h2d_for_method_inputs, y, read by both calls, is placed once.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return self.lowered(self.lowered(x, y), y) + + configs = ( + PropagateDeviceConfig(), + PropagateDeviceConfig(skip_h2d_for_method_inputs=True), + PropagateDeviceConfig(skip_d2h_for_method_outputs=True), + ) + for device_config in configs: + skip_h2d = device_config.skip_h2d_for_method_inputs + skip_d2h = device_config.skip_d2h_for_method_outputs + et = to_edge(export(Model(), inputs)).to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=device_config, + enable_non_cpu_memory_planning=True, + ) + ) + program = et._emitter_output.program + gm = et.exported_program().graph_module + case = f"skip_h2d={skip_h2d}, skip_d2h={skip_d2h}" + 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 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) + ) + cuda_ph, _ = self._collect_placeholders_by_device(gm) + self.assertEqual(len(cuda_ph), 2 if skip_h2d else 0) + # Only the method output can be copied back to the host. + self.assertEqual(len(nodes.d2h_nodes), 0 if skip_d2h else 1) + (output,) = gm.graph.output_node().args[0] + self.assertEqual( + output.target, + ( + executorch_call_delegate + if skip_d2h + else torch.ops.et_copy._d2h_copy.out + ), + ) + self._assert_nodes_buffer_device( + nodes.delegate_nodes, program, DeviceType.CUDA, case, "Delegate" + ) + self._assert_nodes_buffer_device( + nodes.d2h_nodes, program, DeviceType.CPU, case, "D2H copy" + ) + + def test_host_inplace_ops_on_delegate_output_write_host_copy(self): + """Host ops that write a delegate output in place, here relu_ and then + mul_ from reinplace_pass, write its host copy. Their results share the + copy's spec, as memory planning requires, and the method returns the + last one from the host.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Call(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return torch.relu(self.lowered(x, y)) * 2 + + class Add(torch.nn.Module): + def forward(self, x, y): + return torch.relu(x + y) * 2 + + cases = { + "single tensor": (Call(), None), + "getitem": (Add(), PerAddDeviceDelegatePartitioner(["cuda:0"])), + } + reinplaced = frozenset( + {exir_ops.edge.aten.relu.default, exir_ops.edge.aten.mul.Tensor} + ) + for case, (model, partitioner) in cases.items(): + for skip_d2h in (False, True): + edge = to_edge(export(model, inputs)) + if partitioner is not None: + edge = edge.to_backend(partitioner) + et = edge.to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, + reinplace_extra_ops=reinplaced, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=skip_d2h + ), + enable_non_cpu_memory_planning=True, + ) + ) + gm = et.exported_program().graph_module + with self.subTest(case=case, skip_d2h=skip_d2h): + (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes + (relu,) = d2h.users + (mul,) = relu.users + self.assertEqual(relu.target, exir_ops.edge.aten.relu_.default) + self.assertEqual(mul.target, exir_ops.edge.aten.mul_.Tensor) + self.assertIs(relu.meta["spec"], d2h.meta["spec"]) + self.assertIs(mul.meta["spec"], d2h.meta["spec"]) + self.assertEqual(gm.graph.output_node().args[0], (mul,)) + self._assert_nodes_buffer_device( + [mul], + et._emitter_output.program, + DeviceType.CPU, + case, + "Output", + ) + + def test_skip_d2h_keeps_output_also_read_on_device(self): + """With skip_d2h_for_method_outputs, a returned delegate output that a + delegate on the same device also reads stays on the device, like the + other outputs.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Calls(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + a = self.lowered(x, y) + return a, self.lowered(a, y) + + class Adds(torch.nn.Module): + def forward(self, x, y): + a = torch.add(x, y) + return a, torch.add(a, y) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + ) + cases = { + "single tensor": (Calls(), None), + "getitem": (Adds(), PerAddDeviceDelegatePartitioner(["cuda:0"])), + } + for case, (model, partitioner) in cases.items(): + edge = to_edge(export(model, inputs)) + if partitioner is not None: + edge = edge.to_backend(partitioner) + et = edge.to_executorch(deepcopy(et_config)) + gm = et.exported_program().graph_module + with self.subTest(case=case): + self.assertEqual(len(_collect_device_copy_nodes(gm).d2h_nodes), 0) + outputs = gm.graph.output_node().args[0] + self._assert_nodes_device(outputs, DeviceType.CUDA, case, "Output") + self._assert_nodes_buffer_device( + outputs, et._emitter_output.program, DeviceType.CUDA, case, "Output" + ) + + def test_skip_d2h_keeps_output_also_read_off_device(self): + """With skip_d2h_for_method_outputs, a returned delegate output that a + host op or a delegate on another device also reads stays on the device. + Only that other reader gets the host copy.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + lowered_cuda_1 = _lower_add_to_device(inputs, "cuda:1") + + class HostReader(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + a = self.lowered(x, y) + return a, torch.sin(a) + + class OtherDeviceReader(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + self.lowered_cuda_1 = lowered_cuda_1 + + def forward(self, x, y): + a = self.lowered(x, y) + return a, self.lowered_cuda_1(a, y) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + ) + for case, model in { + "host op": HostReader(), + "delegate on cuda:1": OtherDeviceReader(), + }.items(): + et = to_edge(export(model, inputs)).to_executorch(deepcopy(et_config)) + gm = et.exported_program().graph_module + with self.subTest(case=case): + output = gm.graph.output_node().args[0][0] + self.assertEqual(output.target, executorch_call_delegate) + self._assert_nodes_device([output], DeviceType.CUDA, case, "Output", 0) + self._assert_nodes_buffer_device( + [output], + et._emitter_output.program, + DeviceType.CUDA, + case, + "Output", + ) + (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes + self.assertIs(d2h.args[0], output) + self.assertNotIn(gm.graph.output_node(), d2h.users) + + def test_skip_d2h_leaves_returned_output_unplanned(self): + """With skip_d2h_for_method_outputs and unplanned graph outputs, the + caller provides the returned delegate output, while the host copy that + a host op reads is still planned.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + a = x + y + return a, torch.sin(a) + + et = ( + to_edge(export(Model(), (torch.randn(2, 2), torch.randn(2, 2)))) + .to_backend(PerAddDeviceDelegatePartitioner(["cuda:0"])) + .to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_output=False), + ) + ) + ) + gm = et.exported_program().graph_module + output = gm.graph.output_node().args[0][0] + (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes + self.assertIs(d2h.args[0], output) + self.assertIsNone(output.meta["spec"].mem_id) + self.assertIsNotNone(d2h.meta["spec"].mem_id) + + def test_skip_d2h_returns_output_written_on_host_from_host(self): + """With skip_d2h_for_method_outputs, a returned delegate output whose + host copy a host op writes in place, directly or through a view, is + returned as that written host copy, because the device tensor never + sees the write. No copy goes back to the device. A warning names the + op that writes it, since the copy to the host runs on every call.""" + + class Model(torch.nn.Module): + def __init__(self, write_view): + super().__init__() + self.write_view = write_view + + def forward(self, x, y): + a = x + y + if self.write_view: + torch.ops.propagate_device_test.accumulate(a.view(-1), y.view(-1)) + else: + torch.ops.propagate_device_test.accumulate(a, y) + return a + + for write_view in (False, True): + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + gm = _lower_with_host_writer( + Model(write_view), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0"], + _host_writer_config( + write_view, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + ), + ) + with self.subTest(write_view=write_view): + (warning,) = logs.output + self.assertIn("Method output 0 (16 bytes)", warning) + self.assertIn( + "'propagate_device_test.accumulate.default' writes it", warning + ) + (output,) = gm.graph.output_node().args[0] + self.assertIs(output, _host_written(gm)) + self.assertEqual(output.target, torch.ops.et_copy._d2h_copy.out) + (output_spec,) = gm.graph.output_node().meta["spec"] + self.assertEqual(output_spec.device, DeviceType.CPU) + self.assertFalse( + any( + copy.args[0] is output + for copy in _collect_device_copy_nodes(gm).h2d_nodes + ) + ) + + def test_skip_d2h_copies_a_written_output_only_for_a_later_delegate(self): + """With skip_d2h_for_method_outputs, a returned output that a host op + writes before a later delegate on the same device reads it is copied + back to the device for that delegate only. The method returns the + written host copy.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + a = x + y + torch.ops.propagate_device_test.accumulate(a, y) + return a, a + y + + gm = _lower_with_host_writer( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config( + False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + ), + ) + (copy,) = [ + n + for n in _collect_device_copy_nodes(gm).h2d_nodes + if n.args[0] is _host_written(gm) + ] + _, second = _collect_device_copy_nodes(gm).delegate_nodes + self.assertIn(copy, second.args[1:]) + self.assertIs(gm.graph.output_node().args[0][0], _host_written(gm)) + def test_cross_device_copies_preserved(self): """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two directly-connected delegates must be preserved -- the round-trip is a @@ -563,14 +1137,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( @@ -722,6 +1296,47 @@ def forward(self, a, b): f"first delegate's output directly", ) + def test_tensors_shared_by_delegates_on_different_devices(self): + """Delegates on cuda:0 and cuda:1 that read the same method input and + each other's outputs. The input gets one copy for each device, read by + every delegate there, and each output crosses to the other device + through the host.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + a = x + y # delegate 0 (cuda:0) + b = x + a # delegate 1 (cuda:1) + return a, b, x + b # delegate 2 (cuda:0) + + d2h_out = torch.ops.et_copy._d2h_copy.out + for pipeline, gm in _lower_model_to_executorch( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1", "cuda:0"]), + ExecutorchBackendConfig(enable_non_cpu_memory_planning=True), + ): + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + first, second, third = nodes.delegate_nodes + x_copies = [n for n in nodes.h2d_nodes if n.args[0].name == "x"] + first_x, second_x, third_x = ( + next(arg for arg in delegate.args[1:] if arg in x_copies) + for delegate in (first, second, third) + ) + self.assertEqual(len(x_copies), 2) + self.assertIs(first_x, third_x) + self.assertEqual(first_x.meta["spec"].device_index, 0) + self.assertEqual(second_x.meta["spec"].device_index, 1) + crossings = sorted( + ( + h2d.args[0].args[0].meta["spec"].device_index, + h2d.meta["spec"].device_index, + ) + for h2d in nodes.h2d_nodes + if h2d.args[0].target == d2h_out + ) + self.assertEqual(crossings, [(0, 1), (1, 0)]) + # ---- Integration tests: device consistency after to_executorch ---- def test_device_consistency_cuda_1(self): @@ -1265,6 +1880,546 @@ 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, + ) + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + results = self._get_executorch_program( + model, inputs, DeviceAwarePartitioner("cuda:0"), et_config + ) + self.assertIn("Constant 'b_offset' (16 bytes)", logs.output[0]) + for pipeline, program, gm in results: + 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_error_on_per_call_copies(self): + """error_on_per_call_copies turns the warning about a copy that a skip + flag cannot remove into an error: a constant copied to the device and + an input copied for a host reader under skip_h2d_for_method_inputs, and + a returned output written on the host under skip_d2h_for_method_outputs. + Without the skip flag, a constant copy only logs the warning.""" + + class ReadsConstant(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), self.offset * 2 + + class WritesOutputOnHost(torch.nn.Module): + def forward(self, x, y): + a = x + y + torch.ops.propagate_device_test.accumulate(a, y) + return a + + class ReadsInputOnHost(torch.nn.Module): + def forward(self, x, y): + return x + y, torch.sin(x) + + def lower(model, num_inputs, **flags): + return _lower_with_host_writer( + model, + tuple(torch.randn(2, 2) for _ in range(num_inputs)), + ["cuda:0"], + _host_writer_config( + False, + propagate_device_config=PropagateDeviceConfig( + error_on_per_call_copies=True, **flags + ), + ), + ) + + cases = { + "constant": (ReadsConstant(), 1, "skip_h2d_for_method_inputs"), + "input read on host": (ReadsInputOnHost(), 2, "skip_h2d_for_method_inputs"), + "output written on host": ( + WritesOutputOnHost(), + 2, + "skip_d2h_for_method_outputs", + ), + } + for case, (model, num_inputs, flag) in cases.items(): + with self.subTest(case=case): + with self.assertRaisesRegex( + RuntimeError, + "on every call.*Raised because error_on_per_call_copies=True", + ): + lower(model, num_inputs, **{flag: True}) + with self.assertLogs(propagate_device_pass.logger, "WARNING"): + lower(ReadsConstant(), 1) + + 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: contiguous for a transposed parameter, + which to_edge makes contiguous, and channels-last for 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_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(): + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + et = to_edge(export(Model(devices), inputs[:1])).to_executorch( + deepcopy(et_config) + ) + gm = et.exported_program().graph_module + with self.subTest(case=case): + self.assertEqual(len(logs.output), len(copy_indices)) + 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_tensor_written_on_host_is_copied_for_each_delegate(self): + """A tensor that a host op writes in place, directly or through a view, + between two delegates on one device gets a new copy for the second + delegate, so it reads the new value.""" + + class Model(torch.nn.Module): + def __init__(self, write_view): + super().__init__() + self.write_view = write_view + + def forward(self, x, y): + t = x * 2 + a = t + y + if self.write_view: + torch.ops.propagate_device_test.accumulate(t.view(-1), a.view(-1)) + else: + torch.ops.propagate_device_test.accumulate(t, a) + return a, t + y + + for write_view in (False, True): + gm = _lower_with_host_writer( + Model(write_view), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config(write_view), + ) + with self.subTest(write_view=write_view): + writer = _host_writer(gm) + first, second = _collect_device_copy_nodes(gm).delegate_nodes + (second_copy,) = [ + arg + for arg in second.args[1:] + if arg.target == torch.ops.et_copy._h2d_copy.out + and arg.args[0] is _host_written(gm) + ] + nodes = list(gm.graph.nodes) + self.assertLess(nodes.index(first), nodes.index(writer)) + self.assertLess(nodes.index(writer), nodes.index(second_copy)) + + 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 + _h2d_copy.""" + + class Model(torch.nn.Module): + def forward(self, a, b): + c = torch.sin(torch.add(a, b)) + return torch.add(a, c) + + model = Model() + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + 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, inputs, DeviceAwarePartitioner("cuda:0"), et_config + ): + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(len(nodes.delegate_nodes), 2) + # Only sin's output, computed on the host, is copied. + self.assertEqual(len(nodes.h2d_nodes), 1) + cuda_ph, _ = self._collect_placeholders_by_device(gm) + self.assertEqual(sorted(n.name for n in cuda_ph), ["a", "b"]) + self._assert_nodes_buffer_device( + cuda_ph, program, DeviceType.CUDA, pipeline, "Placeholder" + ) + + def test_skip_h2d_copies_input_for_readers_off_device(self): + """With skip_h2d_for_method_inputs=True, a method input that a device + delegate reads stays on that device. A host op, a delegate on the host, + or a delegate on another device that also reads it reads one _d2h_copy + of it, also when more than one delegate on the device reads it, and a + warning says so.""" + + class HostOp(torch.nn.Module): + def forward(self, a, b): + return torch.add(a, b), torch.sin(a) + + class TwoDelegates(torch.nn.Module): + def forward(self, a, b, c): + return torch.add(a, b), torch.add(a, c) + + class TwoDelegatesAndHostOp(torch.nn.Module): + def forward(self, a, b, c): + return torch.add(a, b), torch.add(a, c), torch.sin(a) + + class CopiedIntoBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(2, 2)) + + def forward(self, a, b): + self.state.copy_(a) + return torch.add(a, b) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + cases = { + "host op": (HostOp(), 2, DeviceAwarePartitioner("cuda:0")), + "two delegates on the device and a host op": ( + TwoDelegatesAndHostOp(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:0"]), + ), + "delegate without a device": ( + TwoDelegates(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", None]), + ), + "delegate on the cpu": ( + TwoDelegates(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", "cpu"]), + ), + "other device index": ( + TwoDelegates(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1"]), + ), + # The write-back of the buffer only reads the input. + "buffer write-back": ( + CopiedIntoBuffer(), + 2, + DeviceAwarePartitioner("cuda:0"), + ), + } + for case, (model, num_inputs, partitioner) in cases.items(): + inputs = tuple(torch.randn(2, 2) for _ in range(num_inputs)) + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + results = self._get_executorch_program( + model, inputs, partitioner, et_config + ) + self.assertIn("Method input 'a' (16 bytes) stays on cuda:0", logs.output[0]) + for pipeline, program, gm in results: + with self.subTest(case=case, pipeline=pipeline): + a = next(n for n in gm.graph.nodes if n.name == "a") + self.assertEqual(a.meta["spec"].device, DeviceType.CUDA) + self.assertEqual(a.meta["spec"].device_index, 0) + nodes = _collect_device_copy_nodes(gm) + self.assertIs(nodes.delegate_nodes[0].args[1], a) + # The copy is the only other reader of the input, and is + # not copied again. + (copy,) = [n for n in a.users if n not in nodes.delegate_nodes] + self.assertIn(copy, nodes.d2h_nodes) + self.assertFalse(set(copy.users) & set(nodes.d2h_nodes)) + self._assert_nodes_buffer_device( + [copy], program, DeviceType.CPU, pipeline, "D2H copy" + ) + + def test_skip_h2d_shape_reader_needs_no_copy(self): + """With skip_h2d_for_method_inputs=True, an op that reads only the shape + of an input on the device reads the input there, since its sizes stay + in host memory, so the input gets no copy and no warning.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + return torch.add(x, y), torch.zeros(x.shape[0]) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + batch = torch.export.Dim("batch", min=2, max=8) + ep = export( + Model(), + (torch.randn(4, 2), torch.randn(4, 2)), + dynamic_shapes={"x": {0: batch}, "y": {0: batch}}, + ) + edge = to_edge( + ep, compile_config=EdgeCompileConfig(_check_ir_validity=False) + ).to_backend(DeviceAwarePartitioner("cuda:0")) + with self.assertNoLogs(propagate_device_pass.logger, "WARNING"): + gm = edge.to_executorch(et_config).exported_program().graph_module + (size,) = [n for n in gm.graph.nodes if n.target == torch.ops.aten.sym_size.int] + read = size.args[0] + self.assertEqual(read.op, "placeholder") + self.assertEqual(read.meta["spec"].device, DeviceType.CUDA) + self.assertFalse( + set(read.users) & set(_collect_device_copy_nodes(gm).d2h_nodes) + ) + + def test_skip_both_returns_input_on_device_without_copy(self): + """With both skip flags, a method that returns an input a device + delegate also reads returns it on the device, with no copy and no + warning.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + return torch.add(x, y), x + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True, skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + ) + with self.assertNoLogs(propagate_device_pass.logger, "WARNING"): + results = self._get_executorch_program( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + DeviceAwarePartitioner("cuda:0"), + et_config, + ) + for pipeline, _program, gm in results: + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(nodes.h2d_nodes + nodes.d2h_nodes, []) + x = next(n for n in gm.graph.nodes if n.name == "x") + self.assertIn(x, gm.graph.output_node().args[0]) + self.assertEqual(x.meta["spec"].device, DeviceType.CUDA) + + def test_skip_h2d_refuses_placeholder_written_on_host(self): + """A mutated input or mutable buffer is written back on the host, so the + refusal says so, also when a host op reads it too. An input that a host + op writes in place is refused as well, since the write would only reach + its host copy.""" + + class MutableBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(2, 2)) + + def forward(self, a): + out = torch.add(a, self.state) + self.state.add_(a) + return out + + class MutableBufferReadOnHost(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(2, 2)) + + def forward(self, a): + out = torch.add(a, self.state) + read = torch.sin(self.state) + self.state.add_(a) + return out, read + + class MutatedInput(torch.nn.Module): + def forward(self, a, x): + out = torch.add(a, x) + x.add_(a) + return out + + class InputWrittenByHostOp(torch.nn.Module): + def forward(self, a, x): + out = torch.add(a, x) + torch.ops.propagate_device_test.accumulate(x, a) + return out + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + write_back = r"'copy_' writes it back on the host \(the method mutates it\)" + cases = { + "mutable buffer": (MutableBuffer(), 1, "b_state", write_back), + "mutable buffer read on host": ( + MutableBufferReadOnHost(), + 1, + "b_state", + write_back, + ), + "mutated input": (MutatedInput(), 2, "x", write_back), + "input written by a host op": ( + InputWrittenByHostOp(), + 2, + "x", + "'[a-z_]*accumulate[a-z_]*' writes it in place on the host", + ), + } + for case, (model, num_inputs, name, reason) in cases.items(): + inputs = tuple(torch.randn(2, 2) for _ in range(num_inputs)) + edge = to_edge( + export(model, inputs), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).to_backend(DeviceAwarePartitioner("cuda:0")) + with self.subTest(case=case): + with self.assertRaisesRegex( + RuntimeError, f"placeholder '{name}' on cuda:0, because {reason}" + ): + edge.to_executorch(deepcopy(et_config)) + def test_tensorspec_repr_includes_device(self): spec = TensorSpec(dtype=torch.float32, shape=torch.Size([2, 3])) repr_str = repr(spec)