From 499bd947a4baabbfd070b02ea7b68331e23a90e0 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Sun, 9 Aug 2026 16:15:54 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- exir/passes/propagate_device_config.py | 26 +++++++++ exir/tests/test_propagate_device_pass.py | 70 ++++++++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py index d1896d10b63..6071209ffeb 100644 --- a/exir/passes/propagate_device_config.py +++ b/exir/passes/propagate_device_config.py @@ -25,10 +25,34 @@ @compatibility(is_backward_compatible=False) @dataclass class PropagateDeviceConfig: + """Controls whether the runtime copies method inputs and outputs across the device boundary. + + Skipping a copy also means the runtime must not reserve its own buffer for the tensor, or it + would fill that buffer from the caller's memory and the copy would come back. Memory planning + allocates graph inputs and outputs by default, so a program using either skip below needs:: + + ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True, + skip_d2h_for_method_outputs=True, + ), + enable_non_cpu_memory_planning=True, + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), + ) + + Without ``alloc_graph_input=False`` the runtime reserves device memory for an input the caller + already owns, then copies into it with a host memcpy, which is undefined for device memory. + """ + # 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. + # + # 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 @@ -36,6 +60,8 @@ class PropagateDeviceConfig: # 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. + # + # 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 diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py index e67be7b400a..ba5e87f2eb4 100644 --- a/exir/tests/test_propagate_device_pass.py +++ b/exir/tests/test_propagate_device_pass.py @@ -29,6 +29,7 @@ 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.passes.propagate_device_pass import ( _get_target_device_from_compile_specs, _parse_device_spec_value, @@ -1270,6 +1271,75 @@ def test_tensorspec_repr_includes_device(self): self.assertIn("device=", repr_str) self.assertIn("CPU", repr_str) + def test_skipping_copies_requires_unallocated_graph_io(self): + """Skipping a boundary copy is only complete when the runtime also stops reserving its own + buffer for that tensor. + + Memory planning allocates graph inputs and outputs by default. A planned input carries + allocation info, and the runtime fills a planned input by copying the caller's memory into the + buffer it reserved, so the copy the export asked to skip comes back at run time. For device + memory that copy is also a host memcpy into a device pointer, which crashes. + + This pins the pairing the configuration needs, because the two settings live in different + places and nothing else connects them. + """ + + class Model(torch.nn.Module): + def forward(self, a, b): + return torch.add(a, b) + + model = Model() + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + skip_copies = PropagateDeviceConfig( + skip_h2d_for_method_inputs=True, + skip_d2h_for_method_outputs=True, + ) + + def planned_io(config: ExecutorchBackendConfig): + """Whether the program reserves its own buffer for each graph input and output.""" + lowered = to_edge_transform_and_lower( + torch.export.export(model, inputs), + partitioner=[DeviceAwarePartitioner("cuda:0")], + ).to_executorch(config) + plan = lowered.executorch_program.execution_plan[0] + planned = lambda indices: [ # noqa: E731 + getattr(plan.values[index].val, "allocation_info", None) is not None + for index in indices + ] + return planned(plan.inputs), planned(plan.outputs) + + # Default planning reserves buffers, so the runtime copies into them despite the skip flags. + planned_inputs, planned_outputs = planned_io( + ExecutorchBackendConfig( + propagate_device_config=skip_copies, + enable_non_cpu_memory_planning=True, + ) + ) + self.assertTrue( + all(planned_inputs), + "graph inputs are expected to be planned by default, which is why the pairing below " + f"is needed, but got {planned_inputs}", + ) + + # Asking planning to leave them alone is what actually removes the copy. + planned_inputs, planned_outputs = planned_io( + ExecutorchBackendConfig( + propagate_device_config=skip_copies, + enable_non_cpu_memory_planning=True, + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), + ) + ) + self.assertFalse( + any(planned_inputs), + f"no graph input should be planned when the caller provides it, got {planned_inputs}", + ) + self.assertFalse( + any(planned_outputs), + f"no graph output should be planned when it stays on the device, got {planned_outputs}", + ) + if __name__ == "__main__": unittest.main()