[PyBindings] Make the torch-free Python bindings safe and complete for CUDA inputs and outputs

.ci/scripts/wheel/test_base.py

@@ -141,6 +141,16 @@ def run_tests(model_tests: List[ModelTest]) -> None:
         print(f"✗ Failed to import _C: {e}")
         raise
 
+    # Type checkers need the stub under the extension's own name, and the marker
+    # that says the package is typed.
+    from executorch.extension.pybindings import _C
+
+    bindings_dir = os.path.dirname(_C.__file__)
+    for name in ("_C.pyi", "py.typed"):
+        path = os.path.join(bindings_dir, name)
+        assert os.path.isfile(path), f"{path} is missing from the wheel"
+    print("✓ The _C type stub and py.typed ship beside _C")
+
     # Why are we doing this envvar shenanigans? Since we build the testers, which
     # uses buck, we cannot run as root. This is a sneaky of getting around that
     # test.

.gitignore

@@ -78,6 +78,7 @@ xcuserdata/
 /devtools/bundled_program/serialize/*.fbs
 /exir/_serialize/program.fbs
 /exir/_serialize/scalar_type.fbs
+/extension/pybindings/_C.pyi
 /include/
 /share/
 /version.py

exir/passes/propagate_device_pass.py

@@ -15,6 +15,7 @@
 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
 
@@ -298,6 +299,14 @@ def _insert_d2h_for_getitem(
                 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(

exir/tests/test_propagate_device_pass.py

@@ -1340,6 +1340,47 @@ def planned_io(config: ExecutorchBackendConfig):
             f"no graph output should be planned when it stays on the device, got {planned_outputs}",
         )
 
+    def test_device_output_copied_to_host_is_planned(self):
+        """A delegate output copied to the host is an intermediate tensor on the
+        device, so memory planning must place it even when the graph outputs are
+        left to the caller. Unplanned, the runtime hands the delegate a null
+        pointer for it."""
+
+        class Model(torch.nn.Module):
+            def forward(self, a, b):
+                return torch.add(a, b)
+
+        lowered = to_edge_transform_and_lower(
+            torch.export.export(Model(), (torch.randn(2, 2), torch.randn(2, 2))),
+            partitioner=[DeviceAwarePartitioner("cuda:0")],
+        ).to_executorch(
+            ExecutorchBackendConfig(
+                propagate_device_config=PropagateDeviceConfig(
+                    skip_h2d_for_method_inputs=True
+                ),
+                enable_non_cpu_memory_planning=True,
+                memory_planning_pass=MemoryPlanningPass(
+                    alloc_graph_input=False, alloc_graph_output=False
+                ),
+            )
+        )
+        plan = lowered.executorch_program.execution_plan[0]
+        device_intermediates = [
+            index
+            for index, value in enumerate(plan.values)
+            if index not in plan.inputs
+            and getattr(value.val, "extra_tensor_info", None) is not None
+            and value.val.extra_tensor_info.device_type == DeviceType.CUDA
+        ]
+        self.assertTrue(device_intermediates)
+        for index in device_intermediates:
+            self.assertIsNotNone(
+                plan.values[index].val.allocation_info,
+                f"device tensor {index} is not planned",
+            )
+        for index in plan.outputs:
+            self.assertIsNone(plan.values[index].val.allocation_info)
+
 
 if __name__ == "__main__":
     unittest.main()

extension/pybindings/README.md

@@ -16,7 +16,9 @@ arrays or other dense objects implementing Python's buffer protocol, including
 CUDA `torch.Tensor` inputs are inspected through their public Python attributes
 instead of being cast to `at::Tensor`. DLPack producers, including CuPy, can
 provide CPU or CUDA allocations directly; their device indices are preserved.
-Tensor protocols can be mixed within the same call.
+A DLPack capsule, such as the one `array.__dlpack__()` returns, can be passed
+directly too; an input consumes it, so it cannot be passed again. Tensor
+protocols can be mixed within the same call.
 
 Tensor inputs must have dense storage. In the portable bindings, an input laid
 out in a different dimension order than the exported method, for example a
@@ -35,6 +37,9 @@ input in place, and only such an input sees the method's own in-place updates.
 CUDA inputs require a method
 exported with CUDA inputs and `MemoryPlanningPass(alloc_graph_input=False)` so
 execution can share the caller's device allocation.
+`method_meta().input_tensor_meta(i).device()` tells where the method expects
+input `i`, as `"cpu"` or `"cuda:N"`, and `output_tensor_meta(i).device()` does
+the same for its outputs.
 
 When PyTorch is available, importing `portable_lib` imports it and portable
 bindings return `torch.Tensor` outputs. When it cannot be imported, the default
@@ -45,7 +50,13 @@ on the output's original device. With `clone_outputs=False`, results share
 runtime or input storage and later execution may overwrite that storage.
 
 A CUDA output that fills memory the bindings allocated for it is returned
-without a copy when `clone_outputs=True`. The result keeps that memory, so a
+without a copy when `clone_outputs=True`. That needs a method exported with
+`MemoryPlanningPass(alloc_graph_output=False)`, so the method does not plan its
+outputs into its own memory;
+`method_meta().output_tensor_meta(i).is_memory_planned()` tells whether output
+`i` is planned. A planned CUDA output read with
+`clone_outputs=True` costs a new allocation, a device-to-device copy, and a
+wait for the device. The result keeps that memory, so a
 method writes its next execution into a second buffer, and the two take turns;
 a module allocates new memory each call. Within one execution, reading the
 outputs again with `clone_outputs=True` returns the same object, so an in-place
@@ -57,6 +68,12 @@ through `torch.multiprocessing`. Call `.clone()` for a normal torch tensor.
 `dtype_name` names a result's dtype, including dtypes NumPy lacks, such as
 `"bfloat16"`.
 
+`Method.get_attribute()` returns a view of mutable state. A constant can sit
+in read-only memory, so it comes back read-only: NumPy and the buffer protocol
+see it as read-only, and a DLPack export copies it, unless the consumer asks
+for no copy through DLPack 1.0, which marks it read-only. In torch output mode
+a constant comes back as a copy, since a torch tensor cannot be read-only.
+
 Calls on a shared method or module are serialized. A sequence of separate
 `set_inputs`, `execute`, and `get_outputs` calls on a shared method needs a lock
 held by the caller around the whole sequence: `set_inputs` hides the outputs
@@ -87,7 +104,10 @@ result = np.asarray(output)
 
 CPU and CUDA results implement `__dlpack__` and `__dlpack_device__`. CUDA results
 remain on CUDA and can be passed to another DLPack consumer or back into
-ExecuTorch without moving the data to the CPU.
+ExecuTorch without moving the data to the CPU. `__dlpack__` takes the DLPack
+1.0 keywords `max_version`, `dl_device`, and `copy`, so a call such as
+`np.from_dlpack(result, copy=False)` works. A consumer that passes `stream=-1`
+orders the work itself, so that export skips the device wait.
 
 Without torch or another DLPack library, `to_host()` copies a result to CPU
 memory, and `ExecuTorchResult.from_host()` copies a CPU buffer to a CUDA device
@@ -105,18 +125,28 @@ Copying CUDA memory requires CUDA-enabled bindings: `from_host()`, `to_host()`
 of a CUDA result, and a CUDA `ExecuTorchResult` with `clone_outputs=True`.
 Bindings built without CUDA support synchronize through Python Torch instead,
 so while torch is loaded they still read CUDA inputs, torch tensors or DLPack,
-and return borrowed CUDA results. Execution is synchronous at the Python
-boundary: a call returns after its device work has finished, so its outputs are
-ready on any stream. Each call waits once per device after it runs, rather than
-once per output. Before running, it waits once per device for any CUDA input
-that another library may still be writing: a torch CUDA tensor, a CUDA DLPack
-input, or a result of these bindings passed back in, since that result can share
-memory a caller handed to another library. The same wait covers the memory of a
-CUDA output that a method returned without a copy: once that result is dropped,
-the method writes the memory again, and another library may still be reading
-it. A torch CUDA input in another layout also waits before torch copies it into
-the exported layout. Reading a CUDA result through `__dlpack__` or `to_host()`
-waits once as well. The synchronization has no C++ ATen dependency.
+and return borrowed CUDA results.
+
+Execution is synchronous at the Python boundary: a call returns after its device
+work has finished, so its outputs are ready on any stream. Every wait is for the
+whole device (`cudaDeviceSynchronize`), not for one stream, so it also waits for
+unrelated work queued on that device, for example by torch. Before it runs, a
+call waits once for each device that holds a non-empty CUDA input or output.
+Another library may still be writing such an input: a torch CUDA tensor, a CUDA
+DLPack input, or a result of these bindings passed back in, which can share
+memory a caller handed to that library. It may also still be reading the memory
+of an output from an earlier call, which this call writes again. After it runs,
+a call waits once for each device that holds a CUDA input or output, since the
+method may have updated an input in place. Each cloned CUDA
+output waits once more for its copy. A torch CUDA input in another layout also
+waits before torch copies it into the exported layout. Reading a CUDA result
+through `__dlpack__` or `to_host()` waits once as well. In CUDA-enabled
+bindings a wait holds the GIL, so other Python threads pause until it ends, and
+a queued CUDA host callback that needs the GIL makes the wait hang. Bindings
+built without CUDA wait through `torch.cuda.synchronize`, which releases the
+GIL. On an idle device a wait is short:
+`cudaDeviceSynchronize` took a median of 1.5 microseconds on an A100. The
+synchronization has no C++ ATen dependency.
 
 For bfloat16 inputs, raw byte buffers use the shape declared by the exported
 method. Bfloat16 results expose flat bytes through the buffer protocol because