make device support config method-based

Summary:
Device memory planning and H2D/D2H copy configuration were previously global
flags on ExecutorchBackendConfig, applied uniformly across all methods in a
multi-method program. This made it impossible to configure different behaviors
per method — e.g., skipping H2D copies for one method while keeping them for
another.

This diff introduces PropagateDeviceConfig dataclass that groups
   `skip_h2d_for_method_inputs` and `skip_d2h_for_method_outputs`, with each
   field accepting either a single bool or a Dict[str, bool] for per-method
   overrides. A new `propagate_device_config` field on ExecutorchBackendConfig
   similarly accepts either a single config or Dict[str, PropagateDeviceConfig].

Reviewed By: JacobSzwejbka

Differential Revision: D101243687

Author: Gasoonjia

docs/source/compiler-memory-planning.md

@@ -94,7 +94,7 @@ Users attempting to write a custom memory planning algorithm should start by loo
 
 ## Device-Aware Memory Planning
 
-When `enable_non_cpu_memory_planning=True` is set on `ExecutorchBackendConfig`,
+When `enable_non_cpu_memory_planning=True` is set on `MemoryPlanningPass`,
 the memory planning pass partitions tensor specs by their device type and runs
 the planning algorithm independently for each device.  This produces separate
 memory buffers for each device (e.g. CPU vs. CUDA), ensuring that device memory
@@ -103,7 +103,7 @@ and host memory are never mixed.
 ```python
 program = edge_program.to_executorch(
             exir.ExecutorchBackendConfig(
-                enable_non_cpu_memory_planning=True,
+                memory_planning_pass=exir.passes.MemoryPlanningPass(enable_non_cpu_memory_planning=True),
             )
         )
 ```

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_pass 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

exir/emit/test/test_emit.py

@@ -2185,9 +2185,10 @@ def forward(self, x):
             ExecutorBackendPartitioner()
         ).to_executorch()
 
-        # Check that there is only one delegate because two methods are exactly the same
+        # Check that there are two delegates now because the
+        # passes might apply differently due to per-method config support.
         self.assertEqual(
-            len(edge_program_manager.executorch_program.backend_delegate_data), 1
+            len(edge_program_manager.executorch_program.backend_delegate_data), 2
         )
 
     def test_delegate_deduplicate_with_different_compile_specs(self) -> None:

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:
         """
@@ -237,11 +239,14 @@ def run(
         self,
         graph_module: torch.fx.GraphModule,
         graph_signature: Optional[ExportGraphSignature] = None,
+        enable_non_cpu_memory_planning: Optional[bool] = None,
     ) -> PassResult:
         """
         A pass for memory planning. The actual algorithm used will be picked by
         memory_planning_algo
         """
+        if enable_non_cpu_memory_planning is None:
+            enable_non_cpu_memory_planning = self.enable_non_cpu_memory_planning
         self._set_alloc_node_spec(graph_module)
         # TODO(shunting) if people have concern of adding a field to GraphModule
         # directly, we should define a GraphModule subclass that we can add our
@@ -259,7 +264,7 @@ def run(
             # If mutable buffers are shared, then do not allocate them in the
             # main memory planning algo; they are allocated in run_multimethod.
             self.alloc_mutable_buffers and not self.share_mutable_buffers,
-            self.enable_non_cpu_memory_planning,
+            enable_non_cpu_memory_planning,
         )
 
         if self.share_mutable_buffers and graph_signature is not None:

exir/passes/propagate_device_pass.py

@@ -9,7 +9,8 @@
 import copy
 import logging
 import operator
-from typing import Optional
+from dataclasses import dataclass
+from typing import Dict, Optional, Union
 
 # Import to register the et_copy ops so torch.ops.et_copy is available.
 import executorch.exir.passes._device_copy_ops_registry  # noqa: F401
@@ -20,6 +21,7 @@
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.lowered_backend_module import LoweredBackendModule
 from executorch.exir.tensor import TensorSpec
+from torch.fx._compatibility import compatibility
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -30,6 +32,40 @@
 TARGET_DEVICE_COMPILE_SPEC_KEY = "target_device"
 
 
+@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):
+        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
+        )
+
+
 def _parse_device_spec_value(value: bytes) -> tuple[schema.DeviceType, int]:
     """
     Parse a target_device CompileSpec value (e.g., b"cuda:0") into

exir/program/_program.py

@@ -59,7 +59,10 @@
 from executorch.exir.passes.normalize_view_copy_base_pass import (
     NormalizeViewCopyBasePass,
 )
-from executorch.exir.passes.propagate_device_pass import PropagateDevicePass
+from executorch.exir.passes.propagate_device_pass import (
+    PropagateDeviceConfig,
+    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
 from executorch.exir.passes.remove_graph_asserts_pass import (
@@ -758,15 +761,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

@@ -29,9 +29,11 @@
 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.memory_planning_pass import MemoryPlanningPass
 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,8 +768,12 @@ 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,
-            enable_non_cpu_memory_planning=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            memory_planning_pass=MemoryPlanningPass(
+                enable_non_cpu_memory_planning=True
+            ),
         )
 
         for pipeline, program, gm in self._get_executorch_program(
@@ -822,8 +828,12 @@ 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,
-            enable_non_cpu_memory_planning=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_d2h_for_method_outputs=True
+            ),
+            memory_planning_pass=MemoryPlanningPass(
+                enable_non_cpu_memory_planning=True
+            ),
         )
 
         for pipeline, program, gm in self._get_executorch_program(
@@ -876,9 +886,13 @@ 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,
-            enable_non_cpu_memory_planning=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True,
+                skip_d2h_for_method_outputs=True,
+            ),
+            memory_planning_pass=MemoryPlanningPass(
+                enable_non_cpu_memory_planning=True
+            ),
         )
 
         for pipeline, program, gm in self._get_executorch_program(
@@ -952,8 +966,12 @@ 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,
-            enable_non_cpu_memory_planning=True,
+            propagate_device_config=PropagateDeviceConfig(
+                skip_h2d_for_method_inputs=True
+            ),
+            memory_planning_pass=MemoryPlanningPass(
+                enable_non_cpu_memory_planning=True
+            ),
         )
 
         for pipeline, program, gm in self._get_executorch_program(

runtime/executor/test/method_meta_test.cpp

@@ -248,21 +248,30 @@ TEST_F(MethodMetaTest, MethodMetaBufferDeviceReturnsCudaForDeviceBuffer) {
   ASSERT_EQ(method_meta.error(), Error::Ok);
 
   // ModuleAddWithDevice exports with enable_non_cpu_memory_planning=True.
-  // The model delegates add(a,b) to CUDA, producing:
-  //   non_const_buffer_sizes: [0, 48]  (index 0 reserved)
-  //   non_const_buffer_device: [{buffer_idx=1, device_type=CUDA,
-  //   device_index=0}]
-  // So there is exactly 1 planned buffer (user-facing index 0), on CUDA.
-  ASSERT_EQ(method_meta->num_memory_planned_buffers(), 1);
-
-  // Buffer 0 should be CUDA device.
-  auto device = method_meta->memory_planned_buffer_device(0);
-  ASSERT_TRUE(device.ok());
-  EXPECT_EQ(device->type(), executorch::runtime::etensor::DeviceType::CUDA);
-  EXPECT_EQ(device->index(), 0);
+  // The model delegates add(a,b) to CUDA with H2D/D2H copies:
+  //   - non_const_buffer_sizes: [0, 32, 48]
+  //     (index 0 reserved, buffer 0 = 32 bytes CPU for inputs,
+  //      buffer 1 = 48 bytes CUDA for delegate output)
+  //   - non_const_buffer_device: [{buffer_idx=2, device_type=CUDA,
+  //     device_index=0}]
+  // So there are 2 planned buffers: user-facing index 0 (CPU) and index 1
+  // (CUDA).
+  ASSERT_EQ(method_meta->num_memory_planned_buffers(), 2);
+
+  // Buffer 0 should be CPU device (method inputs).
+  auto device0 = method_meta->memory_planned_buffer_device(0);
+  ASSERT_TRUE(device0.ok());
+  EXPECT_EQ(device0->type(), executorch::runtime::etensor::DeviceType::CPU);
+  EXPECT_EQ(device0->index(), 0);
+
+  // Buffer 1 should be CUDA device (delegate output).
+  auto device1 = method_meta->memory_planned_buffer_device(1);
+  ASSERT_TRUE(device1.ok());
+  EXPECT_EQ(device1->type(), executorch::runtime::etensor::DeviceType::CUDA);
+  EXPECT_EQ(device1->index(), 0);
 
   // Out of range should return error.
   EXPECT_EQ(
-      method_meta->memory_planned_buffer_device(1).error(),
+      method_meta->memory_planned_buffer_device(2).error(),
       Error::InvalidArgument);
 }