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/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/README.md b/extension/pybindings/README.md index 8b031ff734b..aaf15c8c46a 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 @@ -35,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 @@ -45,7 +50,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 @@ -57,6 +68,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 @@ -87,7 +104,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 @@ -105,18 +125,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/py.typed b/extension/pybindings/py.typed new file mode 100644 index 00000000000..e69de29bb2d diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 635a4393184..b2c1e069c9a 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"); } } @@ -945,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 @@ -972,41 +1011,53 @@ 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. 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) - : 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; @@ -1070,21 +1121,13 @@ 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(); } - 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; } @@ -1096,6 +1139,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) { @@ -1138,6 +1240,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; } @@ -1273,6 +1378,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_; @@ -1370,7 +1476,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 +1506,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,12 +1564,14 @@ 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) { + if (PyCapsule_CheckExact(value.ptr())) { + return true; + } if (!py::hasattr(value, "__dlpack__")) { return false; } @@ -1485,7 +1596,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; @@ -1507,57 +1618,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 (const auto& input : torch_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 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 +2001,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 @@ -1957,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, @@ -2010,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()) { @@ -2026,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: @@ -2351,6 +2455,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 +2498,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 +2520,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 +2559,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 +2582,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 +2601,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 +2654,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 +2713,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 +2722,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 +2754,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 +3295,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 +3331,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 +3358,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 +3366,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 +3397,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 +3427,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 +3450,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 +3496,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 +3545,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 +3573,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 +3581,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; } @@ -3524,10 +3632,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 } @@ -3625,7 +3745,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); @@ -3930,8 +4053,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 @@ -3967,12 +4090,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); } } @@ -4053,7 +4174,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( @@ -4068,7 +4189,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) { @@ -4083,7 +4204,7 @@ struct PyProgram final { } bool has_etdump() { - return static_cast(event_tracer_); + return static_cast(state_->event_tracer_); } void write_etdump_result_to_file( @@ -4092,7 +4213,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); @@ -4118,8 +4239,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_; }; @@ -4238,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, @@ -4263,7 +4386,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", @@ -4489,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 e55eecb1e5e..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: @@ -49,8 +50,20 @@ 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, 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.""" @@ -77,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 @@ -171,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: @@ -197,6 +204,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_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_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 3a919f60aba..da1fee1f869 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 { @@ -278,21 +292,98 @@ 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( 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, and this is cheap next to a call. - 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; + 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( + 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 + // 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 = read_only_ ? dlpack::kFlagReadOnly : 0; + if (must_copy || (!copy.is_none() && copy.cast())) { + exported = copy_to(target); + flags = dlpack::kFlagIsCopied; + } + return versioned ? 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 @@ -411,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: @@ -450,12 +543,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 +563,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, @@ -539,6 +681,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 79d88fa73b2..bb1dd706569 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()) @@ -789,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) @@ -812,6 +870,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) @@ -962,14 +1043,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() @@ -1174,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), @@ -1582,7 +1659,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 +1699,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() @@ -1807,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( @@ -1861,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), @@ -1891,6 +1963,117 @@ 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_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() @@ -2679,9 +2862,65 @@ def forward(self, x): cloned.data_ptr(), value.data_ptr() ) - def test_torch_input_device_must_match_method(self): + 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 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 +2944,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..83c772da952 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,39 @@ 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_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() @@ -578,7 +589,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 +608,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) @@ -717,6 +756,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 b5d08deae35..c44546e00eb 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -272,6 +272,76 @@ 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") + + 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) + 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_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" @@ -343,23 +413,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: 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/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/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(); } }; 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() 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