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 ef57aca7ae5..e93186b6941 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -208,6 +208,29 @@ Do not pass a dict to `skip_h2d_for_method_inputs` itself. The pass reads that f truthiness, so any non-empty dict enables the skip for every method regardless of the values inside it. +**Some copies stay even with the skip flags.** In the cases below, `to_executorch` keeps a copy +that runs on every call and logs a warning. To remove the copy, have the partitioner put more of +the graph in the delegate. + +- An input stays on the device of the first delegate that reads it. Other readers, such as an + operator outside the delegate, a delegate on another device, or the method output without + `skip_d2h_for_method_outputs`, read a copy in host memory. A delegate on another device then + gets a second copy, from the host. +- A constant, such as a weight, that a delegate reads from outside is copied to each device that + reads it. +- A returned output that an operator outside the delegate writes in place, such as a mutable + custom operator left on the host, is copied to the host and returned there. The method + metadata gives the device of each output. + +Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail +instead of warning. It fails only on a copy that one of the skip flags you set was meant to remove. + +**`skip_h2d_for_method_inputs` refuses inputs that the method changes.** `to_executorch` fails if +the method changes an input, even when the delegate makes the change. This includes a mutable +buffer, such as a KV cache, even though the caller does not pass it. For a buffer, put it and every +operator that uses it in the delegate. The CUDA partitioner already moves the buffer in, with +`tag_mutated_buffer`, once those operators are there. + **The choice is baked into the `.pte`.** A program exported with copies expects host tensors; one exported without them expects device tensors. Passing the wrong kind is a caller error, not something the runtime corrects, so the two are not interchangeable at run time. diff --git a/exir/passes/BUCK b/exir/passes/BUCK index 6c896ea9f5d..7a25cbaa134 100644 --- a/exir/passes/BUCK +++ b/exir/passes/BUCK @@ -523,7 +523,9 @@ fbcode_target(_kind = runtime.python_library, "propagate_device_pass.py", ], deps = [ + ":convert_constant_dim_order_pass", ":device_copy_ops_registry", + ":insert_write_back_for_buffers_pass", ":propagate_device_config", "//caffe2:torch", "//executorch/exir:delegate", 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 6071209ffeb..49b1ffd20a1 100644 --- a/exir/passes/propagate_device_config.py +++ b/exir/passes/propagate_device_config.py @@ -48,28 +48,52 @@ class PropagateDeviceConfig: # When True, method-level input tensors that feed directly into a device # delegate are NOT wrapped with _h2d_copy. The user must provide tensors - # already on the target device. Useful for pipelines where inputs are - # pre-staged on GPU. + # already on the target device: the device of the first delegate that + # reads the input, in graph order. Any other reader of such an input, such + # as a host op, a delegate on another device, or the method output without + # skip_d2h_for_method_outputs, reads a copy of it in host memory, made on + # every call, and the pass logs a warning. A delegate on another device + # then gets a second copy, from the host. An input that an op writes in + # place on the host makes to_executorch raise. Useful for pipelines where + # inputs are pre-staged on GPU. Constants are not passed + # by the caller, so they are still copied: their data stays on the host, + # in the program or in its external data file. That copy runs on every + # call, once for each device that reads the constant, and stays in device + # memory until the last delegate there has read it. # # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer # for the input and copies the caller's memory into it, which is what this exists to avoid. # A dict can be used to set per-method values, keyed by method name. skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False - # When True, device delegate outputs that are directly method outputs - # are NOT wrapped with _d2h_copy. The method outputs stay on device. - # Useful for cross-method GPU pipelines where the next method consumes - # GPU tensors directly. + # When True, device delegate outputs that the method returns stay on + # device, with one exception: when an op outside the delegate, such as a + # mutable custom op left on the host, writes the host copy of such an + # output in place, the method returns that host copy, since the device + # tensor never sees the write. The copy to the host runs on every call, and the pass logs a + # warning about it. + # Useful for cross-method GPU pipelines where the next method consumes GPU + # tensors directly. # # Pair with MemoryPlanningPass(alloc_graph_output=False) for the same reason as the input flag. # A dict can be used to set per-method values, keyed by method name. skip_d2h_for_method_outputs: Union[bool, Dict[str, bool]] = False + # When True, the copies that the skip flags above cannot remove make + # to_executorch raise instead of logging a warning: with + # skip_h2d_for_method_inputs, a constant copied to the device and a method + # input copied to the host for a reader off its device, and with + # skip_d2h_for_method_outputs, a returned output copied to the host for a + # host op that writes it. Each runs on every call. Use it to make sure a + # program that skips its boundary copies has none left. + error_on_per_call_copies: bool = False + def __hash__(self) -> int: return hash( ( str(self.skip_h2d_for_method_inputs), str(self.skip_d2h_for_method_outputs), + self.error_on_per_call_copies, ) ) @@ -79,4 +103,5 @@ def __eq__(self, other: object) -> bool: return ( self.skip_h2d_for_method_inputs == other.skip_h2d_for_method_inputs and self.skip_d2h_for_method_outputs == other.skip_d2h_for_method_outputs + and self.error_on_per_call_copies == other.error_on_per_call_copies ) 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 ed55b345362..b1d896281c0 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -793,6 +793,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(), @@ -823,8 +824,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 ) @@ -870,7 +873,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 1afaece0f3a..fe3bc675ff4 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -3115,6 +3115,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)