From 9f500f1f2043f9bf7f43ab27e08e802bf910a473 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 00:56:49 -0700 Subject: [PATCH 01/10] [PyBindings] Reject host buffers for device inputs, wait for every tensor of a run, and name input errors consistently Follow-up to the device wait rework in the Python bindings. A NumPy array or another host buffer passed for an input that the method expects on a CUDA device was accepted, and the method then read a host pointer as device memory. Torch tensors and DLPack inputs were already checked. Now a non-empty buffer is rejected the same way, before any host layout work, and the message says how to pass device memory instead. Before a run, the bindings wait in case another library still uses the memory of an input, or of an output from an earlier run. That wait covered only outputs in memory the bindings allocated, and a Module call covered no outputs at all, so a memory-planned output could be overwritten while another library was still reading it. After a run, only outputs were waited for, although a method may update an input in place. Now one helper waits, before and after each run, for every device that holds a non-empty tensor input or output of the method. Input errors now name the input and method the same way: a DLPack device mismatch names both devices; a changed input storage, in both the borrow check and the aliased torch output, says to pass the input again instead of calling set_inputs(); a wrong number or an error while choosing how to read an input gets the name; and a Method checks the input count before it reads any input. A Module call now looks its method up once. The README says that each wait is for the whole device, that it holds the GIL in CUDA-enabled bindings, how short a wait is on an idle device, that CUDA outputs come back without a copy only with MemoryPlanningPass(alloc_graph_output=False), and what a planned CUDA output costs with clone_outputs=True. --- extension/pybindings/README.md | 42 ++- extension/pybindings/pybindings.cpp | 270 +++++++++--------- extension/pybindings/pybindings_cuda.h | 4 +- .../pybindings/pybindings_executorch_result.h | 2 +- extension/pybindings/test/test_pybindings.py | 91 ++++-- .../test/test_pybindings_cuda_dlpack.py | 88 ++++-- .../test/test_pybindings_no_aten.py | 31 +- 7 files changed, 316 insertions(+), 212 deletions(-) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 8b031ff734b..04a019d2bc1 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -45,7 +45,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 @@ -105,18 +111,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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 635a4393184..b7d9b35c6ee 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -720,12 +720,38 @@ void validate_buffer_input( } #ifndef USE_ATEN_LIB +void validate_input_device( + const Method& method, + size_t index, + executorch::aten::Device device, + size_t nbytes) { + if (nbytes == 0) { + return; + } + const auto expected = method.get_input(index).toTensor().device(); + if (device != expected) { + const auto name = [](executorch::aten::Device value) { + return value.is_cpu() ? std::string("cpu") + : "cuda:" + std::to_string(value.index()); + }; + std::string message = "Input " + std::to_string(index) + " of method '" + + method.method_meta().name() + "' is on device " + name(device) + + ", but the method expects " + name(expected); + if (device.is_cpu()) { + message += + ". Pass device memory instead, such as " + "ExecuTorchResult.from_host(...) or a CUDA tensor"; + } + throw py::value_error(message); + } +} + /** A torch.Tensor input read through its public Python API, copied only to * change its layout. */ class TorchTensorView final { public: - explicit TorchTensorView(const py::handle& value) - : owner_(py::reinterpret_borrow(value)) { + TorchTensorView(const py::handle& value, size_t index) + : owner_(py::reinterpret_borrow(value)), index_(index) { if (owner_.attr("is_neg")().cast() || owner_.attr("is_conj")().cast()) { throw py::value_error( @@ -774,6 +800,10 @@ class TorchTensorView final { data_ = reinterpret_cast(address); } + size_t index() const { + return index_; + } + void* data() const { return data_; } @@ -845,23 +875,6 @@ class TorchTensorView final { dim_order_.assign(dim_order.begin(), dim_order.end()); } - void validate_device(const Method& method, size_t index) const { - if (nbytes_ == 0) { - return; - } - const auto expected = method.get_input(index).toTensor().device(); - if (device_ != expected) { - const auto name = [](executorch::aten::Device device) { - return device.is_cpu() ? std::string("cpu") - : "cuda:" + std::to_string(device.index()); - }; - throw py::value_error( - "Input " + std::to_string(index) + " of method '" + - method.method_meta().name() + "' is on device " + name(device_) + - ", but the method expects " + name(expected)); - } - } - private: static std::vector checked_vector( const py::handle& values, @@ -893,6 +906,7 @@ class TorchTensorView final { } py::object owner_; + size_t index_; void* data_ = nullptr; std::vector sizes_; std::vector strides_; @@ -917,13 +931,13 @@ class BorrowedInput final { data_(buffer_->data()), nbytes_(buffer_->nbytes()) {} - void validate() const { + void validate(const char* method_name) const { const auto [current_data, current_nbytes] = buffer_ ? buffer_->current_storage() : torch_storage(); if (current_data != data_ || current_nbytes != nbytes_) { throw py::value_error( - "Input " + std::to_string(index_) + - " storage changed after set_inputs(); call set_inputs() again"); + "Input " + std::to_string(index_) + " of method '" + method_name + + "': its storage changed after it was passed in; pass it again"); } } @@ -1070,16 +1084,6 @@ class DLPackTensorView final { dim_order_.assign(dim_order.begin(), dim_order.end()); } - void validate_device(const Method& method, size_t index) const { - if (nbytes_ != 0 && - device_ != method.get_input(index).toTensor().device()) { - throw py::value_error( - "Input " + std::to_string(index) + " of method '" + - method.method_meta().name() + - "': DLPack input device does not match the method input device"); - } - } - py::object result_backing_owner() const { if (!converted_.empty()) { return py::none(); @@ -1370,7 +1374,8 @@ py::object dlpack_output_owner( py::object module_aliased_torch_output( const executorch::aten::Tensor& tensor, - const std::vector>& torch_inputs) { + const std::vector>& torch_inputs, + const std::string& method_name) { if (torch_inputs.empty() || portable_tensor_output() != PortableTensorOutput::Torch || tensor.nbytes() == 0) { @@ -1399,7 +1404,9 @@ py::object module_aliased_torch_output( py::str(source.attr("dtype")).cast()) != tensor.scalar_type()) { throw py::value_error( - "Input storage changed before get_outputs(); call set_inputs() again"); + "Input " + std::to_string(input->index()) + " of method '" + + method_name + + "': its storage changed before the output was read; pass it again"); } const auto offset = source.attr("storage_offset")().cast() + static_cast((address - start) / itemsize); @@ -1455,9 +1462,8 @@ std::optional scalar_input( return EValue(number.cast()); } throw std::runtime_error( - "Input " + std::to_string(index) + " of method '" + method_meta.name() + - "' is " + runtime::tag_to_string(tag.get()) + ", but got " + - std::string(py::str(value.get_type()))); + std::string("the method expects ") + runtime::tag_to_string(tag.get()) + + ", but got " + std::string(py::str(value.get_type()))); } bool prefer_dlpack_input(const py::handle& value) { @@ -1507,57 +1513,32 @@ std::shared_ptr allocate_device_output( return std::make_shared(std::move(buffer.get())); } -// 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. -void synchronize_device_outputs(const Method& method) { +// The Python boundary is synchronous, and each wait is for the whole device, +// so one wait per device covers every tensor of a run. Before the run, 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 the memory of an +// output from an earlier run, such as a dropped result's buffer or a +// memory-planned output, which this run writes again. After the run, its +// outputs, and any input it updated in place, are ready on every stream. +void synchronize_tensor_devices(const Method& method) { std::vector devices; - for (size_t i = 0; i < method.outputs_size(); ++i) { - const auto& output = method.get_output(i); - if (!output.isTensor() || output.toTensor().nbytes() == 0) { - continue; - } - const auto device = output.toTensor().device(); - if (!device.is_cpu() && - std::find(devices.begin(), devices.end(), device.index()) == - devices.end()) { - devices.push_back(device.index()); + const auto add = [&devices](const EValue& value) { + if (!value.isTensor() || value.toTensor().nbytes() == 0) { + return; } - } - for (const auto index : devices) { - synchronize_device(index); - } -} - -// 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& output_devices = {}) { - std::vector devices; - const auto add = [&devices](executorch::aten::Device device) { + const auto device = value.toTensor().device(); if (!device.is_cpu() && std::find(devices.begin(), devices.end(), device.index()) == devices.end()) { devices.push_back(device.index()); } }; - for (const auto& input : dlpack_inputs) { - if (input->nbytes() != 0) { - add(input->device()); - } + for (size_t i = 0; i < method.inputs_size(); ++i) { + add(method.get_input(i)); } - for (const auto& input : torch_inputs) { - if (input->nbytes() != 0) { - add(input->device()); - } - } - for (const auto device : output_devices) { - add(device); + for (size_t i = 0; i < method.outputs_size(); ++i) { + add(method.get_output(i)); } for (const auto index : devices) { synchronize_device(index); @@ -1915,7 +1896,7 @@ struct PyBundledModule : public BundledModule { method.error(), "Method execution failed with status 0x%" PRIx32, static_cast(method.error())); - synchronize_device_outputs(*method.get()); + synchronize_tensor_devices(*method.get()); } #endif @@ -2351,6 +2332,15 @@ struct PyModule final { "' failed while setting its inputs, so the inputs it would reuse " "may be a mix of two calls. Pass every input on the next call."); } +#ifndef USE_ATEN_LIB + const auto loaded_method = module_->method(method_name); + THROW_IF_ERROR( + loaded_method.error(), + "Failed to execute method %s, error: 0x%" PRIx32, + method_name.c_str(), + static_cast(loaded_method.error())); + const Method& method = *loaded_method.get(); +#endif std::vector new_input_owners(inputs_size); std::vector cpp_inputs; cpp_inputs.reserve(inputs_size); @@ -2385,9 +2375,13 @@ struct PyModule final { for (size_t i = 0; i < inputs_size; ++i) { auto python_input = inputs[i]; const std::string& type_str = py::str(python_input.get_type()); - if (auto scalar = scalar_input(method_meta, i, python_input)) { + if (auto scalar = read_input(method_meta, i, [&] { + return scalar_input(method_meta, i, python_input); + })) { cpp_inputs.push_back(*scalar); - } else if (is_torch_tensor(python_input)) { + } else if (read_input(method_meta, i, [&] { + return is_torch_tensor(python_input); + })) { #ifdef USE_ATEN_LIB auto at_tensor = python_input.cast(); new_input_owners[i].torch = @@ -2403,19 +2397,13 @@ struct PyModule final { } #else torch_inputs.push_back(read_input(method_meta, i, [&] { - return std::make_shared(python_input); + return std::make_shared(python_input, i); })); const auto& tensor = torch_inputs.back(); new_input_owners[i].torch = tensor; validate_tensor_input( method_meta, i, tensor->scalar_type(), tensor->sizes()); - const auto method = module_->method(method_name); - THROW_IF_ERROR( - method.error(), - "Failed to execute method %s, error: 0x%" PRIx32, - method_name.c_str(), - static_cast(method.error())); - tensor->validate_device(*method.get(), i); + validate_input_device(method, i, tensor->device(), tensor->nbytes()); read_input(method_meta, i, [&] { tensor->use_layout(method_meta.input_tensor_meta(i)->dim_order()); }); @@ -2448,7 +2436,9 @@ struct PyModule final { cpp_inputs.push_back(EValue(py::cast(python_input))); } else if (py::isinstance(python_input)) { cpp_inputs.push_back(EValue(py::cast(python_input))); - } else if (prefer_dlpack_input(python_input)) { + } else if (read_input(method_meta, i, [&] { + return prefer_dlpack_input(python_input); + })) { #ifdef USE_ATEN_LIB auto python_tensor = py::module_::import("torch.utils.dlpack") .attr("from_dlpack")(python_input); @@ -2469,13 +2459,7 @@ struct PyModule final { new_input_owners[i].dlpack = view; validate_tensor_input( method_meta, i, view->scalar_type(), view->sizes()); - const auto method = module_->method(method_name); - THROW_IF_ERROR( - method.error(), - "Failed to execute method %s, error: 0x%" PRIx32, - method_name.c_str(), - static_cast(method.error())); - view->validate_device(*method.get(), i); + validate_input_device(method, i, view->device(), view->nbytes()); read_input(method_meta, i, [&] { view->use_layout(method_meta.input_tensor_meta(i)->dim_order()); }); @@ -2494,6 +2478,13 @@ struct PyModule final { })); const auto& buffer = buffer_inputs.back(); new_input_owners[i].buffer = buffer; +#ifndef USE_ATEN_LIB + validate_input_device( + method, + i, + executorch::aten::Device(executorch::aten::DeviceType::CPU), + buffer->nbytes()); +#endif validate_buffer_input(method_meta, i, *buffer); if (!method_meta.input_tensor_meta(i)->is_memory_planned() && buffer->borrows_data()) { @@ -2540,7 +2531,7 @@ struct PyModule final { auto& retained = retained_inputs_[method_name]; for (size_t i = inputs_size; i < retained.size(); ++i) { if (retained[i].validation) { - retained[i].validation->validate(); + retained[i].validation->validate(method_name.c_str()); } } // A load error is the real cause, so it wins over the checks below. @@ -2599,7 +2590,7 @@ struct PyModule final { dlpack_inputs.push_back(owner.dlpack); } } - synchronize_devices_before_run(dlpack_inputs, torch_inputs); + synchronize_tensor_devices(method); #endif auto outputs = module_->execute(method_name); THROW_IF_ERROR( @@ -2608,15 +2599,7 @@ struct PyModule final { method_name.c_str(), static_cast(outputs.error())); #ifndef USE_ATEN_LIB - { - const auto method = module_->method(method_name); - THROW_IF_ERROR( - method.error(), - "Failed to execute method %s, error: 0x%" PRIx32, - method_name.c_str(), - static_cast(method.error())); - synchronize_device_outputs(*method.get()); - } + synchronize_tensor_devices(method); #endif // DLPack producers may supply storage returned directly by the graph. // Keep each output's allocation alive after local adapters go out of scope. @@ -2648,8 +2631,8 @@ struct PyModule final { for (size_t i = 0; i < outputs->size(); ++i) { const auto& output = (*outputs)[i]; if (output.isTensor()) { - auto alias = - module_aliased_torch_output(output.toTensor(), torch_inputs); + auto alias = module_aliased_torch_output( + output.toTensor(), torch_inputs, method_name); if (!alias.is_none()) { result[i] = std::move(alias); } @@ -3189,6 +3172,12 @@ struct PyMethod final { ObjectLock lock(mutex_); const auto inputs = normalize_inputs(python_inputs); const auto inputs_size = py::len(inputs); + if (inputs_size != storage_->method->inputs_size()) { + throw std::runtime_error( + "Method '" + std::string(storage_->method->method_meta().name()) + + "' takes " + std::to_string(storage_->method->inputs_size()) + + " inputs, but got " + std::to_string(inputs_size)); + } std::vector cpp_inputs; cpp_inputs.reserve(inputs_size); std::vector> buffer_inputs; @@ -3219,10 +3208,14 @@ struct PyMethod final { for (size_t i = 0; i < inputs_size; ++i) { auto python_input = inputs[i]; const std::string& type_str = py::str(python_input.get_type()); - if (auto scalar = - scalar_input(storage_->method->method_meta(), i, python_input)) { + if (auto scalar = read_input(storage_->method->method_meta(), i, [&] { + return scalar_input( + storage_->method->method_meta(), i, python_input); + })) { cpp_inputs.push_back(*scalar); - } else if (is_torch_tensor(python_input)) { + } else if (read_input(storage_->method->method_meta(), i, [&] { + return is_torch_tensor(python_input); + })) { #ifdef USE_ATEN_LIB auto at_tensor = python_input.cast(); std::vector tensor_sizes( @@ -3242,7 +3235,7 @@ struct PyMethod final { #else torch_inputs.push_back( read_input(storage_->method->method_meta(), i, [&] { - return std::make_shared(python_input); + return std::make_shared(python_input, i); })); const auto& view = torch_inputs.back(); validate_tensor_input( @@ -3250,7 +3243,8 @@ struct PyMethod final { i, view->scalar_type(), view->sizes()); - view->validate_device(*storage_->method, i); + validate_input_device( + *storage_->method, i, view->device(), view->nbytes()); read_input(storage_->method->method_meta(), i, [&] { view->use_layout(storage_->method->method_meta() .input_tensor_meta(i) @@ -3280,7 +3274,9 @@ struct PyMethod final { cpp_inputs.push_back(EValue(py::cast(python_input))); } else if (py::isinstance(python_input)) { cpp_inputs.push_back(EValue(py::cast(python_input))); - } else if (prefer_dlpack_input(python_input)) { + } else if (read_input(storage_->method->method_meta(), i, [&] { + return prefer_dlpack_input(python_input); + })) { #ifdef USE_ATEN_LIB auto python_tensor = py::module_::import("torch.utils.dlpack") .attr("from_dlpack")(python_input); @@ -3308,7 +3304,8 @@ struct PyMethod final { i, view->scalar_type(), view->sizes()); - view->validate_device(*storage_->method, i); + validate_input_device( + *storage_->method, i, view->device(), view->nbytes()); read_input(storage_->method->method_meta(), i, [&] { view->use_layout(storage_->method->method_meta() .input_tensor_meta(i) @@ -3330,6 +3327,13 @@ struct PyMethod final { python_input, storage_->method->method_meta(), i); })); const auto& buffer = buffer_inputs.back(); +#ifndef USE_ATEN_LIB + validate_input_device( + *storage_->method, + i, + executorch::aten::Device(executorch::aten::DeviceType::CPU), + buffer->nbytes()); +#endif validate_buffer_input(storage_->method->method_meta(), i, *buffer); const auto input_meta = storage_->method->method_meta().input_tensor_meta(i); @@ -3369,17 +3373,6 @@ struct PyMethod final { cpp_inputs.data(), cpp_inputs.size()); Error set_inputs_status = storage_->method->set_inputs(input_evalue_list); - // The runtime fails a wrong count before installing anything, so the - // method keeps its previous inputs and these must not be retained. - if (set_inputs_status != Error::Ok && - cpp_inputs.size() != storage_->method->inputs_size()) { - THROW_IF_ERROR( - set_inputs_status, - "Method '%s' takes %zu inputs, but set_inputs() got %zu", - storage_->method->method_meta().name(), - storage_->method->inputs_size(), - cpp_inputs.size()); - } if (set_inputs_status != Error::Ok) { // set_inputs() installs values one at a time and does not roll back a // prefix when a later value fails, so refuse to execute or read @@ -3429,7 +3422,7 @@ struct PyMethod final { outputs_valid_ = false; handed_outputs_.clear(); for (const auto& input : borrowed_inputs_) { - input.validate(); + input.validate(storage_->method->method_meta().name()); } const auto num_outputs = storage_->method->outputs_size(); allocate_output_storages(); @@ -3457,15 +3450,7 @@ struct PyMethod final { #endif setup_output_storage(*storage_->method, output_storage_spans); #ifndef USE_ATEN_LIB - 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); + synchronize_tensor_devices(*storage_->method); #endif Error execute_status = storage_->method->execute(); THROW_IF_ERROR( @@ -3473,7 +3458,7 @@ struct PyMethod final { "method->execute() failed with error 0x%" PRIx32, static_cast(execute_status)); #ifndef USE_ATEN_LIB - synchronize_device_outputs(*storage_->method); + synchronize_tensor_devices(*storage_->method); #endif outputs_valid_ = true; } @@ -3625,7 +3610,10 @@ struct PyMethod final { py::str(source.attr("dtype")).cast()) != tensor.scalar_type()) { throw py::value_error( - "Input storage changed before get_outputs(); call set_inputs() again"); + "Input " + std::to_string(input->index()) + " of method '" + + storage_->method->method_meta().name() + + "': its storage changed before the output was read; pass it " + "again"); } const auto offset = source.attr("storage_offset")().cast() + static_cast((address - start) / itemsize); diff --git a/extension/pybindings/pybindings_cuda.h b/extension/pybindings/pybindings_cuda.h index a61a7241cff..550a26a85db 100644 --- a/extension/pybindings/pybindings_cuda.h +++ b/extension/pybindings/pybindings_cuda.h @@ -45,8 +45,8 @@ inline std::shared_ptr allocate_cuda_result(size_t nbytes, int index) { #endif } -// The Python boundary is synchronous until the runtime exposes readiness -// events for inputs, output allocations, and caller-selected execution streams. +// The Python API has no way to pass a stream or a readiness event, so the +// Python boundary is synchronous and waits for the whole device. inline void synchronize_cuda_device(int index) { #ifdef EXECUTORCH_PYBINDINGS_CUDA auto guard = executorch::extension::cuda::CUDAGuard::create(index); diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 3a919f60aba..d0c7f90bec0 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -287,7 +287,7 @@ class PyExecuTorchResult final py::capsule to_dlpack(const py::object& stream = py::none()) { // 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. + // or an upload that is still running. if (!device_.is_cpu() && nbytes_ != 0) { synchronize_device(device_.index()); } diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 79d88fa73b2..95b07ad1880 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -418,10 +418,50 @@ def test_numpy_float_scalar_set_inputs(self): def test_method_rejects_tensor_for_scalar_input(self): exported_program, inputs = create_program(ModuleAddScalar()) + module = self.load_fn(exported_program.buffer) method = self.load_prog_fn(exported_program.buffer).load_method("forward") - with self.assertRaises(RuntimeError): - method.set_inputs([inputs[0], torch.ones(1)]) + for run in (module.forward, method.set_inputs): + with self.subTest(run=run): + with self.assertRaisesRegex( + RuntimeError, + r"^Input 1 of method 'forward': the method expects Double, " + r"but got $", + ): + run([inputs[0], torch.ones(1)]) + + def test_method_rejects_aliased_input_whose_storage_changed(self): + # An output that is the input itself is returned as a view of the + # caller's tensor, which must still be the storage the method read. + if self.kernel_mode != "portable": + self.skipTest("only portable bindings return a view of the input") + from executorch.extension.pybindings import _C + + if _C._tensor_output != "torch": + self.skipTest("only torch output mode returns a view of the input") + + class Identity(torch.nn.Module): + def forward(self, x): + return x + + value = torch.ones(2, 2) + program = to_edge(export(Identity(), (value,))).to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ) + ) + ) + method = self.load_prog_fn(program.buffer).load_method("forward") + method.set_inputs([value]) + method.execute() + value.set_(torch.zeros(2, 2)) + with self.assertRaisesRegex( + ValueError, + r"^Input 0 of method 'forward': its storage changed before the output " + r"was read; pass it again$", + ): + method.get_outputs(clone_outputs=False) def test_dlpack_inputs(self): exported_program, inputs = create_program(ModuleAdd()) @@ -762,6 +802,13 @@ def test_bfloat16_raw_buffer_inputs(self): ): with self.assertRaisesRegex(BufferError, "contiguous byte buffers"): executor([noncontiguous, raw_inputs[1]]) + with self.assertRaisesRegex( + ValueError, + "^Input 0 of method 'forward': raw bfloat16 input has " + f"{len(raw_inputs[0]) - 2} bytes, but the method expects " + f"{len(raw_inputs[0])}$", + ): + executor([raw_inputs[0][:-2], raw_inputs[1]]) def test_numpy_array_is_a_single_input(self): exported_program, inputs = create_program(ModuleAddSingleInput()) @@ -962,14 +1009,12 @@ def __str__(self): with RedirectedStderr() as out: try: exported_program, inputs = create_program(ModuleAdd()) - method = self.load_prog_fn(exported_program.buffer).load_method( - "forward" - ) - method.set_inputs((*inputs, 1)) + executorch_module = self.load_fn(exported_program.buffer) + executorch_module.run_method("missing", inputs) self.assertFalse(True) # should be unreachable except Exception: # The runtime's own log, routed to Python's stderr. - self.assertIn("Invalid number of inputs provided", str(out)) + self.assertIn("No method named 'missing' in program", str(out)) def test_quantized_ops(self): eager_module = ModuleAdd() @@ -1582,7 +1627,8 @@ def test_method_rejects_resized_borrowed_input(self): values[0].resize((64, 64), refcheck=False) with self.assertRaisesRegex( - ValueError, "Input 0 storage changed.*call set_inputs" + ValueError, + "Input 0 of method 'forward': its storage changed.*pass it again", ): method.execute() @@ -1621,14 +1667,12 @@ def __str__(self): with RedirectedStderr() as out: try: - program, inputs = create_program(ModuleAdd()) + program, _ = create_program(ModuleAdd()) executorch_program = self.load_prog_fn(program.buffer) - executorch_method = executorch_program.load_method("forward") - inputs = (*inputs, 1) - executorch_output = executorch_method(inputs)[0] # noqa + executorch_program.load_method("missing") self.assertFalse(True) # should be unreachable except Exception: - self.assertIn("Invalid number of inputs provided", str(out)) + self.assertIn("No method named 'missing' in program", str(out)) def test_method_quantized_ops(self): eager_module = ModuleAdd() @@ -2679,9 +2723,9 @@ def forward(self, x): cloned.data_ptr(), value.data_ptr() ) - def test_torch_input_device_must_match_method(self): + def test_input_device_must_match_method(self): if self.kernel_mode != "portable": - self.skipTest("only portable bindings inspect Torch device metadata") + self.skipTest("only portable bindings inspect input device metadata") from executorch.exir._serialize._program import PTEFile, serialize_pte_binary from executorch.exir.schema import ExtraTensorInfo, Tensor @@ -2705,9 +2749,22 @@ def forward(self, x): data = bytes(serialize_pte_binary(PTEFile(schema))) module = self.load_fn(data) method = self.load_prog_fn(data).load_method("forward") + # A host buffer would otherwise be read as device memory. for executor in (module, method): - with self.assertRaisesRegex(ValueError, "cpu.*expects cuda:3"): - executor(value) + for host_input in ( + value, + DLPackOnly(value), + value.numpy(), + memoryview(value.numpy()), + ): + with self.subTest(executor=type(executor), input=type(host_input)): + with self.assertRaisesRegex( + ValueError, + r"^Input 0 of method 'forward' is on device cpu, but the " + r"method expects cuda:3\. Pass device memory instead, such " + r"as ExecuTorchResult\.from_host\(\.\.\.\) or a CUDA tensor$", + ): + executor(host_input) def test_device_planned_method_allocates_on_the_device(self): # The other device test covers the refusal. This one covers what the diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index 990b4bfbfae..86d2680f7cf 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -80,7 +80,9 @@ def setUpClass(cls): ) @classmethod - def _export_delegate(cls, name, inputs_on_device=True): + def _export_delegate( + cls, name, inputs_on_device=True, outputs_on_device=True, plan_outputs=False + ): 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 @@ -107,11 +109,11 @@ def forward(self, value): config=ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( skip_h2d_for_method_inputs=inputs_on_device, - skip_d2h_for_method_outputs=True, + skip_d2h_for_method_outputs=outputs_on_device, ), enable_non_cpu_memory_planning=True, memory_planning_pass=MemoryPlanningPass( - alloc_graph_input=False, alloc_graph_output=False + alloc_graph_input=False, alloc_graph_output=plan_outputs ), ) ) @@ -120,14 +122,14 @@ def forward(self, value): raise AssertionError("test model must execute a CUDA delegate") for index in (*plan.inputs, *plan.outputs): tensor = plan.values[index].val - if tensor.allocation_info is not None: - raise AssertionError( - "test must cover borrowed inputs and unplanned outputs" - ) - 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") + is_output = index in plan.outputs + on_device = outputs_on_device if is_output else inputs_on_device + info = tensor.extra_tensor_info + if on_device != (info is not None and info.device_type == DeviceType.CUDA): + raise AssertionError("test program has a wrong device") + planned = tensor.allocation_info is not None + if planned != (plan_outputs and is_output): + raise AssertionError("test program is planned differently") directory = Path(cls.directory.name) / name directory.mkdir() @@ -140,14 +142,10 @@ def forward(self, value): directory / (data_names[0] + ".ptd") if data_names else None ) - def _load_method(self, host_inputs=False): + def _load_method(self, program=None): from executorch.runtime import Runtime - program_path, data_path = ( - self.host_input_program - if host_inputs - else (self.program_path, self.data_path) - ) + program_path, data_path = program or (self.program_path, self.data_path) return ( Runtime.get() .load_program(program_path, data_path=data_path) @@ -258,26 +256,28 @@ def test_borrowed_result_written_by_another_library_is_waited_for(self): # 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)] + # The reader's outputs are on the host, so only the wait for its input + # can order the read. Its first run sets it up, allocates its outputs, + # and waits for the whole device, so it runs once before the write. + reader = self._load_method( + self._export_delegate( + "host_outputs", outputs_on_device=False, plan_outputs=True + ) ) + reader(torch.zeros(self.shape, device=self.device)) 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) + torch.cuda._sleep(1_000_000_000) shared.fill_(41.0) - self._assert_ready_value(module.forward([passed])[0], 84.0) - del kept + output = reader(passed)[0] + writer.synchronize() + torch.testing.assert_close( + torch.from_dlpack(output), torch.full(self.shape, 84.0) + ) def test_clone_roundtrip_and_borrowed_output_lifetime(self): method = self._load_method() @@ -578,7 +578,9 @@ def on_host(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) + method = self._load_method( + self.host_input_program if host_inputs else None + ) first, second = make_input(1.0), make_input(5.0) torch.cuda.synchronize(self.device) result = method(first)[0] @@ -595,6 +597,32 @@ def on_host(value): reader.synchronize() torch.testing.assert_close(snapshot.cpu(), torch.full(self.shape, 4.0)) + def test_planned_output_memory_waits_for_pending_readers(self): + # Each run writes a memory-planned output again, so it must wait for a + # read another library still has queued there. The inputs are on the + # host, so only the wait for the output's device can order it. + program_path, data_path = self._export_delegate( + "planned_outputs", inputs_on_device=False, plan_outputs=True + ) + module = self.bindings._load_for_executorch_from_buffer( + program_path.read_bytes(), data_path.read_bytes() if data_path else None + ) + method = self._load_method((program_path, data_path)) + for name, run in (("method", method), ("module", module.forward)): + with self.subTest(executor=name): + result = run(torch.full(self.shape, 1.0), clone_outputs=False)[0] + # Imported before the read is queued: an ExecuTorchResult waits + # for the device in __dlpack__, which would finish the read. + view = self._tensor(result) + snapshot = torch.empty_like(view) + reader = torch.cuda.Stream(self.device) + with torch.cuda.stream(reader): + torch.cuda._sleep(1_000_000_000) + snapshot.copy_(view) + run(torch.full(self.shape, 5.0), clone_outputs=False) + 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() value = torch.full(self.shape, 1.0, device=self.device) diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index b5d08deae35..df1e43fc0eb 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -343,23 +343,38 @@ def __dlpack_device__(self): def __dlpack__(self, stream=None): raise UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte") + class FailingDevice(np.ndarray): + def __dlpack_device__(self): + 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): + # A Python error keeps its type, even one whose constructor takes + # more than a message, and carries the name as a note. Choosing how + # to read an input can raise too. + for failing in (FailingProducer(), value.view(FailingDevice)): + with self.subTest(executor=type(executor), input=type(failing)): + with self.assertRaises(UnicodeDecodeError) as caught: + executor((failing, value)) + self.assertIn( + "Input 0 of method 'forward'", + getattr(caught.exception, "__notes__", []), + ) 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 0 of method 'forward': unsupported python type", + ): + executor((object(), value)) + # The count is checked before any input is read. with self.assertRaisesRegex( - RuntimeError, "Input 2 of method 'forward': not a tensor input" + RuntimeError, r"^Method 'forward' takes 2 inputs, but got 3$" ): - method((value, value, value)) + method((value, value, object())) def test_mixed_dlpack_and_buffer_inputs(self) -> None: with open(os.environ["EXECUTORCH_PYBIND_TEST_PTE"], "rb") as program_file: From cc9a0e1bb9e29e266254c1a2f0936e49c9270ab6 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 08:54:44 -0700 Subject: [PATCH 02/10] [PyBindings] Read complex32 DLPack inputs A complex32 result of these bindings exports through DLPack, but the same dtype passed back in, or from another library, was rejected as an unsupported DLPack dtype. DLPack complex with 32 bits now reads as ComplexHalf, matching the export. --- extension/pybindings/pybindings.cpp | 3 +++ extension/pybindings/test/test_pybindings.py | 23 ++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index b7d9b35c6ee..9e55645f8b6 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -1142,6 +1142,9 @@ class DLPackTensorView final { if (dtype.code == dlpack::DataTypeCode::BFloat && dtype.bits == 16) { return ScalarType::BFloat16; } + if (dtype.code == dlpack::DataTypeCode::Complex && dtype.bits == 32) { + return ScalarType::ComplexHalf; + } if (dtype.code == dlpack::DataTypeCode::Complex && dtype.bits == 64) { return ScalarType::ComplexFloat; } diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 95b07ad1880..a4d07c62055 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -859,6 +859,29 @@ def forward(self, x): np.testing.assert_array_equal(np.asarray(output), values) + def test_complex_half_dlpack_input(self): + if self.kernel_mode != "portable": + self.skipTest("only portable bindings read DLPack inputs themselves") + + class Identity(torch.nn.Module): + def forward(self, x): + return x + + value = torch.tensor([1 + 2j, 3 - 4j], dtype=torch.complex32) + program = to_edge(export(Identity(), (value,))).to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ) + ) + ) + method = self.load_prog_fn(program.buffer).load_method("forward") + output = torch.from_dlpack(method(DLPackOnly(value))[0]) + self.assertEqual(output.dtype, torch.complex32) + self.assertTrue( + torch.equal(output.view(torch.float16), value.view(torch.float16)) + ) + def test_unsupported_buffer_format_is_rejected(self): exported_program, _ = create_program(ModuleAddSingleInput()) executorch_module = self.load_fn(exported_program.buffer) From ee6462896f6679f77c1b0445b748d9672c032e61 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:01:21 -0700 Subject: [PATCH 03/10] [PyBindings] Support the DLPack 1.0 keywords on ExecuTorchResult ExecuTorchResult.__dlpack__ took only stream and discarded it. NumPy and other current consumers pass max_version, dl_device and copy, so calls such as np.from_dlpack(result, copy=False) failed with a TypeError, and a consumer that ordered the work itself (stream=-1) still waited for the whole device. __dlpack__ now takes the keyword-only max_version, dl_device and copy. It returns a versioned dltensor_versioned capsule when the consumer asks for DLPack 1.0, copies when asked or when the requested device differs, and raises BufferError when copy=False rules out the copy a request needs. stream=-1 skips the device wait; other streams keep the conservative wait. --- extension/pybindings/README.md | 5 +- extension/pybindings/pybindings.cpp | 6 +- extension/pybindings/pybindings.pyi | 9 +- extension/pybindings/pybindings_dlpack.h | 17 +++ .../pybindings/pybindings_executorch_result.h | 116 ++++++++++++++++-- .../test/test_pybindings_cuda_dlpack.py | 32 +++++ .../test/test_pybindings_no_aten.py | 24 ++++ 7 files changed, 199 insertions(+), 10 deletions(-) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 04a019d2bc1..fc79d5b41a3 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -93,7 +93,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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 9e55645f8b6..0705584a109 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -4254,7 +4254,11 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { .def( "__dlpack__", &PyExecuTorchResult::to_dlpack, - py::arg("stream") = py::none()) + py::arg("stream") = py::none(), + py::kw_only(), + py::arg("max_version") = py::none(), + py::arg("dl_device") = py::none(), + py::arg("copy") = py::none()) .def("__dlpack_device__", &PyExecuTorchResult::dlpack_device) .def( "to_host", diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index e55eecb1e5e..9d376133917 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -50,7 +50,14 @@ class ExecuTorchResult: @property def nbytes(self) -> int: ... def __array__(self, dtype: Any = None, copy: Any = None) -> Any: ... - def __dlpack__(self, stream: Optional[int] = None) -> Any: ... + def __dlpack__( + self, + stream: Optional[int] = None, + *, + max_version: Optional[Tuple[int, int]] = None, + dl_device: Optional[Tuple[int, int]] = None, + copy: Optional[bool] = None, + ) -> Any: ... def __dlpack_device__(self) -> Tuple[int, int]: ... def to_host(self) -> ExecuTorchResult: """Return a copy of this result in CPU memory.""" diff --git a/extension/pybindings/pybindings_dlpack.h b/extension/pybindings/pybindings_dlpack.h index e0dac1afc1d..3e73fe78812 100644 --- a/extension/pybindings/pybindings_dlpack.h +++ b/extension/pybindings/pybindings_dlpack.h @@ -60,6 +60,23 @@ struct ManagedTensor { void (*deleter)(ManagedTensor* self); }; +struct Version { + uint32_t major; + uint32_t minor; +}; + +// DLPack 1.0 capsules carry this struct, named "dltensor_versioned". +struct ManagedTensorVersioned { + Version version; + void* manager_ctx; + void (*deleter)(ManagedTensorVersioned* self); + uint64_t flags; + Tensor dl_tensor; +}; + +constexpr uint64_t kFlagReadOnly = 1ULL << 0; +constexpr uint64_t kFlagIsCopied = 1ULL << 1; + } // namespace executorch::extension::pybindings::dlpack namespace executorch::extension::pybindings { diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index d0c7f90bec0..63c3e855882 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -284,15 +284,68 @@ class PyExecuTorchResult final static_cast(device.device_type), device.device_id); } - py::capsule to_dlpack(const py::object& stream = py::none()) { - // 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. - if (!device_.is_cpu() && nbytes_ != 0) { + // The DLPack __dlpack__ protocol, with the keywords of DLPack 1.0. + py::capsule to_dlpack( + const py::object& stream, + const py::object& max_version, + const py::object& requested_device, + const py::object& copy) { + if (!copy.is_none() && !py::isinstance(copy)) { + throw py::type_error("copy must be True, False, or None"); + } + if (!stream.is_none() && + (!py::isinstance(stream) || + py::isinstance(stream))) { + throw py::type_error("stream must be an integer or None"); + } + if (!device_.is_cpu() && !stream.is_none() && stream.cast() == 0) { + throw py::value_error( + "stream 0 is ambiguous for CUDA; pass 1, 2, -1, or a stream handle"); + } + const auto own = dl_device(); + auto target = own; + if (!requested_device.is_none()) { + const auto requested = requested_device.cast>(); + target = { + static_cast(requested.first), requested.second}; + } + const bool move = target.device_type != own.device_type || + target.device_id != own.device_id; + if (move && !copy.is_none() && !copy.cast()) { + throw py::buffer_error( + "Moving an ExecuTorchResult to another device needs a copy"); + } + // 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. A consumer that passes -1 orders the work itself. + if (!device_.is_cpu() && nbytes_ != 0 && + (stream.is_none() || stream.cast() != -1)) { synchronize_device(device_.index()); } - (void)stream; - return capsule(); + auto exported = shared_from_this(); + uint64_t flags = 0; + if (move || (!copy.is_none() && copy.cast())) { + exported = copy_to(target); + flags |= dlpack::kFlagIsCopied; + } + return wants_versioned_capsule(max_version) + ? exported->versioned_capsule(flags) + : exported->capsule(); + } + + py::capsule versioned_capsule(uint64_t flags) { + auto holder = std::make_unique(shared_from_this()); + holder->versioned.flags = flags; + py::capsule capsule( + &holder->versioned, "dltensor_versioned", [](PyObject* capsule) { + if (PyCapsule_IsValid(capsule, "dltensor_versioned")) { + auto* managed = static_cast( + PyCapsule_GetPointer(capsule, "dltensor_versioned")); + managed->deleter(managed); + } + }); + holder.release(); + return capsule; } // A capsule for these bindings to read the result as a method input. It @@ -450,12 +503,17 @@ class PyExecuTorchResult final 0}; managed.manager_ctx = this; managed.deleter = release_dlpack_holder; + versioned.version = {1, 0}; + versioned.manager_ctx = this; + versioned.deleter = release_versioned_dlpack_holder; + versioned.dl_tensor = managed.dl_tensor; } std::shared_ptr result; std::vector shape; std::vector strides; dlpack::ManagedTensor managed{}; + dlpack::ManagedTensorVersioned versioned{}; }; static void release_dlpack_holder(dlpack::ManagedTensor* managed) { @@ -465,6 +523,50 @@ class PyExecuTorchResult final delete static_cast(managed->manager_ctx); } + static void release_versioned_dlpack_holder( + dlpack::ManagedTensorVersioned* managed) { + py::gil_scoped_acquire gil; + delete static_cast(managed->manager_ctx); + } + + static bool wants_versioned_capsule(const py::object& max_version) { + if (max_version.is_none()) { + return false; + } + const auto version = max_version.cast>(); + return version.first >= 1; + } + + // A copy on `target`, for a DLPack consumer that asks for one. + std::shared_ptr copy_to(dlpack::Device target) const { + if (target.device_type == dlpack::DeviceType::CPU && + target.device_id == 0) { + return to_host(); + } + if (target.device_type != dlpack::DeviceType::CUDA) { + throw py::buffer_error( + "ExecuTorchResult copies only to CPU and CUDA devices"); + } + if (device_.is_cpu()) { + return to_device(target.device_id); + } + if (target.device_id != device_.index()) { + throw py::buffer_error( + "ExecuTorchResult does not copy between CUDA devices"); + } + auto result = std::shared_ptr( + new PyExecuTorchResult(sizes_, strides_, scalar_type_, device_)); + if (nbytes_ != 0) { + result->device_storage_ = allocate_cuda_result(nbytes_, device_.index()); + copy_cuda_storage( + result->device_storage_.get(), data_, nbytes_, device_.index()); + // The consumer may read the copy on any stream. + synchronize_cuda_device(device_.index()); + result->data_ = result->device_storage_.get(); + } + return result; + } + dlpack::Device dl_device() const { return { device_.is_cpu() ? dlpack::DeviceType::CPU : dlpack::DeviceType::CUDA, diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index 86d2680f7cf..947803cb6e4 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -745,6 +745,38 @@ 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 _device_waits(self, body): + from torch.profiler import profile, ProfilerActivity + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as trace: + body() + return sum( + event.count + for event in trace.key_averages() + if event.key == "cudaDeviceSynchronize" + ) + + def test_dlpack_export_keywords(self): + if not self.executorch_output: + self.skipTest("__dlpack__ keywords are an ExecuTorchResult feature") + method = self._load_method() + result = method(torch.full(self.shape, 1.0, device=self.device))[0] + # The profiler waits for the device itself when it stops. + idle = self._device_waits(lambda: None) + # A consumer that passes stream=-1 orders the work itself. + self.assertEqual(self._device_waits(lambda: result.__dlpack__(stream=-1)), idle) + self.assertEqual( + self._device_waits(lambda: result.__dlpack__(stream=1)), idle + 1 + ) + copied = torch.from_dlpack(result.__dlpack__(max_version=(1, 0), copy=True)) + self.assertNotEqual(copied.data_ptr(), torch.from_dlpack(result).data_ptr()) + torch.testing.assert_close(copied.cpu(), torch.full(self.shape, 4.0)) + host = torch.from_dlpack( + result.__dlpack__(max_version=(1, 0), dl_device=(1, 0), copy=True) + ) + self.assertEqual(host.device, torch.device("cpu")) + torch.testing.assert_close(host, torch.full(self.shape, 4.0)) + 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 diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index df1e43fc0eb..2675b8b5ae0 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -272,6 +272,30 @@ def test_dlpack_input_and_output(self) -> None: self.assertEqual(output.__dlpack_device__(), (1, 0)) np.testing.assert_array_equal(np.from_dlpack(output), np.array([3.0])) + def test_dlpack_export_keywords(self) -> None: + with open(os.environ["EXECUTORCH_PYBIND_TEST_PTE"], "rb") as program_file: + program_data = program_file.read() + method = runtime._load_program_from_buffer(program_data).load_method("forward") + values = np.array([1.0], dtype=np.float32) + result = method((values, values))[0] + if not isinstance(result, runtime.ExecuTorchResult): + self.skipTest("torch output mode returns torch tensors") + + shared = np.from_dlpack(result, copy=False) + self.assertTrue(np.shares_memory(shared, np.from_dlpack(result))) + copied = np.from_dlpack(result, copy=True) + self.assertFalse(np.shares_memory(copied, shared)) + np.testing.assert_array_equal(copied, shared) + np.testing.assert_array_equal(np.from_dlpack(result, device="cpu"), shared) + self.assertIn( + '"dltensor_versioned"', repr(result.__dlpack__(max_version=(1, 0))) + ) + self.assertIn('"dltensor"', repr(result.__dlpack__())) + with self.assertRaises(BufferError): + result.__dlpack__(dl_device=(2, 0), copy=False) + with self.assertRaises(TypeError): + result.__dlpack__(stream="default") + def test_bfloat16_raw_buffer_input_and_output(self) -> None: with open( os.environ["EXECUTORCH_PYBIND_TEST_BFLOAT16_PTE"], "rb" From b358fd0b6acd7ab61b6a5f993bfeb7ad94a3a3bd Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:03:04 -0700 Subject: [PATCH 04/10] [PyBindings] Accept raw DLPack capsules as inputs A DLPack capsule, such as the one array.__dlpack__() or a result's __dlpack__() returns, was rejected as an input. Only objects with a __dlpack__ method worked. A dltensor or dltensor_versioned capsule is now accepted on its own or in the input list. Its device comes from the tensor it holds. The capsule is consumed once: it is renamed after a successful read, and passing it again raises ValueError. A read-only DLPack 1.0 CPU tensor is read from a copy, since a method may write its inputs. A producer object is now asked for a DLPack 1.0 capsule first, so it can say when its memory is read-only, with a fallback for producers that do not take max_version. --- extension/pybindings/README.md | 4 +- extension/pybindings/pybindings.cpp | 136 +++++++++++++++--- .../pybindings/pybindings_executorch_result.h | 17 +++ extension/pybindings/test/test_pybindings.py | 11 ++ .../test/test_pybindings_no_aten.py | 45 ++++++ 5 files changed, 192 insertions(+), 21 deletions(-) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index fc79d5b41a3..90f9dca8882 100644 --- a/extension/pybindings/README.md +++ b/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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 0705584a109..8850b84a49b 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -959,13 +959,38 @@ class BorrowedInput final { }; #ifndef USE_ATEN_LIB +// Asks for a DLPack 1.0 capsule, which can mark memory read-only, and falls +// back to a legacy capsule for producers that predate it. +py::object dlpack_capsule( + const py::object& producer, + const py::object& stream) { + const auto request = [&](bool versioned) { + py::dict arguments; + if (!stream.is_none()) { + arguments["stream"] = stream; + } + if (versioned) { + arguments["max_version"] = py::make_tuple(1, 0); + } + return producer.attr("__dlpack__")(**arguments); + }; + try { + return request(/*versioned=*/true); + } catch (py::error_already_set& error) { + if (!error.matches(PyExc_TypeError)) { + throw; + } + } + return request(/*versioned=*/false); +} + // Asks a producer for a capsule of CUDA memory on device `index`. Producers // such as torch export a CUDA tensor only while its device is the current one, // and torch reads that from its own state, so a torch that can use CUDA is // asked to switch. py::object cuda_dlpack_capsule(const py::object& producer, int index) { const auto request = [&producer] { - return producer.attr("__dlpack__")(py::arg("stream") = 1); + return dlpack_capsule(producer, py::int_(1)); }; PyObject* torch = loaded_torch_module(); const auto torch_cuda = torch == nullptr @@ -986,13 +1011,26 @@ py::object cuda_dlpack_capsule(const py::object& producer, int index) { } } -/** A borrowed tensor view obtained through the producer's DLPack protocol. */ +/** A borrowed tensor view obtained through the producer's DLPack protocol, + * or from a DLPack capsule passed directly. */ class DLPackTensorView final { public: explicit DLPackTensorView(const py::handle& value) - : owner_(py::reinterpret_borrow(value)), - advertised_device_( - owner_.attr("__dlpack_device__")().cast>()) { + : owner_(py::reinterpret_borrow(value)) { + if (PyCapsule_CheckExact(owner_.ptr())) { + // A capsule carries its device in the tensor it holds, and has no + // producer to negotiate a stream with. + capsule_ = owner_; + const auto& tensor = take_capsule(); + advertised_device_ = { + static_cast(tensor.device.device_type), tensor.device.device_id}; + validate_advertised_device(); + read(tensor); + mark_used(); + return; + } + advertised_device_ = + owner_.attr("__dlpack_device__")().cast>(); validate_advertised_device(); // A result of these bindings gives its capsule without the wait in // __dlpack__, since the call waits for every CUDA input before it runs. @@ -1001,26 +1039,19 @@ class DLPackTensorView final { } else { capsule_ = advertised_device_.first == dlpack::CUDA ? cuda_dlpack_capsule(owner_, advertised_device_.second) - : owner_.attr("__dlpack__")(); - } - if (!PyCapsule_IsValid(capsule_.ptr(), "dltensor")) { - throw py::value_error("__dlpack__() did not return a dltensor capsule"); - } - managed_ = static_cast( - PyCapsule_GetPointer(capsule_.ptr(), "dltensor")); - if (managed_ == nullptr) { - throw py::error_already_set(); - } - initialize(managed_->dl_tensor); - if (PyCapsule_SetName(capsule_.ptr(), "used_dltensor") != 0) { - throw py::error_already_set(); + : dlpack_capsule(owner_, py::none()); } + read(take_capsule()); + mark_used(); } ~DLPackTensorView() { if (managed_ != nullptr && managed_->deleter != nullptr) { managed_->deleter(managed_); } + if (versioned_ != nullptr && versioned_->deleter != nullptr) { + versioned_->deleter(versioned_); + } } DLPackTensorView(const DLPackTensorView&) = delete; @@ -1088,7 +1119,9 @@ class DLPackTensorView final { if (!converted_.empty()) { return py::none(); } - auto capsule_owner = PyExecuTorchResult::dlpack_backing_owner(managed_); + auto capsule_owner = versioned_ != nullptr + ? PyExecuTorchResult::dlpack_backing_owner(versioned_) + : PyExecuTorchResult::dlpack_backing_owner(managed_); if (!capsule_owner.is_none()) { return capsule_owner; } @@ -1100,6 +1133,65 @@ class DLPackTensorView final { } private: + // The tensor in capsule_, legacy or DLPack 1.0. The capsule keeps owning it + // until mark_used(), so a failure before then lets the capsule free it. + const dlpack::Tensor& take_capsule() { + if (PyCapsule_IsValid(capsule_.ptr(), "dltensor")) { + managed_ = static_cast( + PyCapsule_GetPointer(capsule_.ptr(), "dltensor")); + if (managed_ == nullptr) { + throw py::error_already_set(); + } + return managed_->dl_tensor; + } + if (PyCapsule_IsValid(capsule_.ptr(), "dltensor_versioned")) { + auto* versioned = static_cast( + PyCapsule_GetPointer(capsule_.ptr(), "dltensor_versioned")); + if (versioned == nullptr) { + throw py::error_already_set(); + } + if (versioned->version.major != 1) { + throw py::value_error( + "DLPack version " + std::to_string(versioned->version.major) + + " is not supported"); + } + versioned_ = versioned; + return versioned_->dl_tensor; + } + if (PyCapsule_IsValid(capsule_.ptr(), "used_dltensor") || + PyCapsule_IsValid(capsule_.ptr(), "used_dltensor_versioned")) { + throw py::value_error("This DLPack capsule was already consumed"); + } + throw py::value_error("Expected a dltensor or dltensor_versioned capsule"); + } + + // Takes the tensor over from the capsule, so nothing else can take it again. + void mark_used() { + const auto* name = + versioned_ != nullptr ? "used_dltensor_versioned" : "used_dltensor"; + if (PyCapsule_SetName(capsule_.ptr(), name) != 0) { + managed_ = nullptr; + versioned_ = nullptr; + throw py::error_already_set(); + } + } + + void read(const dlpack::Tensor& tensor) { + initialize(tensor); + if (versioned_ == nullptr || + (versioned_->flags & dlpack::kFlagReadOnly) == 0 || nbytes_ == 0) { + return; + } + // A method may write its inputs, so read-only memory is read from a copy. + if (!device_.is_cpu()) { + throw py::value_error("Read-only DLPack device inputs are not supported"); + } + converted_.assign( + static_cast(data_), + static_cast(data_) + nbytes_); + data_ = converted_.data(); + } + static executorch::aten::ScalarType scalar_type_from_dlpack( const dlpack::DataType& dtype) { if (dtype.lanes != 1) { @@ -1280,6 +1372,7 @@ class DLPackTensorView final { size_t nbytes_ = 0; py::object capsule_; dlpack::ManagedTensor* managed_ = nullptr; + dlpack::ManagedTensorVersioned* versioned_ = nullptr; void* data_ = nullptr; std::vector converted_; std::vector sizes_; @@ -1470,6 +1563,9 @@ std::optional scalar_input( } bool prefer_dlpack_input(const py::handle& value) { + if (PyCapsule_CheckExact(value.ptr())) { + return true; + } if (!py::hasattr(value, "__dlpack__")) { return false; } @@ -1494,7 +1590,7 @@ bool prefer_dlpack_input(const py::handle& value) { py::sequence normalize_inputs(const py::object& inputs) { if (PyObject_CheckBuffer(inputs.ptr()) || is_torch_tensor(inputs) || - py::hasattr(inputs, "__dlpack__")) { + py::hasattr(inputs, "__dlpack__") || PyCapsule_CheckExact(inputs.ptr())) { py::list result; result.append(inputs); return result; diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 63c3e855882..5e0e686a22e 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -278,6 +278,23 @@ class PyExecuTorchResult final return result->owner_.is_none() ? py::cast(result) : result->owner_; } + static py::object dlpack_backing_owner( + dlpack::ManagedTensorVersioned* managed) { + if (managed == nullptr || + managed->deleter != release_versioned_dlpack_holder || + managed->manager_ctx == nullptr) { + return py::none(); + } + const auto* holder = static_cast(managed->manager_ctx); + if (&holder->versioned != managed || + managed->dl_tensor.data != holder->result->data_ || + managed->dl_tensor.byte_offset != 0) { + return py::none(); + } + const auto& result = holder->result; + return result->owner_.is_none() ? py::cast(result) : result->owner_; + } + py::tuple dlpack_device() const { const auto device = dl_device(); return py::make_tuple( diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index a4d07c62055..1abf203b2c9 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -836,6 +836,17 @@ def test_mixed_tensor_protocols(self): output = executor(mixed_inputs)[0] self.assertTrue(torch.allclose(output, inputs[0] + inputs[1])) + def test_raw_dlpack_capsule_is_a_single_input(self): + exported_program, inputs = create_program(ModuleAddSingleInput()) + module = self.load_fn(exported_program.buffer) + method = self.load_prog_fn(exported_program.buffer).load_method("forward") + for executor in (module, method): + with self.subTest(executor=type(executor)): + output = executor(inputs[0].numpy().__dlpack__())[0] + torch.testing.assert_close( + torch.from_dlpack(output), inputs[0] + inputs[0] + ) + def test_non_dense_numpy_input_is_rejected(self): exported_program, inputs = create_program(ModuleAddSingleInput()) executorch_module = self.load_fn(exported_program.buffer) diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index 2675b8b5ae0..e0513716f2a 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -296,6 +296,51 @@ def test_dlpack_export_keywords(self) -> None: with self.assertRaises(TypeError): result.__dlpack__(stream="default") + def test_raw_dlpack_capsule_inputs(self) -> None: + 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") + first = np.array([1.0], dtype=np.float32) + second = np.array([2.0], dtype=np.float32) + read_only = second.copy() + read_only.flags.writeable = False + for executor in (module, method): + for capsules in ( + (first.__dlpack__(), second.__dlpack__()), + ( + first.__dlpack__(max_version=(1, 0)), + read_only.__dlpack__(max_version=(1, 0)), + ), + ): + with self.subTest(executor=type(executor), capsules=capsules): + output = executor(capsules)[0] + np.testing.assert_array_equal(np.from_dlpack(output), [3.0]) + + # A producer that knows DLPack 1.0 is asked for a versioned capsule. + class VersionedProducer: + def __init__(self, array: np.ndarray) -> None: + self.array = array + self.requests = [] + + def __dlpack_device__(self): + return self.array.__dlpack_device__() + + def __dlpack__(self, stream=None, max_version=None): + self.requests.append(max_version) + return self.array.__dlpack__(stream=stream, max_version=max_version) + + producer = VersionedProducer(read_only) + np.testing.assert_array_equal( + np.from_dlpack(method((first, producer))[0]), [3.0] + ) + self.assertEqual(producer.requests, [(1, 0)]) + # A capsule gives its tensor away once. + used = first.__dlpack__() + method((used, second)) + with self.assertRaisesRegex(ValueError, "already consumed"): + method((used, second)) + def test_bfloat16_raw_buffer_input_and_output(self) -> None: with open( os.environ["EXECUTORCH_PYBIND_TEST_BFLOAT16_PTE"], "rb" From 1ec3124e7c3816f8d703c30352e4d18861c0cdd9 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:08:30 -0700 Subject: [PATCH 05/10] [PyBindings] Never return constant attributes as writable memory Method.get_attribute() returned a borrowed view of any attribute, constants included. A constant loaded from a data file can sit in read-only pages, yet the result said its memory was writable, so an ordinary in-place NumPy or torch operation could crash the process. Mutable state lives in the method's planned memory, so an attribute outside it is now treated as a constant. In ExecuTorchResult output mode it is a read-only view: the buffer protocol and NumPy report it read-only, and a DLPack export is a copy, because consumers such as torch ignore the DLPack read-only flag. Only a caller that passes copy=False through DLPack 1.0 shares the memory, marked read-only. Passing the result back in as an input reads a copy. In torch output mode a constant is a copy, since a torch tensor cannot be read-only. --- extension/pybindings/README.md | 6 ++ extension/pybindings/pybindings.cpp | 22 +++++- .../pybindings/pybindings_executorch_result.h | 44 +++++++++--- extension/pybindings/test/test_pybindings.py | 68 +++++++++++++++++++ 4 files changed, 127 insertions(+), 13 deletions(-) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 90f9dca8882..d8c95191eb6 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -65,6 +65,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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 8850b84a49b..d9947dc629c 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -1035,7 +1035,13 @@ class DLPackTensorView final { // 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_)) { - capsule_ = owner_.cast>()->capsule(); + const auto result = owner_.cast>(); + // A method may write its inputs, so read-only memory is read from a + // copy. + capsule_ = result->read_only() + ? result->to_dlpack( + py::none(), py::none(), py::none(), py::bool_(true)) + : result->capsule(); } else { capsule_ = advertised_device_.first == dlpack::CUDA ? cuda_dlpack_capsule(owner_, advertised_device_.second) @@ -3608,10 +3614,22 @@ struct PyMethod final { if (!attr->device().is_cpu() && attr->nbytes() != 0) { synchronize_device(attr->device().index()); } - return portable_tensor_result( + // Mutable state lives in the method's planned memory. Anything else is a + // constant, which can sit in read-only pages, so it is never returned as + // writable: a torch tensor cannot be read-only, so torch gets a copy. + const bool constant = attr->nbytes() != 0 && + !storage_->memory->owns_data(attr->const_data_ptr()); + if (constant && portable_tensor_output() == PortableTensorOutput::Torch) { + return portable_tensor_result(attr.get(), /*clone=*/true); + } + auto result = portable_tensor_result( attr.get(), /*clone=*/false, borrowed_output_owner(attr.get())); + if (constant) { + result.cast().mark_read_only(); + } + return result; #endif } diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 5e0e686a22e..99d81e79394 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -196,7 +196,7 @@ class PyExecuTorchResult final ndim, std::move(shape), std::move(byte_strides), - /*readonly=*/false); + read_only_); } py::buffer_info flat_buffer() { @@ -210,7 +210,7 @@ class PyExecuTorchResult final /*ndim=*/1, {static_cast(nbytes_)}, {1}, - /*readonly=*/false); + read_only_); } py::tuple shape() const { @@ -252,7 +252,21 @@ class PyExecuTorchResult final py::array numpy_array(const py::object& base) { const auto array_dtype = dtype(); auto info = buffer(); - return py::array(array_dtype, info.shape, info.strides, info.ptr, base); + py::array array(array_dtype, info.shape, info.strides, info.ptr, base); + if (read_only_) { + array.attr("setflags")(py::arg("write") = false); + } + return array; + } + + // For memory a method must not write, such as a constant loaded from the + // program or its data file. + void mark_read_only() { + read_only_ = true; + } + + bool read_only() const { + return read_only_; } size_t nbytes() const { @@ -328,9 +342,18 @@ class PyExecuTorchResult final } const bool move = target.device_type != own.device_type || target.device_id != own.device_id; - if (move && !copy.is_none() && !copy.cast()) { + const bool versioned = wants_versioned_capsule(max_version); + // Consumers such as torch ignore the read-only flag, so read-only memory + // is shared only when the caller asks for no copy, and only in a capsule + // that can carry the flag. + const bool share_read_only = + versioned && !copy.is_none() && !copy.cast(); + const bool must_copy = move || (read_only_ && !share_read_only); + if (must_copy && !copy.is_none() && !copy.cast()) { throw py::buffer_error( - "Moving an ExecuTorchResult to another device needs a copy"); + move ? "Moving an ExecuTorchResult to another device needs a copy" + : "A read-only ExecuTorchResult is shared without a copy only " + "through a DLPack 1.0 capsule, which marks it read-only"); } // 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 @@ -340,14 +363,12 @@ class PyExecuTorchResult final synchronize_device(device_.index()); } auto exported = shared_from_this(); - uint64_t flags = 0; - if (move || (!copy.is_none() && copy.cast())) { + uint64_t flags = read_only_ ? dlpack::kFlagReadOnly : 0; + if (must_copy || (!copy.is_none() && copy.cast())) { exported = copy_to(target); - flags |= dlpack::kFlagIsCopied; + flags = dlpack::kFlagIsCopied; } - return wants_versioned_capsule(max_version) - ? exported->versioned_capsule(flags) - : exported->capsule(); + return versioned ? exported->versioned_capsule(flags) : exported->capsule(); } py::capsule versioned_capsule(uint64_t flags) { @@ -658,6 +679,7 @@ class PyExecuTorchResult final size_t nbytes_; executorch::aten::Device device_; py::object owner_; + bool read_only_ = false; }; /** Contiguous buffer facade used to construct a strided torch.Tensor result. */ diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 1abf203b2c9..8bf3d33a35a 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -1969,6 +1969,74 @@ def test_method_method_meta(self) -> None: self.assertEqual(output_tensor.nbytes(), 16) self.assertEqual(str(output_tensor), tensor_info) + def test_constant_attribute_is_never_writable(self): + # A constant from the data file can sit in read-only pages, so a write + # through a borrowed view would crash the process. + if self.kernel_mode != "portable": + self.skipTest("only portable bindings return borrowed attributes") + eager_module = ModuleLinear() + program = to_edge( + export(eager_module, eager_module.get_inputs(), strict=True) + ).to_executorch(ExecutorchBackendConfig(external_constants=True)) + with tempfile.TemporaryDirectory() as directory: + pte_path = os.path.join(directory, "linear.pte") + ptd_path = os.path.join(directory, "linear.ptd") + with open(pte_path, "wb") as file: + file.write(program.buffer) + with open(ptd_path, "wb") as file: + file.write(bytes(program._tensor_data["_default_external_constant"])) + method = self.runtime._load_program( + pte_path, data_path=ptd_path + ).load_method("forward") + weight = method.get_attribute("linear.weight") + expected = eager_module.linear.weight.detach() + if isinstance(weight, torch.Tensor): + # A torch tensor cannot be read-only, so it is a copy. + self.assertNotEqual( + weight.data_ptr(), method.get_attribute("linear.weight").data_ptr() + ) + torch.testing.assert_close(weight, expected) + else: + self.assertTrue(memoryview(weight).readonly) + self.assertFalse(np.asarray(weight).flags.writeable) + np.testing.assert_array_equal(np.asarray(weight), expected.numpy()) + # The method maps the data file, and Windows cannot delete a mapped + # file, so drop it before the directory goes away. + del weight, method + gc.collect() + # torch ignores the DLPack read-only flag, so a DLPack export must be + # a copy that it can write. Checked in ExecuTorchResult output mode. + script = """ +import sys + +import numpy as np +import torch +from executorch.extension.pybindings import portable_lib as runtime + +method = runtime._load_program(sys.argv[1], data_path=sys.argv[2]).load_method( + "forward" +) +weight = method.get_attribute("linear.weight") +assert isinstance(weight, runtime.ExecuTorchResult) +before = np.asarray(weight).copy() +torch.from_dlpack(weight).add_(1) +assert (np.asarray(weight) == before).all() +print("the constant was not written") +""" + completed = subprocess.run( + [sys.executable, "-c", script, pte_path, ptd_path], + env=dict(os.environ, EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT="executorch"), + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual( + completed.returncode, + 0, + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}", + ) + self.assertIn("the constant was not written", completed.stdout) + def test_program_data_separation(self) -> None: eager_module = ModuleLinear() inputs = eager_module.get_inputs() From d10b5afb878659c2d55762c871cb65b2e84df861 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:10:13 -0700 Subject: [PATCH 06/10] [PyBindings] Keep the event tracer alive while a method uses it A Program loaded with enable_etdump=True owned its event tracer and debug buffer, but each Method it loaded kept only a raw pointer to the tracer. So a method from a temporary program, as in load_program(..., enable_etdump=True).load_method("forward"), used freed memory on every execution and could crash. The tracer and its debug buffer now live in the program state that every method already shares, the same way the data map does. --- extension/pybindings/pybindings.cpp | 26 ++++++------ extension/pybindings/test/test_pybindings.py | 43 ++++++++++++++++++++ 2 files changed, 56 insertions(+), 13 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index d9947dc629c..36b27fb92ef 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -2043,6 +2043,10 @@ struct ProgramState final { // the storage. std::vector> backend_option_storage_; LoadBackendOptionsMap backend_options_; + // A method keeps a pointer to the tracer, which logs into the debug buffer, + // so both stay alive as long as any method does. The tracer goes first. + std::unique_ptr debug_buffer_; + std::unique_ptr event_tracer_; explicit ProgramState( std::unique_ptr loader, @@ -4035,8 +4039,8 @@ struct PyProgram final { data_path, std::move(program_buffer), backend_options)), - event_tracer_(std::move(tracer)), debug_buffer_size_(debug_buffer_size) { + state_->event_tracer_ = std::move(tracer); // Figure out the size of each non_const layer we need to support every // method in the program. Map will be easier to use than a list because we // dont know how many non_const arenas there will be @@ -4072,12 +4076,10 @@ struct PyProgram final { } memory_ = std::make_shared(std::move(planned_sizes)); - if (event_tracer_ && debug_buffer_size > 0) { - // If a debug buffer was requested for the ETDump, allocate it and make - // sure its lifetime is as long as the event_tracer. - debug_buffer_ = std::make_unique(debug_buffer_size); - event_tracer_->set_debug_buffer(get_etdump_debug_buffer()); - event_tracer_->set_event_tracer_debug_level( + if (state_->event_tracer_ && debug_buffer_size > 0) { + state_->debug_buffer_ = std::make_unique(debug_buffer_size); + state_->event_tracer_->set_debug_buffer(get_etdump_debug_buffer()); + state_->event_tracer_->set_event_tracer_debug_level( EventTracerDebugLogLevel::kIntermediateOutputs); } } @@ -4158,7 +4160,7 @@ struct PyProgram final { Result res = state_->program_->load_method( method_name.c_str(), memory->mem_manager(), - event_tracer_.get(), + state_->event_tracer_.get(), state_->data_map_.get(), &state_->backend_options_); THROW_IF_ERROR( @@ -4173,7 +4175,7 @@ struct PyProgram final { } Span get_etdump_debug_buffer() { - return Span(debug_buffer_.get(), debug_buffer_size_); + return Span(state_->debug_buffer_.get(), debug_buffer_size_); } std::unique_ptr method_meta(const std::string& method_name) { @@ -4188,7 +4190,7 @@ struct PyProgram final { } bool has_etdump() { - return static_cast(event_tracer_); + return static_cast(state_->event_tracer_); } void write_etdump_result_to_file( @@ -4197,7 +4199,7 @@ struct PyProgram final { if (!has_etdump()) { throw std::runtime_error("No etdump found"); } - auto& etdump = *event_tracer_; + auto& etdump = *state_->event_tracer_; etdump_result result = etdump.get_etdump_data(); if (result.buf != nullptr && result.size > 0) { write_data_to_file(path, result.buf, result.size); @@ -4223,8 +4225,6 @@ struct PyProgram final { private: std::shared_ptr memory_; std::shared_ptr state_; - std::unique_ptr event_tracer_; - std::unique_ptr debug_buffer_; size_t debug_buffer_size_; }; diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 8bf3d33a35a..71b56912ad0 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -2037,6 +2037,49 @@ def test_constant_attribute_is_never_writable(self): ) self.assertIn("the constant was not written", completed.stdout) + def test_method_keeps_its_event_tracer_alive(self): + # The program is dropped right away, while its method keeps running. + # glibc's MALLOC_PERTURB_ overwrites freed memory, so a method still + # pointing at a freed tracer crashes instead of passing by luck. + if self.kernel_mode != "portable" or not sys.platform.startswith("linux"): + self.skipTest("needs the portable bindings and glibc") + script = """ +import gc +import sys + +import torch +from executorch.extension.pybindings import portable_lib as runtime + +with open(sys.argv[1], "rb") as source: + data = source.read() +method = runtime._load_program_from_buffer( + data, enable_etdump=True, debug_buffer_size=1 << 16 +).load_method("forward") +gc.collect() +junk = [bytearray(256) for _ in range(20000)] +for _ in range(5): + method([torch.ones(2, 2), torch.ones(2, 2)]) +print("ran after its program was dropped") +""" + exported_program, _ = create_program(ModuleAdd()) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "add.pte") + with open(path, "wb") as file: + file.write(exported_program.buffer) + completed = subprocess.run( + [sys.executable, "-c", script, path], + env=dict(os.environ, MALLOC_PERTURB_="165"), + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual( + completed.returncode, + 0, + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}", + ) + self.assertIn("ran after its program was dropped", completed.stdout) + def test_program_data_separation(self) -> None: eager_module = ModuleLinear() inputs = eager_module.get_inputs() From 2042929dd700a7f3ccaac5f92eb398740e6c1774 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:12:24 -0700 Subject: [PATCH 07/10] [runtime] Report the device of each method tensor in TensorInfo A torch-free caller had no way to ask whether a method expects an input on the CPU or on a CUDA device, so it had to know how the program was exported. The program already records the device of each tensor. TensorInfo gains device(), read the same way the tensor parser reads it, with tensors that carry no device annotation on the CPU. Inputs, outputs and attributes all report it. TensorInfo has grown before, with name() and is_memory_planned(); it is a value type that only MethodMeta creates. In Python, TensorInfo.device() and the new ExecuTorchResult.device give "cpu" or "cuda:N", the form ExecuTorchResult's repr already used, and TensorInfo's repr includes it. An unknown device type raises instead of being named as CUDA. --- extension/pybindings/README.md | 3 + extension/pybindings/pybindings.cpp | 21 +++++- extension/pybindings/pybindings.pyi | 9 +++ .../pybindings/pybindings_executorch_result.h | 8 ++- extension/pybindings/test/test_pybindings.py | 68 ++++++++++++++++--- .../test/test_pybindings_no_aten.py | 1 + runtime/executor/method_meta.cpp | 47 ++++++++++--- runtime/executor/method_meta.h | 16 ++++- runtime/executor/test/method_meta_test.cpp | 3 +- 9 files changed, 151 insertions(+), 25 deletions(-) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index d8c95191eb6..aaf15c8c46a 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -37,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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 36b27fb92ef..b2c1e069c9a 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -2100,6 +2100,19 @@ struct PyTensorInfo final { return info_.nbytes(); } + std::string device() const { + const auto device = info_.device(); + switch (device.type()) { + case runtime::etensor::DeviceType::CPU: + return "cpu"; + case runtime::etensor::DeviceType::CUDA: + return "cuda:" + std::to_string(device.index()); + } + throw std::runtime_error( + "Unknown device type " + + std::to_string(static_cast(device.type()))); + } + std::string repr() const { std::string size_str = "["; for (const auto& d : info_.sizes()) { @@ -2116,7 +2129,8 @@ struct PyTensorInfo final { std::string(executorch::runtime::toString(info_.scalar_type())) + ", is_memory_planned=" + (info_.is_memory_planned() ? "True" : "False") + - ", nbytes=" + std::to_string(info_.nbytes()) + ")"; + ", nbytes=" + std::to_string(info_.nbytes()) + ", device=" + device() + + ")"; } private: @@ -4343,6 +4357,10 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { "lacks.") .def_property_readonly( "nbytes", &PyExecuTorchResult::nbytes, "Size of the data in bytes.") + .def_property_readonly( + "device", + &PyExecuTorchResult::device_name, + "Where the data lives: \"cpu\" or \"cuda:N\".") .def( "__array__", [](const std::shared_ptr& result, @@ -4598,6 +4616,7 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { .def("dtype", &PyTensorInfo::dtype, call_guard) .def("is_memory_planned", &PyTensorInfo::is_memory_planned, call_guard) .def("nbytes", &PyTensorInfo::nbytes, call_guard) + .def("device", &PyTensorInfo::device, call_guard) .def("__repr__", &PyTensorInfo::repr, call_guard); py::class_(m, "MethodMeta", py::module_local()) .def("name", &PyMethodMeta::name, call_guard) diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 9d376133917..85cd3c8c4ed 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -49,6 +49,11 @@ class ExecuTorchResult: @property def nbytes(self) -> int: ... + @property + def device(self) -> str: + """Where the data lives: ``"cpu"`` or ``"cuda:N"``.""" + ... + def __array__(self, dtype: Any = None, copy: Any = None) -> Any: ... def __dlpack__( self, @@ -204,6 +209,10 @@ class TensorInfo: needs to be provided. False otherwise""" ... + def device(self) -> str: + """Where the program places the tensor: ``"cpu"`` or ``"cuda:N"``.""" + ... + def nbytes(self) -> int: """Number of bytes in the tensor. Not the same as numel if the dtype is larger than 1 byte wide""" diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 99d81e79394..da1fee1f869 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -502,10 +502,12 @@ class PyExecuTorchResult final if (sizes_.size() == 1) { shape += ","; } - const std::string device = - device_.is_cpu() ? "cpu" : "cuda:" + std::to_string(device_.index()); return "ExecuTorchResult(shape=(" + shape + - "), dtype=" + torch_dtype_name() + ", device=" + device + ")"; + "), dtype=" + torch_dtype_name() + ", device=" + device_name() + ")"; + } + + std::string device_name() const { + return device_.is_cpu() ? "cpu" : "cuda:" + std::to_string(device_.index()); } private: diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 71b56912ad0..bb1dd706569 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -1253,9 +1253,7 @@ def test_method_meta(self) -> None: self.assertEqual(meta.num_inputs(), 2) self.assertEqual(meta.num_outputs(), 1) - tensor_info = ( - "TensorInfo(sizes=[2, 2], dtype=Float, is_memory_planned=True, nbytes=16)" - ) + tensor_info = "TensorInfo(sizes=[2, 2], dtype=Float, is_memory_planned=True, nbytes=16, device=cpu)" float_dtype = 6 self.assertEqual( str(meta), @@ -1885,9 +1883,7 @@ def test_program_method_meta(self) -> None: self.assertEqual(meta.num_outputs(), 1) self.assertEqual(meta.num_attributes(), 1) - tensor_info = ( - "TensorInfo(sizes=[2, 2], dtype=Float, is_memory_planned=True, nbytes=16)" - ) + tensor_info = "TensorInfo(sizes=[2, 2], dtype=Float, is_memory_planned=True, nbytes=16, device=cpu)" float_dtype = 6 self.assertEqual( @@ -1939,9 +1935,7 @@ def test_method_method_meta(self) -> None: self.assertEqual(meta.num_inputs(), 2) self.assertEqual(meta.num_outputs(), 1) - tensor_info = ( - "TensorInfo(sizes=[2, 2], dtype=Float, is_memory_planned=True, nbytes=16)" - ) + tensor_info = "TensorInfo(sizes=[2, 2], dtype=Float, is_memory_planned=True, nbytes=16, device=cpu)" float_dtype = 6 self.assertEqual( str(meta), @@ -2868,6 +2862,62 @@ def forward(self, x): cloned.data_ptr(), value.data_ptr() ) + def test_method_meta_reports_tensor_devices(self): + from executorch.exir._serialize._program import PTEFile, serialize_pte_binary + from executorch.exir.schema import ExtraTensorInfo, Tensor + + class Identity(torch.nn.Module): + def forward(self, x): + return x + + program = to_edge(export(Identity(), (torch.ones(2, 2),))).to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ) + ) + ) + host_meta = self.load_prog_fn(program.buffer).method_meta("forward") + self.assertEqual(host_meta.input_tensor_meta(0).device(), "cpu") + self.assertEqual(host_meta.output_tensor_meta(0).device(), "cpu") + schema = program.executorch_program + for entry in schema.execution_plan[0].values: + if isinstance(entry.val, Tensor): + entry.val.extra_tensor_info = ExtraTensorInfo( + device_type=DeviceType.CUDA, device_index=3 + ) + data = bytes(serialize_pte_binary(PTEFile(schema))) + for meta in ( + self.load_prog_fn(data).method_meta("forward"), + self.load_fn(data).method_meta("forward"), + ): + self.assertEqual(meta.input_tensor_meta(0).device(), "cuda:3") + self.assertEqual(meta.output_tensor_meta(0).device(), "cuda:3") + self.assertIn("device=cuda:3", repr(meta.input_tensor_meta(0))) + + # A named buffer reports its device too. + eager_module = ModuleAddWithAttributes() + program = to_edge( + export(eager_module, eager_module.get_inputs(), strict=True) + ).to_executorch(ExecutorchBackendConfig(emit_mutable_buffer_names=True)) + schema = program.executorch_program + for entry in schema.execution_plan[0].values: + if isinstance(entry.val, Tensor): + entry.val.extra_tensor_info = ExtraTensorInfo( + device_type=DeviceType.CUDA, + device_index=3, + fully_qualified_name=( + entry.val.extra_tensor_info.fully_qualified_name + if entry.val.extra_tensor_info + else None + ), + ) + meta = self.load_prog_fn( + bytes(serialize_pte_binary(PTEFile(schema))) + ).method_meta("forward") + self.assertEqual(meta.num_attributes(), 1) + self.assertEqual(meta.attribute_tensor_meta(0).device(), "cuda:3") + def test_input_device_must_match_method(self): if self.kernel_mode != "portable": self.skipTest("only portable bindings inspect input device metadata") diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index e0513716f2a..c44546e00eb 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -281,6 +281,7 @@ def test_dlpack_export_keywords(self) -> None: if not isinstance(result, runtime.ExecuTorchResult): self.skipTest("torch output mode returns torch tensors") + self.assertEqual(result.device, "cpu") shared = np.from_dlpack(result, copy=False) self.assertTrue(np.shares_memory(shared, np.from_dlpack(result))) copied = np.from_dlpack(result, copy=True) diff --git a/runtime/executor/method_meta.cpp b/runtime/executor/method_meta.cpp index cc9ae4a3ed0..c06f87cced0 100644 --- a/runtime/executor/method_meta.cpp +++ b/runtime/executor/method_meta.cpp @@ -22,6 +22,19 @@ namespace executorch { namespace ET_RUNTIME_NAMESPACE { namespace { + +// Like the tensor parser, places a tensor without a device annotation on the +// CPU. +etensor::Device tensor_device(const executorch_flatbuffer::Tensor* tensor) { + if (tensor->extra_tensor_info() == nullptr) { + return etensor::Device{etensor::DeviceType::CPU, 0}; + } + return etensor::Device{ + static_cast( + tensor->extra_tensor_info()->device_type()), + static_cast( + tensor->extra_tensor_info()->device_index())}; +} bool has_required_metadata(const executorch_flatbuffer::ExecutionPlan& plan) { const auto* values = plan.values(); const auto* inputs = plan.inputs(); @@ -154,7 +167,8 @@ Result calculate_nbytes( Span dim_order, executorch::aten::ScalarType scalar_type, const bool is_memory_planned, - std::string_view name) { + std::string_view name, + etensor::Device device) { auto nbytes = calculate_nbytes(sizes, scalar_type); ET_CHECK_OR_RETURN_ERROR( nbytes.ok(), @@ -162,7 +176,13 @@ Result calculate_nbytes( "Failed to calculate nbytes for TensorInfo"); return TensorInfo( - sizes, dim_order, scalar_type, is_memory_planned, name, nbytes.get()); + sizes, + dim_order, + scalar_type, + is_memory_planned, + name, + nbytes.get(), + device); } TensorInfo::TensorInfo( @@ -171,13 +191,15 @@ TensorInfo::TensorInfo( executorch::aten::ScalarType scalar_type, const bool is_memory_planned, std::string_view name, - size_t nbytes) + size_t nbytes, + etensor::Device device) : sizes_(sizes), dim_order_(dim_order), name_(name), scalar_type_(scalar_type), is_memory_planned_(is_memory_planned), - nbytes_(nbytes) {} + nbytes_(nbytes), + device_(device) {} Span TensorInfo::sizes() const { return sizes_; @@ -199,6 +221,10 @@ size_t TensorInfo::nbytes() const { return nbytes_; } +etensor::Device TensorInfo::device() const { + return device_; +} + std::string_view TensorInfo::name() const { return name_; } @@ -284,8 +310,9 @@ Result MethodMeta::input_tensor_meta(size_t index) const { static_cast(tensor_value->scalar_type()), tensor_value->allocation_info() != nullptr || tensor_value->data_buffer_idx() != 0 /* is_memory_planned */, - std::string_view{nullptr, 0}); // Count constant returns as - // memory planned. + std::string_view{nullptr, 0}, // Count constant returns as + // memory planned. + tensor_device(tensor_value)); } size_t MethodMeta::num_outputs() const { @@ -341,8 +368,9 @@ Result MethodMeta::output_tensor_meta(size_t index) const { static_cast(tensor_value->scalar_type()), tensor_value->allocation_info() != nullptr || tensor_value->data_buffer_idx() != 0 /* is_memory_planned */, - std::string_view{nullptr, 0}); // Count constant returns as - // memory planned. + std::string_view{nullptr, 0}, // Count constant returns as + // memory planned. + tensor_device(tensor_value)); } size_t MethodMeta::num_attributes() const { @@ -398,7 +426,8 @@ Result MethodMeta::attribute_tensor_meta(size_t index) const { tensor_value->scalar_type()), tensor_value->allocation_info() != nullptr || tensor_value->data_buffer_idx() != 0 /* is_memory_planned */, - std::string_view{t_name->c_str(), t_name->size()}); + std::string_view{t_name->c_str(), t_name->size()}, + tensor_device(tensor_value)); } ++counter; } diff --git a/runtime/executor/method_meta.h b/runtime/executor/method_meta.h index d863e0125a5..e61e219182a 100644 --- a/runtime/executor/method_meta.h +++ b/runtime/executor/method_meta.h @@ -73,6 +73,12 @@ class TensorInfo final { */ std::string_view name() const; + /** + * Returns the device the program places the tensor on. Tensors without a + * device annotation are on the CPU. + */ + etensor::Device device() const; + private: // Let MethodMeta create TensorInfo. friend class MethodMeta; @@ -86,6 +92,7 @@ class TensorInfo final { * @param[in] scalar_type The scalar type of the tensor. * @param[in] is_memory_planned Whether the tensor's memory was planned. * @param[in] name The fully qualified name of the tensor. + * @param[in] device The device the program places the tensor on. * @returns A Result containing the TensorInfo on success, or an error on * failure. */ @@ -94,7 +101,8 @@ class TensorInfo final { Span dim_order, executorch::aten::ScalarType scalar_type, const bool is_memory_planned, - std::string_view name); + std::string_view name, + etensor::Device device); TensorInfo( Span sizes, @@ -102,7 +110,8 @@ class TensorInfo final { executorch::aten::ScalarType scalar_type, const bool is_memory_planned, std::string_view name, - size_t nbytes); + size_t nbytes, + etensor::Device device); /** * The sizes of the tensor. @@ -131,6 +140,9 @@ class TensorInfo final { /// The size in bytes of the tensor. size_t nbytes_; + + /// The device the program places the tensor on. + etensor::Device device_; }; /** diff --git a/runtime/executor/test/method_meta_test.cpp b/runtime/executor/test/method_meta_test.cpp index 971522bdfcf..fb05f7dfd56 100644 --- a/runtime/executor/test/method_meta_test.cpp +++ b/runtime/executor/test/method_meta_test.cpp @@ -48,7 +48,8 @@ class TensorInfoTestFriend final { Span(dim_order.data(), dim_order.size()), scalar_type, is_memory_planned, - name) + name, + etensor::Device{etensor::DeviceType::CPU, 0}) .get(); } }; From 899bb22dedb254aad19ef2dba08438d3a401b2e5 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:14:52 -0700 Subject: [PATCH 08/10] [PyBindings] Make the installed type stub usable The stub for the _C extension was installed as pybindings.pyi, a name no type checker associates with _C, and the package had no py.typed marker, so the annotations never applied. The stub also imported Enum from typing, declared the Verification members without values, required inputs on Module calls that default them, and put the bundled-program methods on ExecuTorchModule, where one does not exist and the other lives on BundledModule. The wheel now ships the stub as _C.pyi, as the Buck build already names it, and a py.typed marker in executorch.extension.pybindings. The stub fixes the points above. The wheel check fails if the stub or the marker is missing. --- .ci/scripts/wheel/test_base.py | 10 +++++++++ .gitignore | 1 + extension/pybindings/py.typed | 0 extension/pybindings/pybindings.pyi | 35 +++++++++++++---------------- pyproject.toml | 4 ++++ setup.py | 6 +++++ 6 files changed, 36 insertions(+), 20 deletions(-) create mode 100644 extension/pybindings/py.typed diff --git a/.ci/scripts/wheel/test_base.py b/.ci/scripts/wheel/test_base.py index 9e2eb9d6149..950e4c25e08 100644 --- a/.ci/scripts/wheel/test_base.py +++ b/.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. diff --git a/.gitignore b/.gitignore index 49db3eddb54..cab609d2e4c 100644 --- a/.gitignore +++ b/.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 diff --git a/extension/pybindings/py.typed b/extension/pybindings/py.typed new file mode 100644 index 00000000000..e69de29bb2d diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 85cd3c8c4ed..9bfc6728f1a 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -7,7 +7,8 @@ # pyre-strict from __future__ import annotations -from typing import Any, Callable, Dict, Enum, List, Optional, Tuple, TypeVar, Union +from enum import Enum +from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union _T = TypeVar("_T") @@ -21,8 +22,8 @@ class Verification(Enum): This API is experimental and subject to change without notice. """ - Minimal: ... - InternalConsistency: ... + Minimal = 0 + InternalConsistency = 1 @experimental("This API is experimental and subject to change without notice.") class ExecuTorchResult: @@ -89,37 +90,24 @@ class ExecuTorchModule: """ # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. - def __call__(self, inputs: Any, clone_outputs: bool = True) -> List[Any]: ... + def __call__(self, inputs: Any = ..., clone_outputs: bool = True) -> List[Any]: ... # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def run_method( self, method_name: str, - inputs: Any, + inputs: Any = ..., clone_outputs: bool = True, ) -> List[Any]: ... # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def forward( self, - inputs: Any, + inputs: Any = ..., clone_outputs: bool = True, ) -> List[Any]: ... # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def plan_execute( self, method_name: str, clone_outputs: bool = True ) -> List[Any]: ... - # Bundled program methods. - def load_bundled_input( - self, bundle: BundledModule, method_name: str, testset_idx: int - ) -> None: ... - # pyre-ignore[3]: "Any" in return type annotations. - def verify_result_with_bundled_expected_output( - self, - bundle: BundledModule, - method_name: str, - testset_idx: int, - rtol: float = 1e-5, - atol: float = 1e-8, - ) -> List[Any]: ... def has_etdump(self) -> bool: ... def write_etdump_result_to_file( self, path: str, debug_buffer_path: Optional[str] = None @@ -183,7 +171,14 @@ class BundledModule: This API is experimental and subject to change without notice. """ - ... + # pyre-ignore[3]: "Any" in return type annotations. + def verify_result_with_bundled_expected_output( + self, + method_name: str, + testset_idx: int, + rtol: float = 1e-5, + atol: float = 1e-8, + ) -> List[Any]: ... @experimental("This API is experimental and subject to change without notice.") class TensorInfo: diff --git a/pyproject.toml b/pyproject.toml index 9640001e415..7b865171146 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,6 +132,10 @@ license-files = ["LICENSE"] "*.glsl", "*.spirv.b64", ] +# Marks the Python bindings as typed, so type checkers read their stubs. +"executorch.extension.pybindings" = [ + "py.typed", +] [tool.setuptools.exclude-package-data] "*" = ["*.pyc"] diff --git a/setup.py b/setup.py index 36d97b609a0..c3d73287242 100644 --- a/setup.py +++ b/setup.py @@ -2347,6 +2347,12 @@ def run(self): "tools/cmake/executorch-wheel-config.cmake", "lib/cmake/executorch/executorch-config.cmake", ), + # Type checkers find a stub only under the module's own name, so + # the extension's stub ships as _C.pyi, as the Buck build names it. + ( + "extension/pybindings/pybindings.pyi", + "extension/pybindings/_C.pyi", + ), ] # The headers the package installs. Two audiences now: a custom-operator # build, which needs the kernel and tensor helpers, and a C++ application From 8c88fdc73eaa03c72ce2f6af9f8ba319284d7430 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:19:58 -0700 Subject: [PATCH 09/10] [exir] Plan device tensors copied to unplanned host outputs A method whose delegate output is copied to the host, exported with MemoryPlanningPass(alloc_graph_output=False), aborted at run time, whether its inputs were on the host or on the device: the CUDA delegate got a null pointer for its output. The device propagation pass inserts a copy from each delegate output to the host and rewires the graph output to that copy, but it left the output node's spec naming the delegate output. Memory planning reads the graph outputs from that spec, so it skipped the device tensor as if it were a graph output and planned the host copy instead. The pass now updates the output node's spec to the copy. --- exir/passes/propagate_device_pass.py | 9 ++++ exir/tests/test_propagate_device_pass.py | 41 +++++++++++++++++++ .../test/test_pybindings_cuda_dlpack.py | 11 +++++ 3 files changed, 61 insertions(+) diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py index c669fb5ed84..aac2de75593 100644 --- a/exir/passes/propagate_device_pass.py +++ b/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( diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py index ba5e87f2eb4..41902ddced4 100644 --- a/exir/tests/test_propagate_device_pass.py +++ b/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() diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index 947803cb6e4..83c772da952 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -279,6 +279,17 @@ def test_borrowed_result_written_by_another_library_is_waited_for(self): torch.from_dlpack(output), torch.full(self.shape, 84.0) ) + def test_device_inputs_with_unplanned_host_outputs(self): + # The delegate writes a device tensor that the method then copies into + # the caller's host output, so that device tensor must still be planned. + method = self._load_method( + self._export_delegate("unplanned_host_outputs", outputs_on_device=False) + ) + output = method(torch.full(self.shape, 1.0, device=self.device))[0] + torch.testing.assert_close( + torch.from_dlpack(output), torch.full(self.shape, 4.0) + ) + def test_clone_roundtrip_and_borrowed_output_lifetime(self): method = self._load_method() stream = torch.cuda.current_stream(self.device) From 11640ea6a06878fb9ebf3b49bc1128f2eaa73f2c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 20 Oct 2026 09:42:41 -0700 Subject: [PATCH 10/10] [runtime] Let Method.execute borrow outputs Module.forward and the native method take clone_outputs, but the public Method wrapper of executorch.runtime always copied every output, so a caller of the recommended API could not avoid that copy. Method.execute now takes clone_outputs, with the same default and meaning as the bindings: False returns outputs that share the method's memory and may change on the next execution. --- runtime/__init__.py | 16 +++++++++++----- runtime/test/test_runtime.py | 17 +++++++++++++++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/runtime/__init__.py b/runtime/__init__.py index e46c6cb25f3..c8095038719 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -136,16 +136,20 @@ class Method: def __init__(self, method: ExecuTorchMethod) -> None: self._method = method - def execute(self, inputs: Sequence[Any]) -> Sequence[Any]: + def execute( + self, inputs: Sequence[Any], clone_outputs: bool = True + ) -> Sequence[Any]: """Executes the method with the given inputs. Args: inputs: A sequence of input values, typically torch.Tensor objects. + clone_outputs: If False, the outputs share the method's memory and + a later execution may overwrite them, which avoids a copy. Returns: A list of output values, typically torch.Tensor objects. """ - return self._method(inputs) + return self._method(inputs, clone_outputs=clone_outputs) @property def metadata(self) -> MethodMeta: @@ -188,9 +192,11 @@ def load_method(self, name: str) -> Optional[Method]: program placed them. A method whose buffers are all on the host shares one set of arenas with every other host-only method of this program, sized to the largest of them, so two such methods overwrite each - other's intermediate values. Outputs are copied out on every call, so - that sharing does not change what execute returns. A method with any - buffer placed on an accelerator gets its own arenas instead, claimed + other's intermediate values. By default execute copies its outputs out, + so that sharing does not change what it returns. With + ``clone_outputs=False`` the outputs share that memory, so the next + execution of any method of this program may overwrite them. A method + with any buffer placed on an accelerator gets its own arenas instead, claimed when that method is first loaded and owned by that method until the method itself is released. This program caches every method it loads, so in normal use those arenas live as long as the program does. A diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 8dea5708eb6..1b1dcf3cd48 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -52,6 +52,23 @@ def test_smoke(self): outputs = method.execute(inputs) self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) + def test_execute_can_borrow_outputs(self): + ep, inputs = create_program(ModuleAdd()) + method = Runtime.get().load_program(ep.buffer).load_method("forward") + + def address(output): + return torch.from_dlpack(output).data_ptr() + + # Borrowed outputs are views of the method's own memory, so each + # execution writes the same place. Cloned outputs are copies. + borrowed = [method.execute(inputs, clone_outputs=False)[0] for _ in range(2)] + self.assertEqual(address(borrowed[0]), address(borrowed[1])) + cloned = [method.execute(inputs)[0] for _ in range(2)] + self.assertNotEqual(address(cloned[0]), address(cloned[1])) + torch.testing.assert_close( + torch.from_dlpack(borrowed[1]), inputs[0] + inputs[1] + ) + def test_module_with_multiple_method_names(self): ep, inputs = create_program(ModuleMulti()) runtime = Runtime.get()