Skip to content

Commit a4e65bf

Browse files
authored
[PyBindings] Make the torch-free Python bindings safe and complete for CUDA inputs and outputs (#24016)
Follow-up to #23999: makes the Python bindings work well without PyTorch for programs with CUDA inputs and outputs. Each commit fixes one problem and adds a test. ## What is wrong today, and the fix 1. A NumPy array passed for a GPU input is read as GPU memory, which can crash. It is now rejected with a clear message. 2. A run can overwrite a planned output another library still reads. Each `Method` or `Module` run now waits, before and after, for every device that holds one of its tensors. Input errors name the input. 3. A complex32 DLPack input is rejected. It is now read. 4. `ExecuTorchResult.__dlpack__` rejects the DLPack 1.0 keywords, so `np.from_dlpack(result, copy=False)` fails. It now takes `max_version`, `dl_device` and `copy`, can return versioned capsules, and skips the device wait for `stream=-1`. 5. A raw DLPack capsule is not accepted as an input. Legacy and versioned capsules now work, and producers get asked for them. 6. `Method.get_attribute()` returns constants as writable memory, which may be read-only pages. Constants now come back read-only, DLPack exports of them are copies, and torch output mode gets a copy. 7. A method from a temporary `Program` with ETDump on uses a freed event tracer. The tracer now lives as long as any method. 8. A caller cannot ask which device a method expects. `TensorInfo.device()` and `ExecuTorchResult.device` now return `"cpu"` or `"cuda:N"`. 9. The installed type stub is never found and has errors. It now ships fixed, as `_C.pyi` with `py.typed`. 10. A delegate output copied to an unplanned host output (`alloc_graph_output=False`) crashes. The device copy pass now updates the graph output, so memory planning places it. 11. `executorch.runtime`'s `Method.execute` now takes `clone_outputs`, so it can skip the copy. ## Testing Ran on Linux x86_64 with an H100: CUDA tests in both output modes, CPU tests with and without torch. Each new test failed on a build without its fix: - GPU: planned output wait (Method and Module), wait on a result another library wrote, DLPack keywords and `stream=-1`, GPU inputs with unplanned host outputs. - CPU: host buffer check, error messages, complex32, DLPack keywords, raw capsules, read-only constants, tracer lifetime (crashed before), device metadata, device copy planning, `Method.execute(clone_outputs=False)`. - The wheel check fails without `_C.pyi` or `py.typed`; mypy resolves the stub from a built package. Not tested: the wait after a run for inputs updated in place.
1 parent b032012 commit a4e65bf

21 files changed

Lines changed: 1183 additions & 310 deletions

‎.ci/scripts/wheel/test_base.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,16 @@ def run_tests(model_tests: List[ModelTest]) -> None:
141141
print(f"✗ Failed to import _C: {e}")
142142
raise
143143

144+
# Type checkers need the stub under the extension's own name, and the marker
145+
# that says the package is typed.
146+
from executorch.extension.pybindings import _C
147+
148+
bindings_dir = os.path.dirname(_C.__file__)
149+
for name in ("_C.pyi", "py.typed"):
150+
path = os.path.join(bindings_dir, name)
151+
assert os.path.isfile(path), f"{path} is missing from the wheel"
152+
print("✓ The _C type stub and py.typed ship beside _C")
153+
144154
# Why are we doing this envvar shenanigans? Since we build the testers, which
145155
# uses buck, we cannot run as root. This is a sneaky of getting around that
146156
# test.

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ xcuserdata/
7878
/devtools/bundled_program/serialize/*.fbs
7979
/exir/_serialize/program.fbs
8080
/exir/_serialize/scalar_type.fbs
81+
/extension/pybindings/_C.pyi
8182
/include/
8283
/share/
8384
/version.py

‎exir/passes/propagate_device_pass.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import executorch.exir.passes._device_copy_ops_registry # noqa: F401
1616
import executorch.exir.schema as schema
1717
import torch
18+
import torch.utils._pytree as pytree
1819
from executorch.exir.delegate import executorch_call_delegate
1920
from executorch.exir.lowered_backend_module import LoweredBackendModule
2021

@@ -298,6 +299,14 @@ def _insert_d2h_for_getitem(
298299
d2h_node,
299300
delete_user_cb=lambda user, _d2h=d2h_node: user != _d2h,
300301
)
302+
# Memory planning finds the graph outputs through the output node's
303+
# spec, so it has to name the copy now rather than the device tensor.
304+
for user in d2h_node.users:
305+
if user.op == "output" and "spec" in user.meta:
306+
user.meta["spec"] = pytree.tree_map(
307+
lambda value: d2h_spec if value is spec else value,
308+
user.meta["spec"],
309+
)
301310
return True
302311

303312
def _elide_redundant_same_device_copies(

‎exir/tests/test_propagate_device_pass.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1340,6 +1340,47 @@ def planned_io(config: ExecutorchBackendConfig):
13401340
f"no graph output should be planned when it stays on the device, got {planned_outputs}",
13411341
)
13421342

1343+
def test_device_output_copied_to_host_is_planned(self):
1344+
"""A delegate output copied to the host is an intermediate tensor on the
1345+
device, so memory planning must place it even when the graph outputs are
1346+
left to the caller. Unplanned, the runtime hands the delegate a null
1347+
pointer for it."""
1348+
1349+
class Model(torch.nn.Module):
1350+
def forward(self, a, b):
1351+
return torch.add(a, b)
1352+
1353+
lowered = to_edge_transform_and_lower(
1354+
torch.export.export(Model(), (torch.randn(2, 2), torch.randn(2, 2))),
1355+
partitioner=[DeviceAwarePartitioner("cuda:0")],
1356+
).to_executorch(
1357+
ExecutorchBackendConfig(
1358+
propagate_device_config=PropagateDeviceConfig(
1359+
skip_h2d_for_method_inputs=True
1360+
),
1361+
enable_non_cpu_memory_planning=True,
1362+
memory_planning_pass=MemoryPlanningPass(
1363+
alloc_graph_input=False, alloc_graph_output=False
1364+
),
1365+
)
1366+
)
1367+
plan = lowered.executorch_program.execution_plan[0]
1368+
device_intermediates = [
1369+
index
1370+
for index, value in enumerate(plan.values)
1371+
if index not in plan.inputs
1372+
and getattr(value.val, "extra_tensor_info", None) is not None
1373+
and value.val.extra_tensor_info.device_type == DeviceType.CUDA
1374+
]
1375+
self.assertTrue(device_intermediates)
1376+
for index in device_intermediates:
1377+
self.assertIsNotNone(
1378+
plan.values[index].val.allocation_info,
1379+
f"device tensor {index} is not planned",
1380+
)
1381+
for index in plan.outputs:
1382+
self.assertIsNone(plan.values[index].val.allocation_info)
1383+
13431384

13441385
if __name__ == "__main__":
13451386
unittest.main()

‎extension/pybindings/README.md‎

Lines changed: 45 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@ arrays or other dense objects implementing Python's buffer protocol, including
1616
CUDA `torch.Tensor` inputs are inspected through their public Python attributes
1717
instead of being cast to `at::Tensor`. DLPack producers, including CuPy, can
1818
provide CPU or CUDA allocations directly; their device indices are preserved.
19-
Tensor protocols can be mixed within the same call.
19+
A DLPack capsule, such as the one `array.__dlpack__()` returns, can be passed
20+
directly too; an input consumes it, so it cannot be passed again. Tensor
21+
protocols can be mixed within the same call.
2022

2123
Tensor inputs must have dense storage. In the portable bindings, an input laid
2224
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.
3537
CUDA inputs require a method
3638
exported with CUDA inputs and `MemoryPlanningPass(alloc_graph_input=False)` so
3739
execution can share the caller's device allocation.
40+
`method_meta().input_tensor_meta(i).device()` tells where the method expects
41+
input `i`, as `"cpu"` or `"cuda:N"`, and `output_tensor_meta(i).device()` does
42+
the same for its outputs.
3843

3944
When PyTorch is available, importing `portable_lib` imports it and portable
4045
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
4550
runtime or input storage and later execution may overwrite that storage.
4651

4752
A CUDA output that fills memory the bindings allocated for it is returned
48-
without a copy when `clone_outputs=True`. The result keeps that memory, so a
53+
without a copy when `clone_outputs=True`. That needs a method exported with
54+
`MemoryPlanningPass(alloc_graph_output=False)`, so the method does not plan its
55+
outputs into its own memory;
56+
`method_meta().output_tensor_meta(i).is_memory_planned()` tells whether output
57+
`i` is planned. A planned CUDA output read with
58+
`clone_outputs=True` costs a new allocation, a device-to-device copy, and a
59+
wait for the device. The result keeps that memory, so a
4960
method writes its next execution into a second buffer, and the two take turns;
5061
a module allocates new memory each call. Within one execution, reading the
5162
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.
5768
`dtype_name` names a result's dtype, including dtypes NumPy lacks, such as
5869
`"bfloat16"`.
5970

71+
`Method.get_attribute()` returns a view of mutable state. A constant can sit
72+
in read-only memory, so it comes back read-only: NumPy and the buffer protocol
73+
see it as read-only, and a DLPack export copies it, unless the consumer asks
74+
for no copy through DLPack 1.0, which marks it read-only. In torch output mode
75+
a constant comes back as a copy, since a torch tensor cannot be read-only.
76+
6077
Calls on a shared method or module are serialized. A sequence of separate
6178
`set_inputs`, `execute`, and `get_outputs` calls on a shared method needs a lock
6279
held by the caller around the whole sequence: `set_inputs` hides the outputs
@@ -87,7 +104,10 @@ result = np.asarray(output)
87104

88105
CPU and CUDA results implement `__dlpack__` and `__dlpack_device__`. CUDA results
89106
remain on CUDA and can be passed to another DLPack consumer or back into
90-
ExecuTorch without moving the data to the CPU.
107+
ExecuTorch without moving the data to the CPU. `__dlpack__` takes the DLPack
108+
1.0 keywords `max_version`, `dl_device`, and `copy`, so a call such as
109+
`np.from_dlpack(result, copy=False)` works. A consumer that passes `stream=-1`
110+
orders the work itself, so that export skips the device wait.
91111

92112
Without torch or another DLPack library, `to_host()` copies a result to CPU
93113
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()`
105125
of a CUDA result, and a CUDA `ExecuTorchResult` with `clone_outputs=True`.
106126
Bindings built without CUDA support synchronize through Python Torch instead,
107127
so while torch is loaded they still read CUDA inputs, torch tensors or DLPack,
108-
and return borrowed CUDA results. Execution is synchronous at the Python
109-
boundary: a call returns after its device work has finished, so its outputs are
110-
ready on any stream. Each call waits once per device after it runs, rather than
111-
once per output. Before running, it waits once per device for any CUDA input
112-
that another library may still be writing: a torch CUDA tensor, a CUDA DLPack
113-
input, or a result of these bindings passed back in, since that result can share
114-
memory a caller handed to another library. The same wait covers the memory of a
115-
CUDA output that a method returned without a copy: once that result is dropped,
116-
the method writes the memory again, and another library may still be reading
117-
it. A torch CUDA input in another layout also waits before torch copies it into
118-
the exported layout. Reading a CUDA result through `__dlpack__` or `to_host()`
119-
waits once as well. The synchronization has no C++ ATen dependency.
128+
and return borrowed CUDA results.
129+
130+
Execution is synchronous at the Python boundary: a call returns after its device
131+
work has finished, so its outputs are ready on any stream. Every wait is for the
132+
whole device (`cudaDeviceSynchronize`), not for one stream, so it also waits for
133+
unrelated work queued on that device, for example by torch. Before it runs, a
134+
call waits once for each device that holds a non-empty CUDA input or output.
135+
Another library may still be writing such an input: a torch CUDA tensor, a CUDA
136+
DLPack input, or a result of these bindings passed back in, which can share
137+
memory a caller handed to that library. It may also still be reading the memory
138+
of an output from an earlier call, which this call writes again. After it runs,
139+
a call waits once for each device that holds a CUDA input or output, since the
140+
method may have updated an input in place. Each cloned CUDA
141+
output waits once more for its copy. A torch CUDA input in another layout also
142+
waits before torch copies it into the exported layout. Reading a CUDA result
143+
through `__dlpack__` or `to_host()` waits once as well. In CUDA-enabled
144+
bindings a wait holds the GIL, so other Python threads pause until it ends, and
145+
a queued CUDA host callback that needs the GIL makes the wait hang. Bindings
146+
built without CUDA wait through `torch.cuda.synchronize`, which releases the
147+
GIL. On an idle device a wait is short:
148+
`cudaDeviceSynchronize` took a median of 1.5 microseconds on an A100. The
149+
synchronization has no C++ ATen dependency.
120150

121151
For bfloat16 inputs, raw byte buffers use the shape declared by the exported
122152
method. Bfloat16 results expose flat bytes through the buffer protocol because

‎extension/pybindings/py.typed‎

Whitespace-only changes.

0 commit comments

Comments
 (0)