propagate_device_pass.py

# 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.

# pyre-strict

import copy
import logging
import operator
from typing import Optional

# Import to register the et_copy ops so torch.ops.et_copy is available.
import executorch.exir.passes._device_copy_ops_registry  # noqa: F401
import executorch.exir.schema as schema
import torch
import torch.utils._pytree as pytree
from executorch.exir.delegate import executorch_call_delegate
from executorch.exir.lowered_backend_module import LoweredBackendModule

# Re-exported for backward compatibility; the dataclass lives in a lightweight
# module so that ExecutorchBackendConfig can reference it without importing the
# et_copy op registry above.
from executorch.exir.passes.propagate_device_config import (  # noqa: F401
    PropagateDeviceConfig,
)
from executorch.exir.tensor import TensorSpec
from torch.fx.passes.infra.pass_base import PassBase, PassResult

logger: logging.Logger = logging.getLogger(__name__)

# CompileSpec key convention for specifying the target device.
# Partitioners that target a specific device should include a CompileSpec entry
# with this key and a value encoding the device string (e.g., b"cuda:0").
TARGET_DEVICE_COMPILE_SPEC_KEY = "target_device"


def _parse_device_spec_value(value: bytes) -> tuple[schema.DeviceType, int]:
    """
    Parse a target_device CompileSpec value (e.g., b"cuda:0") into
    (DeviceType, device_index).

    The type portion is matched case-insensitively against schema.DeviceType
    member names (e.g., "cpu", "cuda").  Raises ValueError for unknown types.
    """
    device_str = value.decode("utf-8").strip().lower()
    if ":" in device_str:
        type_str, index_str = device_str.split(":", 1)
        device_index = int(index_str)
    else:
        type_str = device_str
        device_index = 0
    device_type = next(
        (dt for dt in schema.DeviceType if dt.name.lower() == type_str),
        None,
    )
    if device_type is None:
        valid = ", ".join(dt.name for dt in schema.DeviceType)
        raise ValueError(f"Unknown device type '{type_str}'. Valid types: {valid}")
    return device_type, device_index


def _get_lowered_module(
    graph_module: torch.fx.GraphModule,
    delegate_call_node: torch.fx.Node,
) -> Optional[LoweredBackendModule]:
    """
    Given an executorch_call_delegate node, retrieve the associated
    LoweredBackendModule from the graph module.
    The first argument to executorch_call_delegate is a get_attr node
    whose target names the LoweredBackendModule attribute.
    """
    if len(delegate_call_node.args) < 1:
        return None
    lowered_node = delegate_call_node.args[0]
    if not isinstance(lowered_node, torch.fx.Node) or lowered_node.op != "get_attr":
        return None
    lowered_module = getattr(graph_module, lowered_node.target, None)
    if isinstance(lowered_module, LoweredBackendModule):
        return lowered_module
    return None


def _get_target_device_from_compile_specs(
    lowered_module: LoweredBackendModule,
) -> Optional[tuple[schema.DeviceType, int]]:
    """
    Look for a CompileSpec with key TARGET_DEVICE_COMPILE_SPEC_KEY and return
    the corresponding (DeviceType, device_index), or None if not found.
    """
    for spec in lowered_module.compile_specs:
        if spec.key == TARGET_DEVICE_COMPILE_SPEC_KEY:
            return _parse_device_spec_value(spec.value)
    return None


def _set_device_on_spec(
    spec: TensorSpec,
    device_type: schema.DeviceType,
    device_index: int = 0,
) -> None:
    """Set the device attribute on a TensorSpec."""
    spec.device = device_type
    spec.device_index = device_index


def _tag_specs_with_device(
    specs: object,
    device_type: schema.DeviceType,
    device_index: int = 0,
) -> bool:
    """Apply device annotation to a TensorSpec or a collection of TensorSpecs.

    Args:
        specs: A TensorSpec, a tuple/list of TensorSpecs, or None.
        device_type: The target device type to set.
        device_index: The device index (e.g., 0 for cuda:0, 1 for cuda:1).

    Returns:
        True if any spec was modified, False otherwise.
    """
    if specs is None:
        return False
    if isinstance(specs, TensorSpec):
        _set_device_on_spec(specs, device_type, device_index)
        return True
    if isinstance(specs, (tuple, list)):
        changed = False
        for s in specs:
            if isinstance(s, TensorSpec):
                _set_device_on_spec(s, device_type, device_index)
                changed = True
        return changed
    return False


def _clone_spec_with_device(
    spec: TensorSpec,
    device_type: schema.DeviceType,
    device_index: int = 0,
) -> TensorSpec:
    """Create a copy of a TensorSpec with a different device."""
    new_spec = copy.copy(spec)
    new_spec.init_mem_planning_fields()
    _set_device_on_spec(new_spec, device_type, device_index)
    return new_spec


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.
    - 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.
    """

    def __init__(
        self,
        skip_h2d_for_method_inputs: bool = False,
        skip_d2h_for_method_outputs: bool = False,
        enable_non_cpu_memory_planning: 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

        if (
            skip_h2d_for_method_inputs or skip_d2h_for_method_outputs
        ) and not enable_non_cpu_memory_planning:
            raise ValueError(
                "skip_h2d_for_method_inputs and skip_d2h_for_method_outputs are "
                "only meaningful when enable_non_cpu_memory_planning=True, since "
                "they control host/device copy insertion which only happens during "
                "device-aware memory planning. Set enable_non_cpu_memory_planning="
                "True, or leave the skip options disabled."
            )

    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 _insert_h2d_copies(
        self,
        graph_module: torch.fx.GraphModule,
        node: torch.fx.Node,
        target_device_type: schema.DeviceType,
        device_index: int,
    ) -> bool:
        """Insert H2D copy nodes for each tensor input to a delegate call."""
        changed = False
        new_args = list(node.args)
        for i, arg in enumerate(node.args[1:], start=1):
            if not isinstance(arg, torch.fx.Node):
                continue
            arg_spec = arg.meta.get("spec")
            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)
                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
                )
                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

        node.args = tuple(new_args)
        return changed

    def _insert_d2h_for_getitem(
        self,
        graph_module: torch.fx.GraphModule,
        node: torch.fx.Node,
    ) -> bool:
        """If *node* is a getitem extracting from a delegate call, tag its spec
        with the delegate device and insert a D2H copy after it."""
        source_node = node.args[0]
        if not (
            isinstance(source_node, torch.fx.Node)
            and source_node.op == "call_function"
            and source_node.target == executorch_call_delegate
        ):
            return False

        spec = node.meta.get("spec")
        source_specs = source_node.meta.get("spec")
        idx = node.args[1]
        if not (
            isinstance(spec, TensorSpec)
            and isinstance(source_specs, (tuple, list))
            and isinstance(idx, int)
            and idx < len(source_specs)
        ):
            return False

        source_spec = source_specs[idx]
        if not isinstance(source_spec, TensorSpec):
            return False

        _set_device_on_spec(spec, source_spec.device, source_spec.device_index)

        if self.skip_d2h_for_method_outputs and self._feeds_directly_to_output(node):
            return True

        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_node.meta["spec"] = d2h_spec
            d2h_node.meta["val"] = node.meta.get("val")
            if "tensor_meta" in node.meta:
                d2h_node.meta["tensor_meta"] = node.meta["tensor_meta"]

            node.replace_all_uses_with(
                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

    def _elide_redundant_same_device_copies(
        self,
        graph_module: torch.fx.GraphModule,
    ) -> None:
        """Remove ``_h2d_copy(_d2h_copy(x))`` round-trips that never leave a device.

        H2D and D2H copies are inserted per delegate boundary independently (an
        H2D before every device delegate input, a D2H after every device delegate
        output), so two consecutive delegates on the *same* device leave a
        redundant ``device -> host -> device`` bounce: the first delegate's output
        is copied to host by a ``_d2h_copy`` and immediately copied back to the
        same device by a ``_h2d_copy`` feeding the second delegate.

        This walks each ``_h2d_copy`` whose input is a ``_d2h_copy`` and, when the
        tensor before the D2H already lives on the device the H2D copies back to,
        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.
        """
        h2d = torch.ops.et_copy._h2d_copy.default
        d2h = torch.ops.et_copy._d2h_copy.default
        changed = False
        for node in list(graph_module.graph.nodes):
            if node.op != "call_function" or node.target != h2d:
                continue
            src = node.args[0]
            if not (
                isinstance(src, torch.fx.Node)
                and src.op == "call_function"
                and src.target == d2h
            ):
                continue

            x = src.args[0]  # the tensor before the d2h (the device delegate output)
            x_spec = x.meta.get("spec") if isinstance(x, torch.fx.Node) else None
            h2d_spec = node.meta.get("spec")
            if not (
                isinstance(x_spec, TensorSpec) and isinstance(h2d_spec, TensorSpec)
            ):
                continue

            # Only elide a proven same-device round-trip.
            if (
                x_spec.device != h2d_spec.device
                or x_spec.device_index != h2d_spec.device_index
            ):
                continue

            node.replace_all_uses_with(x)
            graph_module.graph.erase_node(node)
            changed = True

        if changed:
            # Drop the d2h copies that are now dead (kept if they still feed a
            # host consumer).
            graph_module.graph.eliminate_dead_code()

    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
        #            output specs with the target device.  Delegates without a
        #            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.
        #   Pass 3 – Elide the redundant device->host->device round-trips left
        #            between consecutive same-device delegates by passes 1 and 2.
        changed = False
        device_delegates: set[torch.fx.Node] = set()

        # Pass 1: insert H2D copies and tag delegate output specs.
        for node in list(graph_module.graph.nodes):
            if node.op == "call_function" and node.target == executorch_call_delegate:
                lowered_module = _get_lowered_module(graph_module, node)
                if lowered_module is None:
                    raise RuntimeError(
                        f"executorch_call_delegate node '{node.name}' does not reference "
                        "a valid LoweredBackendModule. The first argument must be a "
                        "get_attr node pointing to a LoweredBackendModule attribute."
                    )

                result = _get_target_device_from_compile_specs(lowered_module)
                if result is None:
                    continue

                target_device_type, device_index = result
                device_delegates.add(node)

                if self.enable_non_cpu_memory_planning:
                    changed |= self._insert_h2d_copies(
                        graph_module, node, target_device_type, device_index
                    )
                else:
                    for arg in node.args[1:]:
                        if isinstance(arg, torch.fx.Node):
                            changed |= _tag_specs_with_device(
                                arg.meta.get("spec"),
                                target_device_type,
                                device_index,
                            )

                changed |= _tag_specs_with_device(
                    node.meta.get("spec"),
                    target_device_type,
                    device_index,
                )

                logger.debug(
                    "PropagateDevicePass: set device=%s on delegate node %s "
                    "(backend=%s)",
                    target_device_type,
                    node.name,
                    lowered_module.backend_id,
                )

        # Second pass: propagate device through getitem nodes and insert D2H
        # 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:
                source = node.args[0]
                if isinstance(source, torch.fx.Node) and source in device_delegates:
                    if self.enable_non_cpu_memory_planning:
                        changed |= self._insert_d2h_for_getitem(graph_module, node)
                    else:
                        spec = node.meta.get("spec")
                        source_specs = source.meta.get("spec")
                        idx = node.args[1]
                        if (
                            isinstance(spec, TensorSpec)
                            and isinstance(source_specs, (tuple, list))
                            and isinstance(idx, int)
                            and idx < len(source_specs)
                        ):
                            source_spec = source_specs[idx]
                            if isinstance(source_spec, TensorSpec):
                                _set_device_on_spec(
                                    spec,
                                    source_spec.device,
                                    source_spec.device_index,
                                )
                                changed = True

        # 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.
        if changed:
            self._elide_redundant_same_device_copies(graph_module)

        graph_module.recompile()
        return PassResult(graph_module, changed)