From 9868f97ab66247a0acd71dd763e6dac07902be86 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 19 Oct 2026 22:54:14 -0700 Subject: [PATCH] [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. --- .github/workflows/pull.yml | 6 + extension/pybindings/README.md | 39 ++-- extension/pybindings/pybindings.cpp | 217 +++++++----------- extension/pybindings/pybindings_cuda.h | 4 +- .../pybindings/pybindings_executorch_result.h | 80 ++++--- extension/pybindings/test/make_test.py | 11 + extension/pybindings/test/test_pybindings.py | 2 + .../test/test_pybindings_cuda_dlpack.py | 139 +++++++---- .../test/test_pybindings_no_aten.py | 50 ++++ 9 files changed, 321 insertions(+), 227 deletions(-) diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index a595a13afa2..f4efbf35b64 100644 --- a/.github/workflows/pull.yml +++ b/.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 diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 514d5709d03..8b031ff734b 100644 --- a/extension/pybindings/README.md +++ b/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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 016fbe43247..635a4393184 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -126,21 +126,6 @@ namespace pybindings { namespace { -PyObject* loaded_torch_module() { - // Retain the module for the interpreter's lifetime. This avoids repeated - // sys.modules lookups without creating a static py::object whose destructor - // could run after Python has finalized. - static PyObject* torch_module = nullptr; - if (torch_module == nullptr) { - PyObject* loaded = PyDict_GetItemString(PyImport_GetModuleDict(), "torch"); - if (loaded != nullptr) { - Py_INCREF(loaded); - torch_module = loaded; - } - } - return torch_module; -} - #ifndef USE_ATEN_LIB constexpr char kTensorOutputEnvironmentVariable[] = "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"; @@ -313,11 +298,7 @@ class BufferTensor final { info_(buffer_.request()) { check_rank(); const auto input_meta = method_meta.input_tensor_meta(input_index); - THROW_IF_ERROR( - input_meta.error(), - "Input %zu of method '%s' is not a tensor input", - input_index, - method_meta.name()); + THROW_IF_ERROR(input_meta.error(), "not a tensor input"); const bool raw_bfloat16 = input_meta->scalar_type() == executorch::aten::ScalarType::BFloat16 && info_.ndim == 1 && @@ -329,8 +310,8 @@ class BufferTensor final { } if (static_cast(info_.size) != input_meta->nbytes()) { throw py::value_error( - "Raw bfloat16 input " + std::to_string(input_index) + " has " + - std::to_string(info_.size) + " bytes, but the method expects " + + "raw bfloat16 input has " + std::to_string(info_.size) + + " bytes, but the method expects " + std::to_string(input_meta->nbytes())); } scalar_type_ = executorch::aten::ScalarType::BFloat16; @@ -658,24 +639,33 @@ bool is_torch_tensor(const py::handle& value) { return result == 1; } -// Runs `read`, which builds one input from a Python object, and starts the -// message of any error it raises with the input and its method, which the -// readers cannot see. The error keeps its Python type. +// Runs `read`, which builds one input from a Python object, and names the +// input and its method in its errors, which the readers cannot see. A Python +// error keeps its exact type and gets the name as a note; the C++ errors +// caught below get it in their message. Others, such as std::bad_alloc, pass +// through unnamed, so pybind11 maps them to their own Python type. template auto read_input(const MethodMeta& method_meta, size_t index, Read&& read) { - const auto named = [&](const char* message) { + const auto name = [&] { return "Input " + std::to_string(index) + " of method '" + - method_meta.name() + "': " + message; + method_meta.name() + "'"; + }; + const auto named = [&name](const char* message) { + return name() + ": " + message; }; try { return read(); } catch (py::error_already_set& error) { - // Chained to the original, so its traceback stays available. - const auto type = error.type(); - const auto message = - named(py::str(error.value()).cast().c_str()); - py::raise_from(error, type.ptr(), message.c_str()); - throw py::error_already_set(); + // Rebuilding the exception as type(message) would fail for any exception + // whose constructor is not a single string. + if (error.value()) { + try { + error.value().attr("add_note")(name()); + } catch (...) { + // add_note is best effort; never hide the original error. + } + } + throw; } catch (const py::buffer_error& error) { throw py::buffer_error(named(error.what())); } catch (const py::value_error& error) { @@ -684,7 +674,7 @@ auto read_input(const MethodMeta& method_meta, size_t index, Read&& read) { throw py::type_error(named(error.what())); } catch (const std::invalid_argument& error) { throw py::value_error(named(error.what())); - } catch (const std::exception& error) { + } catch (const std::runtime_error& error) { throw std::runtime_error(named(error.what())); } } @@ -730,12 +720,6 @@ void validate_buffer_input( } #ifndef USE_ATEN_LIB -void synchronize_torch_cuda_device(int index) { - py::reinterpret_borrow(loaded_torch_module()) - .attr("cuda") - .attr("synchronize")(index); -} - /** A torch.Tensor input read through its public Python API, copied only to * change its layout. */ class TorchTensorView final { @@ -872,7 +856,8 @@ class TorchTensorView final { : "cuda:" + std::to_string(device.index()); }; throw py::value_error( - "Input " + std::to_string(index) + " is on device " + name(device_) + + "Input " + std::to_string(index) + " of method '" + + method.method_meta().name() + "' is on device " + name(device_) + ", but the method expects " + name(expected)); } } @@ -995,16 +980,10 @@ class DLPackTensorView final { advertised_device_( owner_.attr("__dlpack_device__")().cast>()) { validate_advertised_device(); - // A result of these bindings is written either by an execution, which - // finished before it returned, or by a copy on the legacy default stream, - // which the delegates' streams wait for. So it needs no wait unless it - // was handed to another library. Anything else may still be written on - // its producer's streams until the call waits for it. + // A result of these bindings gives its capsule without the wait in + // __dlpack__, since the call waits for every CUDA input before it runs. if (py::isinstance(owner_)) { - own_result_ = owner_.cast>(); - } - if (own_result_) { - capsule_ = own_result_->capsule(); + capsule_ = owner_.cast>()->capsule(); } else { capsule_ = advertised_device_.first == dlpack::CUDA ? cuda_dlpack_capsule(owner_, advertised_device_.second) @@ -1061,14 +1040,6 @@ class DLPackTensorView final { return nbytes_; } - // Whether the producer may still be writing the tensor on its own streams. - // Asked when the call runs, since a result set as an input can be handed to - // another library before that. - bool may_have_pending_writes() const { - return advertised_device_.first == dlpack::CUDA && - !(own_result_ && !own_result_->handed_out()); - } - // Describes the tensor in the exported dimension order. A CPU tensor dense // in another order is copied into that order, since reading it with the // exported strides would permute its values. A device tensor would need a @@ -1103,8 +1074,9 @@ class DLPackTensorView final { if (nbytes_ != 0 && device_ != method.get_input(index).toTensor().device()) { throw py::value_error( - "DLPack input " + std::to_string(index) + - " device does not match the method input device"); + "Input " + std::to_string(index) + " of method '" + + method.method_meta().name() + + "': DLPack input device does not match the method input device"); } } @@ -1298,7 +1270,6 @@ class DLPackTensorView final { py::object owner_; std::pair advertised_device_; - std::shared_ptr own_result_; size_t nbytes_ = 0; py::object capsule_; dlpack::ManagedTensor* managed_ = nullptr; @@ -1496,6 +1467,12 @@ bool prefer_dlpack_input(const py::handle& value) { if (!PyObject_CheckBuffer(value.ptr())) { return true; } + // A result of these bindings always takes the DLPack path, which keeps its + // real dtype, shape, and strides. A bfloat16 result exports flat bytes + // instead, which the buffer path would read in the method's layout. + if (py::isinstance(value)) { + return true; + } // NumPy and other host buffers keep the buffer path, including its // read-only copy behavior. CUDA results expose only their DLPack storage. if (py::hasattr(value, "__dlpack_device__")) { @@ -1530,19 +1507,6 @@ std::shared_ptr allocate_device_output( return std::make_shared(std::move(buffer.get())); } -// Waits for every queued operation on a device. Bindings built without CUDA -// support have no CUDA runtime of their own, so they wait through torch's, -// which is loaded whenever such bindings return CUDA tensors. -void synchronize_device(int index) { -#ifndef EXECUTORCH_PYBINDINGS_CUDA - if (loaded_torch_module() != nullptr) { - synchronize_torch_cuda_device(index); - return; - } -#endif - synchronize_cuda_device(index); -} - // The Python boundary is synchronous: an execution returns only after its // device work has finished, so its outputs are ready wherever they are read. // One wait per execution covers every output, instead of one per output. @@ -1565,13 +1529,15 @@ void synchronize_device_outputs(const Method& method) { } } -// Device inputs from another library may still be written on that library's -// streams, torch tensors included. One wait per call covers every such input, -// instead of one per input. An ExecuTorchResult these bindings returned needs -// no wait unless its memory was handed to another library since. -void synchronize_device_inputs( +// Another library may still be writing a CUDA input on its own streams: a +// torch tensor, a DLPack tensor, or a result of these bindings, which can share +// memory a caller gave to that library. It may also still be reading an output +// buffer that a dropped result used, which the run writes again. One wait per +// device covers all of them. +void synchronize_devices_before_run( const std::vector>& dlpack_inputs, - const std::vector>& torch_inputs) { + const std::vector>& torch_inputs, + const std::vector& output_devices = {}) { std::vector devices; const auto add = [&devices](executorch::aten::Device device) { if (!device.is_cpu() && @@ -1581,7 +1547,7 @@ void synchronize_device_inputs( } }; for (const auto& input : dlpack_inputs) { - if (input->nbytes() != 0 && input->may_have_pending_writes()) { + if (input->nbytes() != 0) { add(input->device()); } } @@ -1590,6 +1556,9 @@ void synchronize_device_inputs( add(input->device()); } } + for (const auto device : output_devices) { + add(device); + } for (const auto index : devices) { synchronize_device(index); } @@ -2558,7 +2527,8 @@ struct PyModule final { #endif } else { throw std::runtime_error( - "Unsupported python type " + type_str + + "Input " + std::to_string(i) + " of method '" + method_meta.name() + + "': unsupported python type " + type_str + ". Ensure that inputs are passed as a flat list of tensors."); } } @@ -2629,7 +2599,7 @@ struct PyModule final { dlpack_inputs.push_back(owner.dlpack); } } - synchronize_device_inputs(dlpack_inputs, torch_inputs); + synchronize_devices_before_run(dlpack_inputs, torch_inputs); #endif auto outputs = module_->execute(method_name); THROW_IF_ERROR( @@ -3388,7 +3358,9 @@ struct PyMethod final { #endif } else { throw std::runtime_error( - "Unsupported python type " + type_str + + "Input " + std::to_string(i) + " of method '" + + storage_->method->method_meta().name() + + "': unsupported python type " + type_str + ". Ensure that inputs are passed as a flat list of tensors."); } } @@ -3485,7 +3457,15 @@ struct PyMethod final { #endif setup_output_storage(*storage_->method, output_storage_spans); #ifndef USE_ATEN_LIB - synchronize_device_inputs(dlpack_inputs_, torch_inputs_); + std::vector output_devices; + for (size_t i = 0; i < storage_->device_output_storages.size(); ++i) { + if (storage_->device_output_storages[i]) { + output_devices.push_back( + storage_->method->get_output(i).toTensor().device()); + } + } + synchronize_devices_before_run( + dlpack_inputs_, torch_inputs_, output_devices); #endif Error execute_status = storage_->method->execute(); THROW_IF_ERROR( @@ -3598,21 +3578,13 @@ struct PyMethod final { std::vector handed_outputs_; // Per device output buffer, alive while such a result still uses it. std::vector> handed_buffers_; - // Per device output buffer, whether a result was handed it since the last - // execution. - std::vector handed_slots_; // Per device output, the buffer a still-alive earlier result uses, kept so // it can be reused once that result is dropped. struct SpareDeviceOutput final { std::shared_ptr buffer; std::weak_ptr handed; - std::shared_ptr handed_out; }; std::vector spare_device_outputs_; - // Per device output buffer, whether another library may have read it, so - // reusing it must wait for that library's queued work. Torch outputs always - // may; an ExecuTorchResult only once exported through DLPack. - std::vector> handed_out_; #ifndef USE_ATEN_LIB struct BorrowedOutputOwner final { @@ -3690,10 +3662,7 @@ struct PyMethod final { std::make_shared>(buffer); const std::shared_ptr memory(holder, buffer->data()); py::object result; - const auto torch_output = - portable_tensor_output() == PortableTensorOutput::Torch; - handed_out_[*slot] = std::make_shared(torch_output); - if (torch_output) { + if (portable_tensor_output() == PortableTensorOutput::Torch) { auto* retained = new std::shared_ptr(memory); result = portable_tensor_result( tensor, @@ -3702,11 +3671,9 @@ struct PyMethod final { delete static_cast*>(pointer); })); } else { - result = py::cast( - PyExecuTorchResult::hand_over(tensor, memory, handed_out_[*slot])); + result = py::cast(PyExecuTorchResult::hand_over(tensor, memory)); } handed_buffers_[*slot] = holder; - handed_slots_[*slot] = true; handed_outputs_.resize(storage_->method->outputs_size()); handed_outputs_[index] = result; return result; @@ -3853,9 +3820,7 @@ struct PyMethod final { storage_->output_storages = std::move(output_storages); storage_->device_output_storages = std::move(device_output_storages); handed_buffers_.assign(num_outputs, std::weak_ptr()); - handed_slots_.assign(num_outputs, false); spare_device_outputs_.assign(num_outputs, SpareDeviceOutput()); - handed_out_.assign(num_outputs, nullptr); } // A handed-over result still uses its buffer, so the method writes into @@ -3866,46 +3831,26 @@ struct PyMethod final { #ifndef USE_ATEN_LIB for (size_t i = 0; i < storage_->device_output_storages.size(); ++i) { auto& buffer = storage_->device_output_storages[i]; - if (!buffer || !handed_slots_[i]) { + if (!buffer || handed_buffers_[i].expired()) { continue; } const auto device = storage_->method->get_output(i).toTensor().device(); - if (!handed_buffers_[i].expired()) { - // In a loop such as `out = method(x)[0]` the previous result is still - // alive here, so its buffer and this one take turns instead of a new - // allocation each call. - auto& spare = spare_device_outputs_[i]; - if (spare.buffer && spare.handed.expired()) { - std::swap(buffer, spare.buffer); - std::swap(handed_buffers_[i], spare.handed); - std::swap(handed_out_[i], spare.handed_out); - wait_for_other_readers(i, device); - } else { - spare = {buffer, handed_buffers_[i], handed_out_[i]}; - buffer = allocate_device_output(i, buffer->size(), device); - } + auto& spare = spare_device_outputs_[i]; + if (spare.buffer && spare.handed.expired()) { + std::swap(buffer, spare.buffer); + std::swap(handed_buffers_[i], spare.handed); } else { - wait_for_other_readers(i, device); + // Allocated into a local first, so a failed allocation leaves the slot + // and the spare untouched. + auto fresh = allocate_device_output(i, buffer->size(), device); + spare = {std::move(buffer), handed_buffers_[i]}; + buffer = std::move(fresh); + handed_buffers_[i].reset(); } - handed_buffers_[i].reset(); - handed_out_[i].reset(); - handed_slots_[i] = false; } #endif } -#ifndef USE_ATEN_LIB - // A dropped result's memory may still be read by work another library - // queued on its own streams. Freeing it would have waited for that work, so - // reusing it waits too. Work these bindings queued has already finished, - // since every execution waits for its device before it returns. - void wait_for_other_readers(size_t index, executorch::aten::Device device) { - if (handed_out_[index] && *handed_out_[index]) { - synchronize_device(device.index()); - } - } -#endif - py::list get_outputs_as_py_list( const std::vector& outputs, bool clone_outputs = true) { @@ -3933,7 +3878,11 @@ struct PyMethod final { list[i] = py::cast(v.toTensor()); } #else - if (i < handed_outputs_.size() && handed_outputs_[i]) { + // Within one execution, repeated default reads return the same handed + // over object, so a caller sees stable identity and the buffer is not + // handed over twice. A clone_outputs=False read is a separate borrowed + // view and never returns that cached object. + if (clone_outputs && i < handed_outputs_.size() && handed_outputs_[i]) { list[i] = handed_outputs_[i]; continue; } diff --git a/extension/pybindings/pybindings_cuda.h b/extension/pybindings/pybindings_cuda.h index 1512f719862..a61a7241cff 100644 --- a/extension/pybindings/pybindings_cuda.h +++ b/extension/pybindings/pybindings_cuda.h @@ -121,8 +121,8 @@ auto with_cuda_device(int index, Function&& function) { // Copies on the legacy default stream. A copy to the host returns once it has // finished. A copy to the device may still be in flight when this returns, but // it runs before anything queued afterwards on a blocking stream, including the -// delegates' cudaStreamPerThread; a result holding such memory waits for it -// only when it hands the memory to another library. +// delegates' cudaStreamPerThread; a result holding such memory still waits for +// the device before another library or a call reads it. inline void copy_cuda_storage( void* destination, const void* source, diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 465e9330e3f..3a919f60aba 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -91,6 +91,40 @@ inline executorch::aten::ScalarType scalar_type_from_torch_dtype( throw py::value_error("Unsupported torch dtype: " + torch_dtype); } +inline PyObject* loaded_torch_module() { + // Retain the module for the interpreter's lifetime. This avoids repeated + // sys.modules lookups without creating a static py::object whose destructor + // could run after Python has finalized. + static PyObject* torch_module = nullptr; + if (torch_module == nullptr) { + PyObject* loaded = PyDict_GetItemString(PyImport_GetModuleDict(), "torch"); + if (loaded != nullptr) { + Py_INCREF(loaded); + torch_module = loaded; + } + } + return torch_module; +} + +inline void synchronize_torch_cuda_device(int index) { + py::reinterpret_borrow(loaded_torch_module()) + .attr("cuda") + .attr("synchronize")(index); +} + +// Waits for every queued operation on a device. Bindings built without CUDA +// support have no CUDA runtime of their own, so they wait through torch's, +// which is loaded whenever such bindings see CUDA memory. +inline void synchronize_device(int index) { +#ifndef EXECUTORCH_PYBINDINGS_CUDA + if (loaded_torch_module() != nullptr) { + synchronize_torch_cuda_device(index); + return; + } +#endif + synchronize_cuda_device(index); +} + /** Tensor result returned by portable Python bindings. */ class PyExecuTorchResult final : public std::enable_shared_from_this { @@ -251,34 +285,19 @@ class PyExecuTorchResult final } py::capsule to_dlpack(const py::object& stream = py::none()) { - // A result is ready on any stream once its upload, if any, has finished. - if (copy_in_flight_ && !device_.is_cpu() && nbytes_ != 0) { - synchronize_cuda_device(device_.index()); + // Always wait for the device before handing CUDA memory to another + // library: the result may share memory a call still has work queued on, + // or an upload that is still running, and this is cheap next to a call. + if (!device_.is_cpu() && nbytes_ != 0) { + synchronize_device(device_.index()); } - wait_for_other_writers(); (void)stream; - copy_in_flight_ = false; - *handed_out_ = true; return capsule(); } - // Memory handed to another library may since have been written by it, on a - // stream these bindings do not know, so every read of it waits first. - void wait_for_other_writers() const { - if (*handed_out_ && !device_.is_cpu() && nbytes_ != 0) { - synchronize_cuda_device(device_.index()); - } - } - - // Whether another library has held this result's memory and so may have - // written to it on its own streams. - bool handed_out() const { - return *handed_out_; - } - - // A capsule for these bindings to read the result as a method input, which - // does not count as handing it out: see DLPackTensorView for why such an - // input needs no wait. + // A capsule for these bindings to read the result as a method input. It + // skips the wait in to_dlpack(), since the call waits for every CUDA input + // before it runs. py::capsule capsule() { auto holder = std::make_unique(shared_from_this()); py::capsule capsule(&holder->managed, "dltensor", [](PyObject* capsule) { @@ -304,7 +323,9 @@ class PyExecuTorchResult final if (device_.is_cpu()) { std::memcpy(storage.data(), data_, nbytes_); } else { - wait_for_other_writers(); + // Wait before the copy: the result may share memory a call still has + // work queued on, or an upload that is still running. + synchronize_cuda_device(device_.index()); copy_cuda_storage( storage.data(), data_, @@ -368,21 +389,17 @@ class PyExecuTorchResult final device_index, CudaCopy::HostToDevice); result->data_ = result->device_storage_.get(); - result->copy_in_flight_ = true; } return result; } // `storage` must own the memory `tensor` points at; the result reads that - // memory and keeps `storage` alive. `handed_out` lets the caller see, after - // the result is gone, whether its memory was handed to another library. + // memory and keeps `storage` alive. static std::shared_ptr hand_over( const executorch::aten::Tensor& tensor, - std::shared_ptr storage, - std::shared_ptr handed_out) { + std::shared_ptr storage) { auto result = std::make_shared(tensor, /*clone=*/false); result->device_storage_ = std::move(storage); - result->handed_out_ = std::move(handed_out); return result; } @@ -522,9 +539,6 @@ class PyExecuTorchResult final size_t nbytes_; executorch::aten::Device device_; py::object owner_; - std::shared_ptr handed_out_ = std::make_shared(false); - // A copy into this result's device memory may still be in flight. - bool copy_in_flight_ = false; }; /** Contiguous buffer facade used to construct a strided torch.Tensor result. */ diff --git a/extension/pybindings/test/make_test.py b/extension/pybindings/test/make_test.py index c11178aaf16..a57ad0ab7d2 100644 --- a/extension/pybindings/test/make_test.py +++ b/extension/pybindings/test/make_test.py @@ -38,6 +38,17 @@ def get_inputs(self): ) +class ModuleAddBFloat16ToChannelsLast(ModuleAdd): + def forward(self, x, y): + return (x + y).to(memory_format=torch.channels_last) + + def get_inputs(self): + return ( + torch.ones(1, 2, 2, 2, dtype=torch.bfloat16), + torch.ones(1, 2, 2, 2, dtype=torch.bfloat16), + ) + + class ModuleChannelsLast(torch.nn.Module): """The module to serialize and execute.""" diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 4e5da6cf330..79d88fa73b2 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -2013,6 +2013,7 @@ def test_method_rejects_input_on_unrepresentable_device(self): with self.assertRaises(RuntimeError) as caught: executorch_module.forward([inputs[0].to("meta"), inputs[1]]) message = str(caught.exception) + self.assertIn("Input 0 of method 'forward'", message) # Asserts the device is named as Python spells it, since an uppercased or # index-less name would not match what the caller passed. self.assertIn("is on device meta", message) @@ -2482,6 +2483,7 @@ def test_program_method_rejects_input_on_unrepresentable_device(self): with self.assertRaises(RuntimeError) as caught: method.set_inputs([inputs[0].to("meta"), inputs[1]]) message = str(caught.exception) + self.assertIn("Input 0 of method 'forward'", message) self.assertIn("is on device meta", message) self.assertIn("only CPU and CUDA tensors", message) diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index d04cd3e8652..990b4bfbfae 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -74,10 +74,13 @@ def setUpClass(cls): ctypes.c_uint64, ] cls.driver.cuPointerGetAttribute.restype = ctypes.c_int - cls._export_delegate() + cls.program_path, cls.data_path = cls._export_delegate("cuda_dlpack") + cls.host_input_program = cls._export_delegate( + "host_inputs", inputs_on_device=False + ) @classmethod - def _export_delegate(cls): + def _export_delegate(cls, name, inputs_on_device=True): from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.exir import ExecutorchBackendConfig, to_edge_transform_and_lower from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -103,7 +106,8 @@ def forward(self, value): program = edge.to_executorch( config=ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( - skip_h2d_for_method_inputs=True, skip_d2h_for_method_outputs=True + skip_h2d_for_method_inputs=inputs_on_device, + skip_d2h_for_method_outputs=True, ), enable_non_cpu_memory_planning=True, memory_planning_pass=MemoryPlanningPass( @@ -120,24 +124,33 @@ def forward(self, value): raise AssertionError( "test must cover borrowed inputs and unplanned outputs" ) - if tensor.extra_tensor_info.device_type != DeviceType.CUDA: + if (inputs_on_device or index in plan.outputs) and ( + tensor.extra_tensor_info.device_type != DeviceType.CUDA + ): raise AssertionError("test must retain CUDA inputs and outputs") - cls.program_path = Path(cls.directory.name) / "cuda_dlpack.pte" - with cls.program_path.open("wb") as source: + directory = Path(cls.directory.name) / name + directory.mkdir() + program_path = directory / f"{name}.pte" + with program_path.open("wb") as source: program.write_to_file(source) data_names = sorted(program._tensor_data or {}) - program.write_tensor_data_to_file(cls.directory.name) - cls.data_path = ( - Path(cls.directory.name) / (data_names[0] + ".ptd") if data_names else None + program.write_tensor_data_to_file(str(directory)) + return program_path, ( + directory / (data_names[0] + ".ptd") if data_names else None ) - def _load_method(self): + def _load_method(self, host_inputs=False): from executorch.runtime import Runtime + program_path, data_path = ( + self.host_input_program + if host_inputs + else (self.program_path, self.data_path) + ) return ( Runtime.get() - .load_program(self.program_path, data_path=self.data_path) + .load_program(program_path, data_path=data_path) .load_method("forward") ._method ) @@ -180,6 +193,8 @@ def _assert_ready_value(self, output, expected): ready.record() # Only wait for the consumer. A global wait before the read would hide # missing ordering between the producer, delegate, and consumer streams. + # An ExecuTorchResult waits for the device in __dlpack__ already, so in + # that output mode this catches a wrong value, not a missing wait. ready.synchronize() torch.testing.assert_close(snapshot.cpu(), torch.full(self.shape, expected)) return tensor @@ -237,6 +252,33 @@ def test_result_handed_out_after_set_inputs_is_waited_for(self): method.execute() self._assert_ready_value(method.get_outputs()[0], 84.0) + def test_borrowed_result_written_by_another_library_is_waited_for(self): + # Two clone_outputs=False reads share the same memory. Another library + # writes it through one of them, so passing the other back in must + # wait for that write. + method = self._load_method() + stream = torch.cuda.current_stream(self.device) + # A module reads it back: it writes each call into new memory, so only + # the wait for its input can order the read. Its first run sets it up + # and waits for the whole device, so it runs once before the write. Its + # output stays alive, since freeing it in the next call waits as well. + program = self.program_path.read_bytes() + data = self.data_path.read_bytes() if self.data_path else None + module = self.bindings._load_for_executorch_from_buffer(program, data) + kept = module.forward( + [CudaDLPackOnly(torch.zeros(self.shape, device=self.device), stream)] + ) + value = torch.full(self.shape, 1.0, device=self.device) + written = method(CudaDLPackOnly(value, stream), clone_outputs=False)[0] + passed = method.get_outputs(clone_outputs=False)[0] + writer = torch.cuda.Stream(self.device) + with torch.cuda.stream(writer): + shared = torch.from_dlpack(written) + torch.cuda._sleep(50_000_000) + shared.fill_(41.0) + self._assert_ready_value(module.forward([passed])[0], 84.0) + del kept + def test_clone_roundtrip_and_borrowed_output_lifetime(self): method = self._load_method() stream = torch.cuda.current_stream(self.device) @@ -301,7 +343,7 @@ def run(value, clone_outputs=True): self.assertEqual(handed_copies, borrowed_copies) self._assert_ready_value(first, 8.0) self.assertIs(method.get_outputs()[0], first) - self.assertIs(method.get_outputs(clone_outputs=False)[0], first) + self.assertIsNot(method.get_outputs(clone_outputs=False)[0], first) # Feeding the result back must not write the next output into it. second = method(first)[0] @@ -520,29 +562,38 @@ def forward(self, value): self._assert_ready_value(third, 9.0) def test_reused_output_memory_waits_for_pending_readers(self): - method = self._load_method() + # Once the result is dropped, the next run writes into its memory, so it + # must wait for a read another library still has queued there, whether + # its own input is on the device or on the host. count = torch.Size(self.shape).numel() - host = memoryview(array.array("f", [1.0]) * count).cast("B") - # Inputs from these bindings need no wait of their own, so only the - # reuse wait can keep the next run from overwriting a pending read. - first = self.bindings.ExecuTorchResult.from_host( - host.cast("f", self.shape), self.device.index - ) - second = self.bindings.ExecuTorchResult.from_host( - memoryview(array.array("f", [5.0]) * count).cast("B").cast("f", self.shape), - self.device.index, - ) - torch.cuda.synchronize(self.device) - result = method(first)[0] - reader = torch.cuda.Stream(self.device) - with torch.cuda.stream(reader): - torch.cuda._sleep(1_000_000_000) - snapshot = self._tensor(result).clone() - del result - gc.collect() - method(second) - reader.synchronize() - torch.testing.assert_close(snapshot.cpu(), torch.full(self.shape, 4.0)) + + def on_device(value): + host = memoryview(array.array("f", [value]) * count).cast("B") + return self.bindings.ExecuTorchResult.from_host( + host.cast("f", self.shape), self.device.index + ) + + def on_host(value): + return torch.full(self.shape, value) + + for host_inputs, make_input in ((False, on_device), (True, on_host)): + with self.subTest(host_inputs=host_inputs): + method = self._load_method(host_inputs=host_inputs) + first, second = make_input(1.0), make_input(5.0) + torch.cuda.synchronize(self.device) + result = method(first)[0] + # Imported before the read is queued: an ExecuTorchResult waits + # for the device in __dlpack__, which would finish the read. + view = self._tensor(result) + reader = torch.cuda.Stream(self.device) + with torch.cuda.stream(reader): + torch.cuda._sleep(1_000_000_000) + snapshot = view.clone() + del result, view + gc.collect() + method(second) + reader.synchronize() + torch.testing.assert_close(snapshot.cpu(), torch.full(self.shape, 4.0)) def test_method_outputs_take_turns_between_two_buffers(self): method = self._load_method() @@ -666,12 +717,12 @@ def test_bfloat16_host_result_moves_to_the_device(self): with self.assertRaisesRegex(ValueError, "already on a device"): self.bindings.ExecuTorchResult.from_host(on_device, self.device.index) - def test_device_chain_waits_for_the_device_once_per_call(self): - # Feeding each output into the next call, with a host upload per call, - # needs one device wait per call. Waiting per tensor made a chain of - # small programs measurably slower. - if not self.executorch_output: - self.skipTest("chains ExecuTorchResult outputs") + def test_device_chain_waits_for_the_device_twice_per_call(self): + # Feeding each output into the next call, with a host upload per call. + # Each call waits for the device twice: once before it runs, in case + # another library still uses its input or output memory, and once after, + # so its output is ready on any stream. The count is fixed per call, not + # per tensor. from torch.profiler import profile, ProfilerActivity method = self._load_method() @@ -700,14 +751,14 @@ def chain(): self.bindings.ExecuTorchResult.from_host(host, self.device.index) value = method(value, clone_outputs=True)[0] - # The profiler waits for the device itself when it stops. Each call - # waits exactly once: none would return outputs still being written. - self.assertEqual(waits(chain) - waits(lambda: None), calls) + # The profiler waits for the device itself when it stops. + self.assertEqual(waits(chain) - waits(lambda: None), 2 * calls) expected = 8.0 for _ in range(calls): expected = (expected + 1.0) * 2.0 + final = value.to_host() if self.executorch_output else value.cpu() torch.testing.assert_close( - torch.from_dlpack(value.to_host()), torch.full(self.shape, expected) + torch.from_dlpack(final), torch.full(self.shape, expected) ) def _require_libgomp_without_torch(self): diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index 7bdc9ae54b3..b5d08deae35 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -311,6 +311,56 @@ def test_bfloat16_raw_buffer_input_and_output(self) -> None: tensor, torch.full((2, 2), 2.0, dtype=torch.bfloat16) ) + def test_bfloat16_result_in_another_layout_feeds_back_correctly(self) -> None: + # The method takes contiguous inputs and returns a channels-last result. + # A bfloat16 result exports flat bytes through the buffer protocol, so + # read that way it would take the method's dim order and come out + # permuted, with no error. + if runtime._tensor_output != "executorch": + self.skipTest("ExecuTorchResult output was not selected") + with open( + os.environ["EXECUTORCH_PYBIND_TEST_BFLOAT16_CHANNELS_LAST_PTE"], "rb" + ) as program_file: + program_data = program_file.read() + method = runtime._load_program_from_buffer(program_data).load_method("forward") + # bfloat16 is the upper half of float32, so these are exactly 0 to 7. + values = (np.arange(8, dtype=np.float32).view(np.uint32) >> 16).astype( + np.uint16 + ) + zeros = bytes(values.nbytes) + channels_last = values.reshape(1, 2, 2, 2).transpose(0, 2, 3, 1).tobytes() + + result = method((values.tobytes(), zeros))[0] + self.assertEqual(result.strides, (16, 2, 8, 4)) + fed_back = method((result, zeros))[0] + self.assertEqual(memoryview(fed_back).tobytes(), channels_last) + + def test_input_errors_name_the_input(self) -> None: + class FailingProducer: + def __dlpack_device__(self): + return (1, 0) + + def __dlpack__(self, stream=None): + raise UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte") + + with open(os.environ["EXECUTORCH_PYBIND_TEST_PTE"], "rb") as program_file: + program_data = program_file.read() + module = runtime._load_for_executorch_from_buffer(program_data) + method = runtime._load_program_from_buffer(program_data).load_method("forward") + value = np.array([1.0], dtype=np.float32) + for executor in (module, method): + with self.subTest(executor=type(executor)): + # A Python error keeps its type, even one whose constructor takes + # more than a message, and carries the name as a note. + with self.assertRaises(UnicodeDecodeError) as caught: + executor((FailingProducer(), value)) + self.assertIn("Input 0 of method 'forward'", caught.exception.__notes__) + # A method reads every input before it checks their count. + with self.assertRaisesRegex( + RuntimeError, "Input 2 of method 'forward': not a tensor input" + ): + method((value, value, value)) + def test_mixed_dlpack_and_buffer_inputs(self) -> None: with open(os.environ["EXECUTORCH_PYBIND_TEST_PTE"], "rb") as program_file: program_data = program_file.read()