Document and test the memory planning a device-activation export needs

exir/passes/propagate_device_config.py

@@ -25,17 +25,43 @@
 @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
 
     # 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.
+    #
+    # 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
 

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()