make device support config method-based (#19970)

Differential Revision: D101243687

Pull Request resolved: #19970

Author: Gasoonjia

backends/cuda/tests/test_cuda_export.py

@@ -385,11 +385,13 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         # Both input and output tensors should be on CUDA device for now.
         self.assertEqual(
             len(cpu_tensors),
-            3,
-            f"Expecteed three CPU tensors for method inputs and outputs, but found {len(cpu_tensors)}",
+            0,
+            f"Expected no CPU tensors: method inputs/outputs should be tagged "
+            f"CUDA, but found {len(cpu_tensors)}",
         )
         self.assertEqual(
             len(cuda_tensors),
             3,
-            "Expected CUDA tensors for delegate outputs",
+            f"Expected 3 CUDA tensors (2 method inputs + 1 method output), "
+            f"but found {len(cuda_tensors)}",
         )

exir/capture/BUCK

@@ -47,6 +47,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/exir:pass_manager",
         "//executorch/exir:tracer",
         "//executorch/exir/passes:lib",
+        "//executorch/exir/passes:propagate_device_config",
         "//executorch/exir/passes:sym_shape_eval_pass",
     ],
 )

exir/capture/_config.py

@@ -13,6 +13,7 @@
 from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
 from executorch.exir.pass_manager import PassType
 from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
+from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
 from executorch.exir.tracer import ExirDynamoConfig
 from torch.fx._compatibility import compatibility
@@ -60,6 +61,13 @@ class ExecutorchBackendConfig:
     # A single memory planning pass can be defined for all the programs in the
     # EdgeProgramManager or can be defined per program.
     memory_planning_pass: Union[PassType, Dict[str, PassType]] = MemoryPlanningPass()
+
+    # A single propagate device config can be defined for all the programs in the
+    # EdgeProgramManager or can be defined per program.
+    propagate_device_config: Union[
+        PropagateDeviceConfig, Dict[str, PropagateDeviceConfig]
+    ] = field(default_factory=PropagateDeviceConfig)
+
     to_out_var_pass: PassType = ToOutVarPass(ignore_to_out_var_failure=False)
     dynamic_memory_planning_mode: DynamicMemoryPlanningMode = (
         DynamicMemoryPlanningMode.UPPER_BOUND
@@ -124,18 +132,6 @@ class ExecutorchBackendConfig:
     # where all tensors are planned into CPU memory regardless of device.
     enable_non_cpu_memory_planning: bool = False
 
-    # 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.
-    skip_h2d_for_method_inputs: 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.
-    skip_d2h_for_method_outputs: bool = False
-
     # Add ops to the set of re-inplace ops to be used by the reinplace pass.
     # Re-inplace pass checks the eligibility of an op to be re-inplaced and
     # memory planning pass allcoates the output buffer of the op to be the same

exir/passes/BUCK

@@ -460,13 +460,24 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "propagate_device_config",
+    srcs = [
+        "propagate_device_config.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "propagate_device_pass",
     srcs = [
         "propagate_device_pass.py",
     ],
     deps = [
         ":device_copy_ops_registry",
+        ":propagate_device_config",
         "//caffe2:torch",
         "//executorch/exir:delegate",
         "//executorch/exir:lowered_backend_module",

exir/passes/memory_planning_pass.py

@@ -153,7 +153,6 @@ def __init__(
         alloc_mutable_buffers: bool = True,
         share_mutable_buffers: bool = False,
         alignment: int = ALIGNMENT,
-        enable_non_cpu_memory_planning: bool = False,
     ) -> None:
         r"""
         alloc_graph_input/alloc_graph_output will have 4 different combinations
@@ -174,8 +173,11 @@ def __init__(
         self.alloc_mutable_buffers = alloc_mutable_buffers
         self.share_mutable_buffers = share_mutable_buffers
         self.alignment = alignment
-        self.enable_non_cpu_memory_planning = enable_non_cpu_memory_planning
         self.state = _MemoryPlanningState()
+        # Set by EdgeProgramManager.to_executorch() from the top-level
+        # ExecutorchBackendConfig. When True, apply_algo partitions specs by
+        # device so non-CPU buffers get their own memory arenas.
+        self.enable_non_cpu_memory_planning: bool = False
 
     def _set_alloc_node_spec(self, graph_module: torch.fx.GraphModule) -> None:
         """

exir/passes/propagate_device_config.py

@@ -0,0 +1,56 @@
+# 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
+
+"""
+Configuration for PropagateDevicePass.
+
+This is intentionally kept in a lightweight module (no heavy imports such as
+the et_copy op registry) so that ``ExecutorchBackendConfig`` -- which is
+imported throughout the codebase -- can reference ``PropagateDeviceConfig``
+without pulling in the device-copy op registration as an import-time side
+effect.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, Union
+
+from torch.fx._compatibility import compatibility
+
+
+@compatibility(is_backward_compatible=False)
+@dataclass
+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.
+    # 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.
+    # 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
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                str(self.skip_h2d_for_method_inputs),
+                str(self.skip_d2h_for_method_outputs),
+            )
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, PropagateDeviceConfig):
+            return False
+        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
+        )

exir/passes/propagate_device_pass.py

@@ -19,6 +19,13 @@
 import torch
 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
 
@@ -172,6 +179,17 @@ def __init__(
         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"

exir/program/BUCK

@@ -41,6 +41,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/exir/passes:insert_write_back_for_buffers_pass",
         "//executorch/exir/passes:lib",
         "//executorch/exir/passes:normalize_view_copy_base_pass",
+        "//executorch/exir/passes:propagate_device_config",
         "//executorch/exir/passes:propagate_device_pass",
         "//executorch/exir/passes:remove_graph_asserts_pass",
         "//executorch/exir/passes:remove_mixed_type_operators",

exir/program/_program.py

@@ -59,6 +59,7 @@
 from executorch.exir.passes.normalize_view_copy_base_pass import (
     NormalizeViewCopyBasePass,
 )
+from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig
 from executorch.exir.passes.propagate_device_pass import PropagateDevicePass
 from executorch.exir.passes.quant_fusion_pass import quant_fusion_and_const_prop_pass
 from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS, reinplace_pass
@@ -758,15 +759,22 @@ def edge_to_executorch_passes(
     Returns a list of passes to lower from edge to executorch.
     Get the pre memory planning passes based on the method name, if the pass is not in the dict, use the default pass.
     """
+    # Handle propagate device config
+    propagate_device_config = config.propagate_device_config
+    if isinstance(propagate_device_config, dict):
+        device_cfg = propagate_device_config.get(name, PropagateDeviceConfig())
+    else:
+        device_cfg = propagate_device_config
+
     passes: List[PassType] = [
         # ExecuTorch backend ops are unable to handle unbacked symints. So after
         # this pass, passes cannot be Interpreter-based, because it will fail if
         # there exists an unbacked symint operation.
         *config.passes,
         SpecPropPass(),
         PropagateDevicePass(
-            skip_h2d_for_method_inputs=config.skip_h2d_for_method_inputs,
-            skip_d2h_for_method_outputs=config.skip_d2h_for_method_outputs,
+            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,
         ),
         EdgeToBackendOpsPass(),

exir/tests/test_propagate_device_pass.py

@@ -32,6 +32,7 @@
 from executorch.exir.passes.propagate_device_pass import (
     _get_target_device_from_compile_specs,
     _parse_device_spec_value,
+    PropagateDeviceConfig,
     TARGET_DEVICE_COMPILE_SPEC_KEY,
 )
 from executorch.exir.schema import DeviceType
@@ -766,7 +767,9 @@ def forward(self, a, b):
         inputs = (torch.randn(2, 2), torch.randn(2, 2))
         et_config = ExecutorchBackendConfig(
             emit_stacktrace=False,
-            skip_h2d_for_method_inputs=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
             enable_non_cpu_memory_planning=True,
         )
 
@@ -822,7 +825,9 @@ def forward(self, a, b):
         inputs = (torch.randn(2, 2), torch.randn(2, 2))
         et_config = ExecutorchBackendConfig(
             emit_stacktrace=False,
-            skip_d2h_for_method_outputs=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_d2h_for_method_outputs=True
+            ),
             enable_non_cpu_memory_planning=True,
         )
 
@@ -876,8 +881,10 @@ def forward(self, a, b):
         inputs = (torch.randn(2, 2), torch.randn(2, 2))
         et_config = ExecutorchBackendConfig(
             emit_stacktrace=False,
-            skip_h2d_for_method_inputs=True,
-            skip_d2h_for_method_outputs=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True,
+                skip_d2h_for_method_outputs=True,
+            ),
             enable_non_cpu_memory_planning=True,
         )
 
@@ -952,7 +959,9 @@ def forward(self, a, b):
         inputs = (torch.randn(2, 2), torch.randn(2, 2))
         et_config = ExecutorchBackendConfig(
             emit_stacktrace=False,
-            skip_h2d_for_method_inputs=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
             enable_non_cpu_memory_planning=True,
         )