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

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. Whether
another library might still be writing that memory was tracked 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 data from before
that write.

The flags are gone. Before a call runs, it now waits once for each device that
holds one of its CUDA inputs, including results of these bindings. A method
also waits for the device of each output buffer it writes again, because
another library may still be reading a dropped result that used it. After the
call runs, it waits once for each device that holds an output, as before.
Handing a CUDA result to another library through __dlpack__ or to_host() always
waits too. 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 now always takes the DLPack path,
which keeps its real dtype, shape and strides. A bfloat16 result used to take
the buffer path, whose flat bytes were read in the method's layout, so a result
in another dim order came back permuted with no error. A bfloat16 result now
also keeps its own shape instead of taking the method's.

A failed device allocation no longer turns off the no-copy output return for
the rest of a method's life, because the new buffer is allocated into a local
first. A clone_outputs=False read is now its own borrowed view, not the object
a default read handed over.

An error raised while reading an input keeps its Python type. A Python
exception gets the input and method as a note instead of being rebuilt, which
failed for types such as UnicodeDecodeError. C++ runtime errors, the device
checks and the unsupported-type error name the input and method in their
message, and std::bad_alloc still becomes MemoryError.

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