[PyBindings] Wait for the device instead of per-result flags, and fix bfloat16 result layout and input errors (#23999)

Follow-up to the Python bindings change that added torch-free execution
and device memory support.

## What was wrong

A CUDA result borrowed from a method, through
`get_outputs(clone_outputs=False)` or a method that returns its own
input, shares the method's memory. The bindings tracked "might another
library still be writing this?" with a flag on each result object, and
every borrowed result started with a fresh flag. So when another library
wrote the memory through one result, and a second view of the same
memory was passed back in, the call could read the old data.

A bfloat16 result on the host, passed back in, was read in the method's
layout instead of its own. A result in another dim order came back
permuted, with no error.

## What this changes

The flags are replaced by device waits. Before a call runs, it waits
once per device for its CUDA inputs, including results of these
bindings. A method also waits for the device of an output buffer it
writes again, since another library may still be reading a dropped
result there. After a call, it waits once per output device.
`__dlpack__` and `to_host()` on a CUDA result always wait. In bindings
built without CUDA, `__dlpack__` now waits through torch, as calls
already do, instead of failing.

A result of these bindings passed back in always takes the DLPack path,
which keeps its dtype, shape and strides. This fixes the bfloat16 layout
bug, and a bfloat16 result now keeps its own shape instead of taking the
method's.

A failed device allocation no longer turns off the no-copy output
return. A `clone_outputs=False` read is now its own borrowed view.

An error raised while reading an input keeps its Python type. A Python
exception gets the input and method as a note, and C++ errors name them
in the message.

The torch-free CI job now also builds a channels-last bfloat16 program,
so the layout test runs there.

## Tested

On Linux x86_64 with H100 GPUs, built the bindings with CUDA, without
CUDA, and as the standalone torch-free module. Ran
`test_pybindings_cuda_dlpack.py`, `test_pybindings_no_aten.py`,
`test_pybindings.py` and `runtime/test/test_runtime.py` in both output
modes. Every failure there also fails without this change: tests written
for the other output mode, two quantized-op tests whose kernels are not
registered here, and three subprocess tests that need the CUDA runtime
library outside torch. Each new or changed test fails with its fix
removed.

Author: shoumikhin

.github/workflows/pull.yml

@@ -182,10 +182,15 @@ jobs:
         from executorch.extension.pybindings.test.make_test import (
             create_program,
             ModuleAddBFloat16,
+            ModuleAddBFloat16ToChannelsLast,
         )
 
         program, _ = create_program(ModuleAddBFloat16())
         Path("cmake-out-no-aten/add_bfloat16.pte").write_bytes(bytes(program.buffer))
+        program, _ = create_program(ModuleAddBFloat16ToChannelsLast())
+        Path("cmake-out-no-aten/add_bfloat16_channels_last.pte").write_bytes(
+            bytes(program.buffer)
+        )
         PY
         EXECUTORCH_PYBIND_TEST_PTE=cmake-out-no-aten/add.pte \
           PYTHONPATH=cmake-out-no-aten:.. \
@@ -205,6 +210,7 @@ jobs:
         EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch \
           EXECUTORCH_PYBIND_TEST_PTE=cmake-out-no-aten/add.pte \
           EXECUTORCH_PYBIND_TEST_BFLOAT16_PTE=cmake-out-no-aten/add_bfloat16.pte \
+          EXECUTORCH_PYBIND_TEST_BFLOAT16_CHANNELS_LAST_PTE=cmake-out-no-aten/add_bfloat16_channels_last.pte \
           PYTHONPATH=cmake-out-no-aten:.. \
           python extension/pybindings/test/test_pybindings_no_aten.py
         if ldd cmake-out-no-aten/_C*.so | grep -Eiq 'torch|aten|c10'; then

extension/pybindings/README.md

@@ -47,10 +47,12 @@ 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
 method writes its next execution into a second buffer, and the two take turns;
-a module allocates new memory each call. Reading the outputs again before the
-next execution returns the same object, so an in-place change to one read
-shows in the other. In torch output mode such a result is a view of memory
-torch does not own, so it cannot grow with `resize_()`.
+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
+change to one read shows in the other; a `clone_outputs=False` read is a
+separate borrowed view. In torch output mode such a result is a view of memory
+torch does not own, so it cannot grow with `resize_()` and cannot be shared
+through `torch.multiprocessing`. Call `.clone()` for a normal torch tensor.
 
 `dtype_name` names a result's dtype, including dtypes NumPy lacks, such as
 `"bfloat16"`.
@@ -99,20 +101,29 @@ output = method([device_input])[0]
 result = np.asarray(output.to_host())
 ```
 
-CUDA DLPack inputs and `ExecuTorchResult` outputs require CUDA-enabled bindings.
-Torch CUDA inputs and outputs synchronize through Python Torch, including in
-CPU builds of the bindings. 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 for the device once, rather than once per tensor.
-Before running, it also waits once if an input may still be written on another
-library's stream: a torch CUDA tensor, a DLPack input from another library, or
-an `ExecuTorchResult` whose memory was handed out through `__dlpack__`. The
-synchronization has no C++ ATen dependency.
+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.
 
 For bfloat16 inputs, raw byte buffers use the shape declared by the exported
 method. Bfloat16 results expose flat bytes through the buffer protocol because
 NumPy and PEP 3118 do not define that numeric dtype. DLPack consumers receive the
-original bfloat16 dtype and tensor shape.
+original bfloat16 dtype and tensor shape. A bfloat16 result passed back in is
+read through DLPack too, so it keeps its own shape and dim order instead of
+taking the method's.
 
 ```python
 import cupy as cp