diff --git a/docs/source/runtime-python-api-reference.rst b/docs/source/runtime-python-api-reference.rst index 87c8c8facd1..ea2243059e7 100644 --- a/docs/source/runtime-python-api-reference.rst +++ b/docs/source/runtime-python-api-reference.rst @@ -43,7 +43,21 @@ Other buffers, such as ``memoryview``, work as inputs too. NumPy has no bfloat16 For data on a GPU, use DLPack, the standard way to share tensors between Python libraries. This needs ExecuTorch built with CUDA support. Pass a CuPy array to a method that takes CUDA inputs. Read a CUDA result with ``cupy.from_dlpack(out)``, or copy it to the CPU with ``out.to_host()``. ``ExecuTorchResult.from_host()`` copies CPU data to a CUDA device. -To get ``ExecuTorchResult`` values even when PyTorch is installed, set ``EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch`` before importing ExecuTorch. +Output types +~~~~~~~~~~~~ + +Importing ``executorch.runtime`` does not import PyTorch, unless the setting below asks for ``torch``. When PyTorch is installed, ``Method.execute`` returns ``torch.Tensor`` outputs, and the first call that returns one imports ``torch``. Without PyTorch, or when importing it raises an ``ImportError``, it returns ``ExecuTorchResult`` outputs. Any other error from importing PyTorch is raised by every call that returns a tensor; set the setting below to ``executorch`` to run without PyTorch. + +A future release will make the outputs follow the tensor inputs of each call. A call with any ``torch.Tensor`` input will return ``torch.Tensor`` outputs. A call whose tensor inputs are all NumPy arrays, other buffers, DLPack objects, or ``ExecuTorchResult`` values will return ``ExecuTorchResult`` outputs, even when PyTorch is installed. A call with no tensor inputs keeps returning ``torch.Tensor`` outputs whenever ``torch`` can be imported. Each call whose outputs will change warns with a ``FutureWarning``. + +To choose the output type now, set ``EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT`` before importing ExecuTorch. With ``executorch``, every call returns ``ExecuTorchResult`` outputs and ExecuTorch never imports PyTorch. With ``torch``, every call returns ``torch.Tensor`` outputs, and importing ExecuTorch imports PyTorch. + +.. code-block:: python + + import os + + os.environ["EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"] = "executorch" + from executorch.runtime import Runtime CUDA streams ~~~~~~~~~~~~ diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 62ccff747c4..7e1ba466e88 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -41,9 +41,13 @@ execution can share the caller's device allocation. 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 -output type is `ExecuTorchResult`. Output conversion uses Python APIs and +Unless `EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT` is `torch`, importing +`portable_lib` does not import PyTorch. When PyTorch is installed, portable +bindings return `torch.Tensor` outputs, and the first call that returns one +imports `torch`. When PyTorch is not installed, or importing it raises an +`ImportError`, outputs are `ExecuTorchResult` values. Any other error from +importing it is raised by every call that returns a tensor; set the variable +to `executorch` to run without PyTorch. Output conversion uses Python APIs and DLPack; the extension does not link against PyTorch. `clone_outputs=True` returns storage that later executions do not overwrite, on the output's original device. With `clone_outputs=False`, results share @@ -91,9 +95,17 @@ with `clone_outputs=False` lives in the method's output memory, so passing it back in as an input means each run reads what the run before it wrote there; pass a copied result instead. -To select `ExecuTorchResult` outputs even when PyTorch is installed, set the -environment variable before importing the bindings. The setting is fixed when -the extension loads, so importing PyTorch later does not change the output type. +A future release will make tensor outputs follow the tensor inputs of each +call, including inputs kept from an earlier call. A call with any torch tensor +input will return torch tensors. A call whose tensor inputs are all NumPy +arrays, other buffers, DLPack objects, or `ExecuTorchResult` values will return +`ExecuTorchResult` values, even when PyTorch is installed. A call with no tensor +inputs keeps returning torch tensors whenever `torch` can be imported. Each call +whose outputs will change warns with a `FutureWarning`. To choose the output +type now, set `EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT` to `executorch` or `torch` +before importing the bindings. The setting is fixed when the extension loads. +`executorch` never imports PyTorch, and `torch` requires it to import +successfully. ```python import os @@ -103,11 +115,10 @@ from executorch.extension.pybindings import portable_lib ``` `ExecuTorchResult` preserves shape, strides, dtype, and device. CPU results expose -the Python buffer protocol so NumPy can read them without another copy. Without -the environment setting, the bindings fall back to `ExecuTorchResult` if -PyTorch cannot be imported. Explicitly selecting `torch` requires it to import -successfully. ATen-mode bindings continue to return native PyTorch tensors -directly. +the Python buffer protocol so NumPy can read them without another copy. A result +is not a number or a truth value: `bool()`, `int()`, and `float()` raise +`TypeError`, so convert it first, for example with `np.asarray(output)`. +ATen-mode bindings continue to return native PyTorch tensors directly. ```python output = module.forward([numpy_input])[0] @@ -133,11 +144,22 @@ output = method([device_input])[0] result = np.asarray(output.to_host()) ``` +The buffer protocol cannot describe dtypes such as bfloat16, so pass their data +as unsigned integers of the same width and name the dtype. For example, with an +`ml_dtypes.bfloat16` array: + +```python +device_input = ExecuTorchResult.from_host( + bfloat16_input.view(np.uint16), device_index=0, dtype="bfloat16" +) +``` + 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. +and import torch to do so unless `EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT` is +`executorch`. So with torch installed, or imported by you, 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. How long a call with diff --git a/extension/pybindings/portable_lib.py b/extension/pybindings/portable_lib.py index b936acaeb2f..3dc40d0d455 100644 --- a/extension/pybindings/portable_lib.py +++ b/extension/pybindings/portable_lib.py @@ -28,23 +28,19 @@ class _LightweightExperimentalWarning(DeprecationWarning): _TENSOR_OUTPUT_ENV = "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT" _tensor_output = os.environ.get(_TENSOR_OUTPUT_ENV) -if _tensor_output in (None, "torch"): +# Without the setting, the first call that returns torch.Tensor outputs imports +# torch, so importing ExecuTorch does not load it. +if _tensor_output == "torch": try: - # Preserve the existing API: importing portable_lib also imports Torch, - # and tensor outputs are torch.Tensor objects. - import torch as _torch + _import_module("torch") except ImportError as error: - if _tensor_output == "torch": - raise ImportError( - "ExecuTorch portable pybindings were configured to return " - "torch.Tensor outputs, but PyTorch could not be imported. " - f"Install PyTorch or set {_TENSOR_OUTPUT_ENV}=executorch " - "before importing portable_lib." - ) from error - _torch = None -elif _tensor_output == "executorch": - _torch = None -else: + raise ImportError( + "ExecuTorch portable pybindings were configured to return " + "torch.Tensor outputs, but PyTorch could not be imported. " + f"Install PyTorch or set {_TENSOR_OUTPUT_ENV}=executorch " + "before importing portable_lib." + ) from error +elif _tensor_output not in (None, "executorch"): raise ValueError( f"{_TENSOR_OUTPUT_ENV} must be either 'torch' or 'executorch', " f"but was {_tensor_output!r}" @@ -152,7 +148,6 @@ class _LightweightExperimentalWarning(DeprecationWarning): del _bindings del _tensor_output del _TENSOR_OUTPUT_ENV -del _torch del _warning_category del _import_module del _exir_warnings diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 17b4935a73d..071c8c10657 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -728,38 +728,153 @@ struct RuntimeErrorScope final { std::unique_ptr logs_; }; +// The Python type of a call's tensor outputs. ATen bindings always return +// torch tensors. +enum class PortableTensorOutput { Torch, ExecuTorch }; + #ifndef USE_ATEN_LIB constexpr char kTensorOutputEnvironmentVariable[] = "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"; -enum class PortableTensorOutput { Torch, ExecuTorch }; - -PortableTensorOutput& portable_tensor_output() { - static PortableTensorOutput output = PortableTensorOutput::Torch; +constexpr char kTensorOutputWarning[] = + "ExecuTorch returns torch.Tensor outputs for a call whose tensor inputs " + "are not torch.Tensor. A future release will return ExecuTorchResult " + "outputs for such a call. To keep torch.Tensor outputs, pass torch.Tensor " + "inputs. To choose the type for every call, set " + "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT to torch or executorch before " + "importing ExecuTorch. With executorch, PyTorch is not loaded."; + +// Set from EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT when the extension loads. +std::optional& tensor_output_override() { + static std::optional output; return output; } void configure_portable_tensor_output() { const char* value = std::getenv(kTensorOutputEnvironmentVariable); - if (value != nullptr && std::strcmp(value, "executorch") == 0) { - portable_tensor_output() = PortableTensorOutput::ExecuTorch; + if (value == nullptr) { + return; + } + if (std::strcmp(value, "executorch") == 0) { + tensor_output_override() = PortableTensorOutput::ExecuTorch; return; } - if (value != nullptr && std::strcmp(value, "torch") != 0) { + if (std::strcmp(value, "torch") != 0) { throw py::value_error( std::string(kTensorOutputEnvironmentVariable) + " must be either 'torch' or 'executorch', but was '" + value + "'"); } - try { - py::module_::import("torch"); - (void)loaded_torch_module(); - portable_tensor_output() = PortableTensorOutput::Torch; - } catch (py::error_already_set& error) { - if (value != nullptr || !error.matches(PyExc_ImportError)) { - throw; + py::module_::import("torch"); + tensor_output_override() = PortableTensorOutput::Torch; +} + +// Points a warning at the code that called ExecuTorch rather than at +// ExecuTorch's own Python wrappers, such as Method.execute. +int caller_stacklevel() { + int level = 1; + for (auto frame = py::reinterpret_borrow( + reinterpret_cast(PyEval_GetFrame())); + frame && !frame.attr("f_back").is_none() && + py::str(frame.attr("f_globals").attr("get")("__name__", "")) + .cast() + .rfind("executorch.", 0) == 0; + frame = frame.attr("f_back")) { + ++level; + } + return level; +} + +// Whether torch is imported, importing it if it can be. +bool torch_available() { + // Read and written with the GIL held. An import inside a function-local + // static's initializer could deadlock, since the import lets other threads + // run and one of them could then wait for that static while holding the GIL. + static bool torch_missing = false; + if (loaded_torch_module() == nullptr) { + // Checked only while torch is not loaded, so a torch imported after a call + // found it missing is used. + if (torch_missing) { + return false; + } + // A None in sys.modules blocks torch only until it is removed. + if (PyDict_GetItemString(PyImport_GetModuleDict(), "torch") == Py_None) { + return false; + } + try { + // Before Python 3.13, PyImport_ImportModule raises KeyError, not the + // ImportError, in a thread that waited for another thread's import. + py::module_::import("importlib").attr("import_module")("torch"); + } catch (py::error_already_set& error) { + // As when torch was imported with ExecuTorch, a torch that fails with + // an ImportError, such as for a missing library, counts as missing. + // Any other error is reported. + if (!error.matches(PyExc_ImportError)) { + try { + error.value().attr("add_note")( + "ExecuTorch imports torch here to return torch.Tensor outputs " + "or to wait for CUDA work. Set " + "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch before " + "importing ExecuTorch to skip it."); + } catch (...) { + // add_note is best effort; never hide the original error. + } + throw; + } + torch_missing = true; + return false; + } + } + return true; +} + +// Picks the type of a call's tensor outputs. A later release will follow the +// call's tensor inputs, including inputs kept from earlier calls: torch.Tensor +// if any of them is one, and ExecuTorchResult if all are other tensors. Until +// then, a call whose type that rule would change keeps torch.Tensor outputs +// whenever torch can be imported, as before, and warns. A call with no tensor +// inputs returns torch.Tensor whenever torch can be imported, now and later, +// so its type never depends on what imported torch first. Picking the type for +// a call that returns no tensor never imports torch. +PortableTensorOutput portable_tensor_output( + bool returns_tensor, + bool torch_input, + bool other_tensor_input) { + if (const auto& output = tensor_output_override()) { + return *output; + } + if (torch_input) { + return PortableTensorOutput::Torch; + } + if (!returns_tensor || !torch_available()) { + return PortableTensorOutput::ExecuTorch; + } + // Every call whose type will change warns, so Python's warning filters + // decide what shows: a warning hidden at one call still shows at a later one. + if (other_tensor_input && + PyErr_WarnEx( + PyExc_FutureWarning, kTensorOutputWarning, caller_stacklevel()) != + 0) { + // An error filter made the warning an exception. A call that runs the + // method picks the type first, so it raises before the run. + throw py::error_already_set(); + } + return PortableTensorOutput::Torch; +} + +// From the declared outputs, so the type is picked before the method runs. +// An output bound to a caller's tensor comes back as that tensor, so it does +// not count. +bool returns_tensor( + const MethodMeta& meta, + const std::vector& bound_outputs = {}) { + for (size_t i = 0; i < meta.num_outputs(); ++i) { + const auto tag = meta.output_tag(i); + const bool bound = i < bound_outputs.size() && bound_outputs[i]; + if (tag.ok() && tag.get() == Tag::Tensor && !bound) { + return true; } - portable_tensor_output() = PortableTensorOutput::ExecuTorch; } + return false; } #endif @@ -1222,6 +1337,39 @@ class BufferTensor final { executorch::aten::ScalarType scalar_type_; }; +// The dtype of host data that ExecuTorchResult.from_host() is told to read as +// `dtype`. The data comes the way results of that dtype export it: in the +// buffer format of that dtype, which int8 shares with qint8, or, for dtypes +// such as bfloat16 that PEP 3118 cannot describe, as unsigned integers of the +// same width. Any other mismatch is refused, so data is never reinterpreted +// by mistake. +executorch::aten::ScalarType host_data_dtype( + executorch::aten::ScalarType buffer_type, + const std::string& dtype) { + using executorch::aten::ScalarType; + const auto named = scalar_type_from_name(dtype); + if (!named) { + throw py::value_error("Unsupported dtype '" + dtype + "'"); + } + const char* format = python_dtype(*named).buffer_format; + const bool readable = format != nullptr + ? std::strcmp(format, python_dtype(buffer_type).buffer_format) == 0 + : runtime::elementSize(buffer_type) == runtime::elementSize(*named) && + (buffer_type == ScalarType::Byte || + buffer_type == ScalarType::UInt16 || + buffer_type == ScalarType::UInt32 || + buffer_type == ScalarType::UInt64); + if (!readable) { + throw py::value_error( + "dtype " + dtype + " cannot read a buffer of " + + python_dtype(buffer_type).torch_name + + ". Pass data in the buffer format of that dtype or, for a dtype the " + "buffer protocol cannot describe, unsigned integers of the same " + "width."); + } + return *named; +} + bool is_torch_tensor(const py::handle& value) { static PyObject* tensor_type = nullptr; if (tensor_type == nullptr) { @@ -2079,9 +2227,9 @@ py::object dlpack_output_owner( py::object module_aliased_torch_output( const executorch::aten::Tensor& tensor, const std::vector>& torch_inputs, - const std::string& method_name) { - if (torch_inputs.empty() || - portable_tensor_output() != PortableTensorOutput::Torch || + const std::string& method_name, + PortableTensorOutput output) { + if (torch_inputs.empty() || output != PortableTensorOutput::Torch || tensor.nbytes() == 0) { return py::none(); } @@ -2220,6 +2368,18 @@ std::shared_ptr allocate_device_output( return std::make_shared(std::move(buffer.get())); } +// Bindings built without CUDA support wait for CUDA work through torch, which +// used to be imported with ExecuTorch. So a wait imports it, unless the +// setting keeps torch out. +void synchronize_device_for_call(int index) { +#ifndef EXECUTORCH_PYBINDINGS_CUDA + if (tensor_output_override() != PortableTensorOutput::ExecuTorch) { + torch_available(); + } +#endif + synchronize_device(index); +} + // 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, @@ -2248,7 +2408,7 @@ void synchronize_tensor_devices(const Method& method) { add(method.get_output(i)); } for (const auto index : devices) { - synchronize_device(index); + synchronize_device_for_call(index); } } @@ -2515,9 +2675,10 @@ class PyCallerStreamScope final { py::object portable_tensor_result( const executorch::aten::Tensor& tensor, + PortableTensorOutput output, bool clone = true, py::object owner = py::none()) { - if (portable_tensor_output() == PortableTensorOutput::ExecuTorch) { + if (output == PortableTensorOutput::ExecuTorch) { return py::cast(std::make_shared( tensor, clone, clone ? py::none() : std::move(owner))); } @@ -2759,6 +2920,7 @@ inline std::unique_ptr load_module_from_data_loader( inline py::list get_outputs_as_py_list( const std::vector& outputs, + PortableTensorOutput output, bool clone_outputs = true, py::object owner = py::none(), const std::vector& output_owners = {}, @@ -2791,9 +2953,10 @@ inline py::list get_outputs_as_py_list( i < handed_over.size() && handed_over[i]; list[i] = returns_without_copy ? portable_tensor_result( - v.toTensor(), /*clone=*/false, handed_over[i]) + v.toTensor(), output, /*clone=*/false, handed_over[i]) : portable_tensor_result( v.toTensor(), + output, clone_outputs, output_owners.empty() ? owner : output_owners[i]); #endif @@ -2849,6 +3012,21 @@ struct PyBundledModule : public BundledModule { size_t testset_idx, double rtol = 1e-5, double atol = 1e-8) { + auto tensor_output = PortableTensorOutput::Torch; +#ifndef USE_ATEN_LIB + { + const auto meta = this->method_meta(method_name); + THROW_IF_ERROR( + meta.error(), + "Failed to get metadata for method %s", + method_name.c_str()); + // The inputs come from the bundled program, so Python passed no tensors. + tensor_output = portable_tensor_output( + returns_tensor(meta.get()), + /*torch_input=*/false, + /*other_tensor_input=*/false); + } +#endif // Execute the method auto result = BundledModule::execute(method_name, testset_idx); if (!result.ok()) { @@ -2857,6 +3035,7 @@ struct PyBundledModule : public BundledModule { "Method execution failed with status 0x%" PRIx32, static_cast(result.error())); } + const auto& outputs = result.get(); #ifndef USE_ATEN_LIB { const auto method = this->method(method_name); @@ -2869,8 +3048,7 @@ struct PyBundledModule : public BundledModule { #endif // Convert outputs to py::list - const auto& outputs = result.get(); - py::list py_outputs = get_outputs_as_py_list(outputs); + py::list py_outputs = get_outputs_as_py_list(outputs, tensor_output); Error status = BundledModule::verify_method_outputs( method_name, testset_idx, rtol, atol); @@ -3593,9 +3771,11 @@ struct PyModule final { for (size_t i = 0; i < inputs_size; ++i) { retained[i] = std::move(new_input_owners[i]); } + auto tensor_output = PortableTensorOutput::Torch; #ifndef USE_ATEN_LIB - // The run reads inputs kept from earlier calls too, so the wait and the - // output owners below cover every installed input, not only this call's. + // The run reads inputs kept from earlier calls too, so the output type, + // the wait and the output owners below cover every installed input, not + // only this call's. buffer_inputs.clear(); torch_inputs.clear(); dlpack_inputs.clear(); @@ -3610,6 +3790,10 @@ struct PyModule final { dlpack_inputs.push_back(owner.dlpack); } } + tensor_output = portable_tensor_output( + returns_tensor(method_meta), + !torch_inputs.empty(), + !buffer_inputs.empty() || !dlpack_inputs.empty()); synchronize_tensor_devices(method); #endif auto outputs = module_->execute(method_name); @@ -3641,6 +3825,7 @@ struct PyModule final { #endif auto result = get_outputs_as_py_list( outputs.get(), + tensor_output, clone_outputs, std::move(output_owner), output_owners, @@ -3652,7 +3837,7 @@ struct PyModule final { const auto& output = (*outputs)[i]; if (output.isTensor()) { auto alias = module_aliased_torch_output( - output.toTensor(), torch_inputs, method_name); + output.toTensor(), torch_inputs, method_name, tensor_output); if (!alias.is_none()) { result[i] = std::move(alias); } @@ -4633,6 +4818,18 @@ struct PyMethod final { "' has no outputs to return: it has not executed since its last " "set_inputs(), or that execution failed. Call execute() first."); } + return read_outputs(tensor_output(caller_outputs_), clone_outputs); + } + + py::list call(const py::object& inputs, bool clone_outputs = true) { + ObjectLock lock(mutex_); + set_inputs(inputs); + const auto output = tensor_output(bound_outputs_); + execute_installed(/*inputs_just_installed=*/true); + return read_outputs(output, clone_outputs); + } + + py::list read_outputs(PortableTensorOutput output, bool clone_outputs) { std::vector result(storage_->method->outputs_size()); Error get_outputs_status = storage_->method->get_outputs( @@ -4669,7 +4866,7 @@ struct PyMethod final { } // Retrieve outputs - py::list outputs = get_outputs_as_py_list(result, clone_outputs); + py::list outputs = get_outputs_as_py_list(result, output, clone_outputs); for (size_t i = 0; i < caller_backed.size(); ++i) { if (caller_backed[i]) { outputs[i] = caller_backed[i]; @@ -4678,13 +4875,6 @@ struct PyMethod final { return outputs; } - py::list call(const py::object& inputs, bool clone_outputs = true) { - ObjectLock lock(mutex_); - set_inputs(inputs); - execute_installed(/*inputs_just_installed=*/true); - return get_outputs(clone_outputs); - } - py::object get_attribute(const std::string& name) { ObjectLock lock(mutex_); Result attr = @@ -4700,18 +4890,24 @@ struct PyMethod final { #else // No execute() in this call has waited for the work that last wrote it. if (!attr->device().is_cpu() && attr->nbytes() != 0) { - synchronize_device(attr->device().index()); + synchronize_device_for_call(attr->device().index()); } // 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); + // An attribute does not come from the inputs of set_inputs(). + const auto output = portable_tensor_output( + /*returns_tensor=*/true, + /*torch_input=*/false, + /*other_tensor_input=*/false); + if (constant && output == PortableTensorOutput::Torch) { + return portable_tensor_result(attr.get(), output, /*clone=*/true); } auto result = portable_tensor_result( attr.get(), + output, /*clone=*/false, borrowed_output_owner(attr.get())); if (constant) { @@ -4776,6 +4972,20 @@ struct PyMethod final { return index < bound_outputs_.size() && bool(bound_outputs_[index]); } + // Uses the inputs of the last successful set_inputs(), and the outputs bound + // for the run that the outputs come from. + PortableTensorOutput tensor_output( + const std::vector& bound_outputs) { +#ifdef USE_ATEN_LIB + return PortableTensorOutput::Torch; +#else + return portable_tensor_output( + returns_tensor(storage_->method->method_meta(), bound_outputs), + !torch_inputs_.empty(), + !buffer_inputs_.empty() || !dlpack_inputs_.empty()); +#endif + } + std::optional aliased_input(size_t index) const { return index < inputs_aliased_by_outputs_.size() ? inputs_aliased_by_outputs_[index] @@ -4855,9 +5065,10 @@ struct PyMethod final { std::vector> dlpack_inputs; }; - py::object aliased_torch_output(const executorch::aten::Tensor& tensor) { - if (torch_inputs_.empty() || - portable_tensor_output() != PortableTensorOutput::Torch || + py::object aliased_torch_output( + const executorch::aten::Tensor& tensor, + PortableTensorOutput output) { + if (torch_inputs_.empty() || output != PortableTensorOutput::Torch || tensor.nbytes() == 0) { return py::none(); } @@ -4909,7 +5120,8 @@ struct PyMethod final { // result stays independent. std::optional hand_over_device_output( size_t index, - const executorch::aten::Tensor& tensor) { + const executorch::aten::Tensor& tensor, + PortableTensorOutput output) { const auto slot = device_buffer_slot(tensor); if (!slot) { return std::nullopt; @@ -4925,10 +5137,11 @@ struct PyMethod final { std::make_shared>(buffer); const std::shared_ptr memory(holder, buffer->data()); py::object result; - if (portable_tensor_output() == PortableTensorOutput::Torch) { + if (output == PortableTensorOutput::Torch) { auto* retained = new std::shared_ptr(memory); result = portable_tensor_result( tensor, + output, /*clone=*/false, py::capsule(retained, [](void* pointer) { delete static_cast*>(pointer); @@ -5109,6 +5322,7 @@ struct PyMethod final { py::list get_outputs_as_py_list( const std::vector& outputs, + PortableTensorOutput output, bool clone_outputs = true) { const auto outputs_size = outputs.size(); py::list list(outputs_size); @@ -5143,13 +5357,14 @@ struct PyMethod final { continue; } if (clone_outputs) { - if (const auto handed = hand_over_device_output(i, v.toTensor())) { + if (const auto handed = + hand_over_device_output(i, v.toTensor(), output)) { list[i] = *handed; continue; } } if (!clone_outputs) { - auto alias = aliased_torch_output(v.toTensor()); + auto alias = aliased_torch_output(v.toTensor(), output); if (!alias.is_none()) { list[i] = std::move(alias); continue; @@ -5157,6 +5372,7 @@ struct PyMethod final { } list[i] = portable_tensor_result( v.toTensor(), + output, clone_outputs, clone_outputs && v.toTensor().device().is_cpu() ? py::none() @@ -5447,6 +5663,15 @@ py::dict get_option(const std::string& backend_name, const py::dict& options) { return result; } +auto refuse_scalar_conversion(const char* problem) { + return [problem](const PyExecuTorchResult&) -> py::object { + throw py::type_error( + std::string("ExecuTorchResult ") + problem + + ". Convert it first, for example with numpy.asarray(result) or " + "torch.from_dlpack(result)."); + }; +} + } // namespace PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { @@ -5491,9 +5716,12 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { #else configure_portable_tensor_output(); m.attr("_uses_aten") = false; - m.attr("_tensor_output") = - portable_tensor_output() == PortableTensorOutput::Torch ? "torch" - : "executorch"; + const auto& tensor_output = tensor_output_override(); + m.attr("_tensor_output") = !tensor_output + ? py::object(py::none()) + : py::str( + *tensor_output == PortableTensorOutput::Torch ? "torch" + : "executorch"); #endif py::class_>( m, @@ -5559,10 +5787,18 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { "Return a copy of this result in CPU memory.") .def_static( "from_host", - [](const py::object& source, int device_index) { + [](const py::object& source, + int device_index, + const std::optional& dtype) { if (py::isinstance(source)) { - return source.cast>() - ->to_device(device_index); + const auto result = + source.cast>(); + if (dtype && *dtype != result->dtype()) { + throw py::value_error( + "dtype " + *dtype + " does not match the result's dtype " + + result->dtype()); + } + return result->to_device(device_index); } const BufferTensor host(source); return PyExecuTorchResult::from_host( @@ -5571,14 +5807,26 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { host.sizes().begin(), host.sizes().end()), std::vector( host.strides().begin(), host.strides().end()), - host.scalar_type(), + dtype ? host_data_dtype(host.scalar_type(), *dtype) + : host.scalar_type(), device_index); }, py::arg("source"), py::arg("device_index") = 0, + py::kw_only(), + py::arg("dtype") = py::none(), "Copy a dense CPU buffer, or a CPU ExecuTorchResult, to a CUDA " - "device, for use as a method input. Requires bindings built with " - "CUDA support.") + "device, for use as a method input. dtype, such as \"bfloat16\", " + "is the dtype of the copy, for data that comes the way results of " + "that dtype export it: for dtypes the buffer protocol cannot " + "describe, as unsigned integers of the same width. Requires " + "bindings built with CUDA support.") + // Python would otherwise treat a result as always true, and int() and + // float() would parse its bytes as text. + .def("__bool__", refuse_scalar_conversion("has no truth value")) + .def("__int__", refuse_scalar_conversion("cannot be converted to int")) + .def( + "__float__", refuse_scalar_conversion("cannot be converted to float")) .def("__repr__", &PyExecuTorchResult::repr); #ifndef USE_ATEN_LIB py::class_< diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index b78a657434e..3f0c0a5402f 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -70,10 +70,16 @@ class ExecuTorchResult: ... @staticmethod - def from_host(source: Any, device_index: int = 0) -> ExecuTorchResult: + def from_host( + source: Any, device_index: int = 0, *, dtype: Optional[str] = None + ) -> ExecuTorchResult: """Copy a dense CPU buffer, or a CPU ExecuTorchResult, to a CUDA device, for use as a method input. + ``dtype``, such as ``"bfloat16"``, is the dtype of the copy, for data + that comes the way results of that dtype export it: for dtypes the + buffer protocol cannot describe, as unsigned integers of the same width. + Requires bindings built with CUDA support. """ ... diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 1397c46186a..6a3db374d15 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/extension/pybindings/pybindings_executorch_result.h @@ -76,17 +76,30 @@ inline const PyTensorDType& python_dtype( return kDTypes[index]; } -inline executorch::aten::ScalarType scalar_type_from_torch_dtype( - const std::string& torch_dtype) { +// The dtype that ExecuTorchResult.dtype calls `name`, such as "bfloat16". +inline std::optional scalar_type_from_name( + const std::string& name) { using executorch::aten::ScalarType; - for (int8_t value = 0; value < static_cast(ScalarType::NumOptions); + // The dtype table ends at UInt64, while ATen's ScalarType goes on past it. + for (int8_t value = 0; value <= static_cast(ScalarType::UInt64); ++value) { const auto scalar_type = static_cast(value); - const auto& dtype = python_dtype(scalar_type); - if (torch_dtype == std::string("torch.") + dtype.torch_name) { + if (name == python_dtype(scalar_type).torch_name) { return scalar_type; } } + return std::nullopt; +} + +inline executorch::aten::ScalarType scalar_type_from_torch_dtype( + const std::string& torch_dtype) { + const std::string prefix = "torch."; + if (torch_dtype.compare(0, prefix.size(), prefix) == 0) { + if (const auto scalar_type = + scalar_type_from_name(torch_dtype.substr(prefix.size()))) { + return *scalar_type; + } + } throw py::value_error("Unsupported torch dtype: " + torch_dtype); } @@ -97,10 +110,19 @@ inline PyObject* loaded_torch_module() { static PyObject* torch_module = nullptr; if (torch_module == nullptr) { PyObject* loaded = PyDict_GetItemString(PyImport_GetModuleDict(), "torch"); - if (loaded != nullptr) { - Py_INCREF(loaded); - torch_module = loaded; + if (loaded == nullptr || loaded == Py_None) { + return nullptr; + } + // Owned first: the lookups below run Python code, which can remove it + // from sys.modules. + auto module = py::reinterpret_borrow(loaded); + // A torch that another thread is still importing is not usable yet. + const auto spec = py::getattr(module, "__spec__", py::none()); + if (!spec.is_none() && + py::getattr(spec, "_initializing", py::bool_(false)).cast()) { + return nullptr; } + torch_module = module.release().ptr(); } return torch_module; } @@ -112,16 +134,19 @@ inline void synchronize_torch_cuda_device(int index) { } // Waits for every queued operation on a device. Bindings built without CUDA -// support have no CUDA runtime of their own, so they wait through torch's, -// which is loaded whenever such bindings see CUDA memory. +// support have no CUDA runtime of their own, so they wait through torch's +// once torch is imported. inline void synchronize_device(int index) { #ifndef EXECUTORCH_PYBINDINGS_CUDA - if (loaded_torch_module() != nullptr) { - synchronize_torch_cuda_device(index); - return; + if (loaded_torch_module() == nullptr) { + throw std::runtime_error( + "CUDA tensor interchange requires bindings built with CUDA support, " + "or torch imported to wait for CUDA work"); } -#endif + synchronize_torch_cuda_device(index); +#else synchronize_cuda_device(index); +#endif } /** Tensor result returned by portable Python bindings. */ diff --git a/extension/pybindings/test/BUCK b/extension/pybindings/test/BUCK index 0602bfc2de2..d6928718dfc 100644 --- a/extension/pybindings/test/BUCK +++ b/extension/pybindings/test/BUCK @@ -44,6 +44,7 @@ fbcode_target( "//executorch/exir:schema", "//executorch/exir/backend/test:device_util", "//executorch/extension/pybindings:portable_lib", + "//executorch/runtime:runtime", ], ) diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index ea75811b113..2a63c4edf53 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -14,8 +14,10 @@ import subprocess import sys import tempfile +import textwrap import time import unittest +import warnings import weakref from io import StringIO @@ -163,115 +165,525 @@ def test_numpy_float16_buffer_input(self): self.assertTrue(torch.equal(output, inputs[0] + inputs[0])) - def test_buffer_input_does_not_import_torch(self): - exported_program, inputs = create_program(ModuleAddSingleInput()) - executorch_module = self.load_fn(exported_program.buffer) - torch_module = sys.modules.pop("torch") - try: - output = executorch_module(inputs[0].numpy())[0] - self.assertNotIn("torch", sys.modules) - finally: - sys.modules["torch"] = torch_module + def test_tensor_output_type(self): + if self.kernel_mode != "portable": + self.skipTest("only portable bindings choose Python tensor outputs") + from executorch.devtools.bundled_program.config import ( + MethodTestCase, + MethodTestSuite, + ) + from executorch.devtools.bundled_program.core import BundledProgram + from executorch.devtools.bundled_program.serialize import ( + serialize_from_bundled_program_to_flatbuffer, + ) - self.assertTrue(torch.equal(output, inputs[0] + inputs[0])) + inputs = ModuleAddWithAttributes().get_inputs() + program = to_edge( + export(ModuleAddWithAttributes(), inputs, strict=True), + constant_methods={"tensor_constant": torch.ones(2), "int_constant": 3}, + ).to_executorch(ExecutorchBackendConfig(emit_mutable_buffer_names=True)) + bundled_program, _ = create_program(ModuleAdd()) + bundled = serialize_from_bundled_program_to_flatbuffer( + BundledProgram( + bundled_program, + [ + MethodTestSuite( + method_name="forward", + test_cases=[ + MethodTestCase( + inputs=inputs, expected_outputs=(inputs[0] + inputs[1],) + ) + ], + ) + ], + ) + ) + prelude = """ +import importlib.abc +import sys +import threading +import warnings + +import numpy as np - def test_tensor_output_import_policy(self): +program_path, bundled_path = sys.argv[1:] +x = np.ones((2, 2), dtype=np.float32) + +class BlockTorch(importlib.abc.MetaPathFinder): + def __init__(self, error=ModuleNotFoundError): + self.error = error + # The module the error names: torch itself, or a part of it. + self.missing = "torch" + self.attempts = 0 + + def find_spec(self, name, path=None, target=None): + if name == "torch" or name.startswith("torch."): + self.attempts += 1 + error = self.error("torch is unavailable") + if isinstance(error, ImportError): + error.name = self.missing + raise error + +def kinds(outputs): + torch = sys.modules.get("torch") + return [ + "torch" if torch is not None and isinstance(value, torch.Tensor) + else type(value).__name__ + for value in outputs + ] + +# Importing torch can show warnings of its own, so only this one counts. +OURS = "ExecuTorch returns torch.Tensor outputs" + +def ours(caught): + return [w for w in caught if str(w.message).startswith(OURS)] + +def check(call, expected, warnings_expected=0): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + outputs = call() + future = ours(caught) + assert kinds(outputs) == expected, (kinds(outputs), expected) + assert len(future) == warnings_expected, [str(w.message) for w in future] + assert all(w.category is FutureWarning for w in future) + return future +""" + # Each runs in a fresh process, since torch imports are process-wide. + scenarios = { + "NumPy inputs before torch is imported": ( + None, + """ + from executorch.runtime import Runtime + + assert "torch" not in sys.modules + program = Runtime.get().load_program(program_path) + check(lambda: program.load_method("int_constant").execute([]), ["int"]) + assert "torch" not in sys.modules + forward = program.load_method("forward") + (warning,), line = check(lambda: forward.execute([x, x]), ["torch"], 1), sys._getframe().f_lineno + # The warning points at the caller, not at ExecuTorch's wrapper. + assert (warning.filename, warning.lineno) == ("", line), warning + assert "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT" in str(warning.message) + check(lambda: forward.execute([x, x]), ["torch"], 1) + """, + ), + "NumPy inputs kept by set_inputs after torch is imported": ( + None, + """ + import torch + from executorch.extension.pybindings import portable_lib + + native = sys.modules["executorch.extension.pybindings._C"] + assert native._tensor_output is None + method = portable_lib._load_program(program_path).load_method("forward") + method.set_inputs([x, x]) + method.execute() + # An attribute does not come from the inputs. + check(lambda: [method.get_attribute("state")], ["torch"]) + check(method.get_outputs, ["torch"], 1) + """, + ), + "NumPy inputs with every output bound to a torch tensor": ( + None, + """ + import torch + from executorch.exir import ExecutorchBackendConfig, to_edge + from executorch.exir.passes import MemoryPlanningPass + from executorch.extension.pybindings import portable_lib + + class Add(torch.nn.Module): + def forward(self, a, b): + return a + b + + t = torch.ones(2, 2) + config = ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass(alloc_graph_output=False) + ) + buffer = to_edge(torch.export.export(Add(), (t, t))).to_executorch(config).buffer + program = portable_lib._load_program_from_buffer(buffer) + method = program.load_method("forward") + bound = torch.empty(2, 2) + method.set_output(bound, 0) + # A bound output is the caller's own tensor, now and later. + check(lambda: method([x, x]), ["torch"]) + assert method([x, x])[0] is bound + # Bound after the run, so the outputs of that run are not. + late = program.load_method("forward") + late.set_inputs([x, x]) + late.execute() + late.set_output(torch.empty(2, 2), 0) + check(late.get_outputs, ["torch"], 1) + """, + ), + "torch inputs, kept or mixed, and no tensor inputs once torch is imported": ( + None, + """ + import torch + from executorch.extension.pybindings import portable_lib + + t = torch.ones(2, 2) + module = portable_lib._load_for_executorch(program_path) + check(lambda: module.forward([t, t]), ["torch"]) + # Input 1 is still the torch tensor of the last call. + check(lambda: module.run_method("forward", [x]), ["torch"]) + check(lambda: module.forward([x, t]), ["torch"]) + check(lambda: module.run_method("tensor_constant"), ["torch"]) + with open(bundled_path, "rb") as source: + bundled = portable_lib._load_bundled_program_from_buffer( + source.read() + ) + check( + lambda: bundled.verify_result_with_bundled_expected_output( + "forward", 0 + ), + ["torch"], + ) + """, + ), + "the warning turned into an error": ( + None, + """ + from executorch.extension.pybindings import portable_lib + + module = portable_lib._load_for_executorch(program_path) + method = portable_lib._load_program(program_path).load_method("forward") + with warnings.catch_warnings(): + warnings.filterwarnings("error", OURS) + for call in (module.forward, method) * 2: + try: + call([x, x]) + except FutureWarning: + pass + else: + raise AssertionError("every such call must raise") + # The raising calls did not run: state is still 0, so 1 + 1 + 1. + for call in (module.forward, method): + assert call([x, x])[0].tolist() == [[3.0, 3.0]] * 2 + """, + ), + "the first warning filtered out": ( + None, + """ + from executorch.extension.pybindings import portable_lib + + module = portable_lib._load_for_executorch(program_path) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + module.forward([x, x]) + # A warning nobody saw does not use up the later ones. + check(lambda: module.forward([x, x]), ["torch"], 1) + """, + ), + "two threads before torch is imported": ( + None, + """ + from executorch.extension.pybindings import portable_lib + + modules = [portable_lib._load_for_executorch(program_path) for _ in range(2)] + barrier = threading.Barrier(2) + outputs = [None, None] + + def run(index): + barrier.wait() + outputs[index] = modules[index].forward([x, x])[0] + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + threads = [threading.Thread(target=run, args=(i,)) for i in (0, 1)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert kinds(outputs) == ["torch", "torch"], kinds(outputs) + assert len(ours(caught)) == 2, [str(w.message) for w in caught] + """, + ), + "torch fails to import or is not installed": ( + None, + """ + blocker = BlockTorch(RuntimeError) + sys.meta_path.insert(0, blocker) + from executorch.extension.pybindings import portable_lib + + module = portable_lib._load_for_executorch(program_path) + # A torch that fails with an error other than ImportError is + # reported at every call. + for _ in range(2): + try: + module.forward([x, x]) + except RuntimeError as error: + assert "torch is unavailable" in str(error), error + notes = "".join(getattr(error, "__notes__", [])) + assert "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT" in notes, notes + else: + raise AssertionError("a broken torch must not be hidden") + # An ImportError, even for a part of torch, counts as missing, + # as when torch was imported with ExecuTorch. + blocker.error, blocker.missing = ImportError, "torch._C" + check(lambda: module.forward([x, x]), ["ExecuTorchResult"]) + check(lambda: module.run_method("tensor_constant"), ["ExecuTorchResult"]) + # A missing torch is not looked for again on every call. + assert blocker.attempts == 3, blocker.attempts + assert "torch" not in sys.modules + """, + ), + "a torch that fails with an ImportError, in two threads": ( + None, + """ + import os + import tempfile + import time + + package = os.path.join(tempfile.mkdtemp(), "torch") + os.makedirs(package) + with open(os.path.join(package, "__init__.py"), "w") as source: + source.write( + "import __main__\\n" + "__main__.started.set()\\n" + "__main__.release.wait()\\n" + "raise ImportError('libtorch_cpu.so is missing')\\n" + ) + sys.path.insert(0, os.path.dirname(package)) + started, release = threading.Event(), threading.Event() + from executorch.extension.pybindings import portable_lib + + modules = [portable_lib._load_for_executorch(program_path) for _ in range(2)] + outputs = [None, None] + + def run(index): + outputs[index] = modules[index].forward([x, x])[0] + + threads = [threading.Thread(target=run, args=(i,)) for i in (0, 1)] + threads[0].start() + started.wait() + # The second call waits for the import that the first one runs. + threads[1].start() + time.sleep(0.5) + release.set() + for thread in threads: + thread.join() + assert kinds(outputs) == ["ExecuTorchResult"] * 2, kinds(outputs) + """, + ), + "the setting asks for ExecuTorchResult": ( + "executorch", + """ + from executorch.extension.pybindings import portable_lib + + native = sys.modules["executorch.extension.pybindings._C"] + assert native._tensor_output == "executorch" + module = portable_lib._load_for_executorch(program_path) + check(lambda: module.forward([x, x]), ["ExecuTorchResult"]) + check(lambda: module.run_method("tensor_constant"), ["ExecuTorchResult"]) + assert "torch" not in sys.modules + import torch + + t = torch.ones(2, 2) + check(lambda: module.forward([t, t]), ["ExecuTorchResult"]) + """, + ), + "the setting asks for torch.Tensor": ( + "torch", + """ + import importlib.util + import os + from pathlib import Path + + # The extension imported on its own, without portable_lib, which + # imports torch itself and on Windows adds the folders of the + # extension's libraries. + name = "executorch.extension.pybindings._C" + if sys.platform == "win32": + native_directory = Path(importlib.util.find_spec(name).origin).parent + for directory in (native_directory, native_directory.parent.parent / "lib"): + if directory.is_dir(): + os.add_dll_directory(str(directory.resolve())) + native = importlib.import_module(name) + assert "torch" in sys.modules + assert native._tensor_output == "torch" + from executorch.extension.pybindings import portable_lib + + module = portable_lib._load_for_executorch(program_path) + check(lambda: module.forward([x, x]), ["torch"]) + """, + ), + "the setting asks for torch.Tensor without torch": ( + "torch", + """ + sys.meta_path.insert(0, BlockTorch()) + try: + from executorch.extension.pybindings import portable_lib + except ImportError as error: + assert "PyTorch could not be imported" in str(error), error + else: + raise AssertionError("importing must fail without torch") + """, + ), + } + # Calls with tensor inputs of other kinds return torch.Tensor outputs + # now and will return ExecuTorchResult outputs later, so they warn, + # even after torch is imported. Calls without tensor inputs keep + # torch.Tensor outputs, so they do not, even before it is. + for call, torch_imported, warnings_expected in ( + ("module.run_method('tensor_constant')", False, 0), + ("[method.get_attribute('state')]", False, 0), + ( + "bundled.verify_result_with_bundled_expected_output('forward', 0)", + False, + 0, + ), + ("module.forward([x, x])", True, 1), + ("module.forward(capsules)", True, 1), + ("method(capsules)", True, 1), + ): + scenarios[f"{call}, torch imported: {torch_imported}"] = ( + None, + f""" + {'import torch' if torch_imported else ''} + from executorch.extension.pybindings import portable_lib + + module = portable_lib._load_for_executorch(program_path) + method = portable_lib._load_program(program_path).load_method("forward") + with open(bundled_path, "rb") as source: + bundled = portable_lib._load_bundled_program_from_buffer(source.read()) + capsules = [x.__dlpack__(), x.__dlpack__()] + caught, line = check(lambda: {call}, ["torch"], {warnings_expected}), sys._getframe().f_lineno + assert all(warning.lineno == line for warning in caught), caught + """, + ) + with tempfile.TemporaryDirectory() as directory: + program_path = os.path.join(directory, "program.pte") + with open(program_path, "wb") as handle: + handle.write(program.buffer) + bundled_path = os.path.join(directory, "program.bpte") + with open(bundled_path, "wb") as handle: + handle.write(bundled) + for name, (setting, code) in scenarios.items(): + with self.subTest(name): + environment = os.environ.copy() + environment.pop("EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT", None) + if setting is not None: + environment["EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"] = setting + completed = subprocess.run( + [ + sys.executable, + "-c", + prelude + textwrap.dedent(code), + program_path, + bundled_path, + ], + env=environment, + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual( + completed.returncode, 0, completed.stdout + completed.stderr + ) + + def test_torch_that_is_not_ready_is_ignored(self): if self.kernel_mode != "portable": - self.skipTest("only portable bindings configure Python tensor outputs") + self.skipTest("only portable bindings look torch up in sys.modules") program, _ = create_program(ModuleAddSingleInput()) script = """ -import importlib import importlib.abc import os import sys -from pathlib import Path +import threading -class BlockTorch(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "torch" or fullname.startswith("torch."): - raise ImportError("torch is unavailable") - -entrypoint, native_name, policy, available, program_path, native_path = sys.argv[1:] -if sys.platform == "win32": - native_directory = Path(native_path).parent - os.add_dll_directory(str(native_directory)) - library_directory = native_directory.parent.parent / "lib" - if library_directory.is_dir(): - os.add_dll_directory(str(library_directory)) -blocker = BlockTorch() -if available == "no": +import numpy as np + +program_path, case = sys.argv[1:] +if case in ("none", "restored"): + sys.modules["torch"] = None +elif case == "missing": + class BlockTorch(importlib.abc.MetaPathFinder): + def find_spec(self, name, path=None, target=None): + if name == "torch": + raise ModuleNotFoundError(f"No module named {name!r}", name=name) + + blocker = BlockTorch() sys.meta_path.insert(0, blocker) +elif case == "importing": + # A torch that another thread is still importing. + package = os.path.join(os.path.dirname(program_path), "partial", "torch") + os.makedirs(package) + with open(os.path.join(package, "__init__.py"), "w") as source: + source.write("import __main__\\n__main__.started.set()\\n__main__.release.wait()\\n") + sys.path.insert(0, os.path.dirname(package)) + started, release = threading.Event(), threading.Event() + importer = threading.Thread(target=__import__, args=("torch",), daemon=True) + importer.start() + started.wait() +from executorch.extension.pybindings import portable_lib +with open(program_path, "rb") as source: + module = portable_lib._load_for_executorch_from_buffer(source.read()) +if case == "vanishing": + # Reading its __spec__ drops the only other reference to it. + class Vanishing: + @property + def __spec__(self): + del sys.modules["torch"] + + sys.modules["torch"] = Vanishing() try: - bindings = importlib.import_module(entrypoint) -except (ImportError, SystemError) as error: - # Built-in extension loaders may wrap initialization errors. - assert policy == "torch" and available == "no" - while error.__cause__ is not None: - error = error.__cause__ - assert isinstance(error, ImportError) - assert "torch is unavailable" in str(error) + output = module(np.ones((2, 2), dtype=np.float32))[0] +except AttributeError: + # That object has no Tensor to check inputs against; it only must not + # crash the process. + assert case == "vanishing" else: - assert policy != "torch" or available == "yes" - native = sys.modules[native_name] - expected = "torch" if policy != "executorch" and available == "yes" else "executorch" - assert native._tensor_output == expected - assert ("torch" in sys.modules) == (expected == "torch") - import numpy as np - program_data = Path(program_path).read_bytes() - module = bindings._load_for_executorch_from_buffer(program_data) + assert isinstance(output, portable_lib.ExecuTorchResult), type(output) +if case == "importing": + release.set() + importer.join() +if case == "restored": + # A torch that sys.modules blocked at one call is imported by a later one. + del sys.modules["torch"] +elif case == "missing": + # A torch that was missing at one call is used once it is imported. + sys.meta_path.remove(blocker) + import torch +if case in ("restored", "missing"): output = module(np.ones((2, 2), dtype=np.float32))[0] - if expected == "executorch": - assert isinstance(output, bindings.ExecuTorchResult) - np.testing.assert_array_equal(np.asarray(output), np.full((2, 2), 2.0)) - if available == "no": - sys.meta_path.remove(blocker) - import torch - assert isinstance(module(torch.ones(2, 2))[0], bindings.ExecuTorchResult) - else: - import torch - assert isinstance(output, torch.Tensor) - torch.testing.assert_close(output, torch.full((2, 2), 2.0)) + assert isinstance(output, sys.modules["torch"].Tensor), type(output) """ - native_name = self.load_fn.__module__ - native_path = sys.modules[native_name].__spec__.origin with tempfile.TemporaryDirectory() as directory: program_path = os.path.join(directory, "add.pte") with open(program_path, "wb") as handle: handle.write(program.buffer) - for entrypoint in dict.fromkeys((self.runtime.__name__, native_name)): - for policy in (None, "torch", "executorch"): - for available in ("yes", "no"): - with self.subTest( - entrypoint=entrypoint, policy=policy, available=available - ): - environment = os.environ.copy() - environment.pop("EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT", None) - if policy is not None: - environment["EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"] = ( - policy - ) - completed = subprocess.run( - [ - sys.executable, - "-c", - script, - entrypoint, - native_name, - str(policy), - available, - program_path, - native_path, - ], - env=environment, - capture_output=True, - text=True, - timeout=60, - ) - self.assertEqual( - completed.returncode, - 0, - completed.stdout + completed.stderr, - ) + # Without the setting, the call that finds torch missing is the + # one that would import it. + for case, setting in ( + ("none", "executorch"), + ("importing", "executorch"), + ("vanishing", "executorch"), + ("restored", None), + ("missing", None), + ): + with self.subTest(case=case): + environment = os.environ.copy() + environment.pop("EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT", None) + if setting is not None: + environment["EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"] = setting + completed = subprocess.run( + [sys.executable, "-c", script, program_path, case], + env=environment, + capture_output=True, + text=True, + timeout=60, + ) + self.assertEqual( + completed.returncode, 0, completed.stdout + completed.stderr + ) + + def test_executorch_result_from_host_refuses_unknown_dtype(self): + # Checked before any upload, so it needs no CUDA. Unlike the no-ATen + # suite, this also runs against ATen bindings. + with self.assertRaisesRegex(ValueError, "Unsupported dtype 'bfloat'"): + self.runtime.ExecuTorchResult.from_host( + np.zeros(2, dtype=np.uint16), dtype="bfloat" + ) def test_empty_numpy_buffer_input(self): exported_program, inputs = create_program(ModuleAddEmpty()) @@ -1319,7 +1731,10 @@ def __dlpack_device__(self): saved_stderr = sys.stderr sys.stderr = captured try: - module(LogsDuringTheCall()) + # Only the runtime log goes to stderr, not the output type warning. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + module(LogsDuringTheCall()) finally: sys.stderr = saved_stderr @@ -1420,7 +1835,7 @@ def test_method_rejects_aliased_input_whose_storage_changed(self): self.skipTest("only portable bindings return a view of the input") from executorch.extension.pybindings import _C - if _C._tensor_output != "torch": + if _C._tensor_output == "executorch": self.skipTest("only torch output mode returns a view of the input") class Identity(torch.nn.Module): @@ -3817,6 +4232,128 @@ def forward(self, x, y): method.execute(cuda_inputs)[0].cpu(), torch.full((4, 4), 4.0) ) + def test_cuda_input_imports_torch_to_wait_without_cuda_support(self): + if self.kernel_mode != "portable" or not sys.platform.startswith("linux"): + self.skipTest("needs the portable bindings and the Linux CUDA driver") + if not torch.cuda.is_available(): + self.skipTest("needs a visible CUDA device") + from executorch.extension.pybindings import portable_lib + + if "CudaBackend" in portable_lib._get_registered_backend_names(): + self.skipTest("bindings built with CUDA support wait without torch") + + 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(4),))).to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ) + ) + ) + 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=0 + ) + programs = {"identity": bytes(serialize_pte_binary(PTEFile(schema)))} + # The same method without outputs, which export cannot make. + schema.execution_plan[0].outputs = [] + programs["no outputs"] = bytes(serialize_pte_binary(PTEFile(schema))) + # Device memory from the CUDA driver, so nothing imports torch before + # the bindings wait for the device. + script = """ +import ctypes +import sys + +from executorch.extension.pybindings import portable_lib + +class Device(ctypes.Structure): + _fields_ = [("device_type", ctypes.c_int32), ("device_id", ctypes.c_int32)] + +class DataType(ctypes.Structure): + _fields_ = [("code", ctypes.c_uint8), ("bits", ctypes.c_uint8), ("lanes", ctypes.c_uint16)] + +class DLTensor(ctypes.Structure): + _fields_ = [ + ("data", ctypes.c_void_p), + ("device", Device), + ("ndim", ctypes.c_int32), + ("dtype", DataType), + ("shape", ctypes.POINTER(ctypes.c_int64)), + ("strides", ctypes.POINTER(ctypes.c_int64)), + ("byte_offset", ctypes.c_uint64), + ] + +class ManagedTensor(ctypes.Structure): + _fields_ = [("dl_tensor", DLTensor), ("manager_ctx", ctypes.c_void_p), ("deleter", ctypes.c_void_p)] + +driver = ctypes.CDLL("libcuda.so.1") +context, pointer = ctypes.c_void_p(), ctypes.c_uint64() +assert driver.cuInit(0) == 0 +assert driver.cuDevicePrimaryCtxRetain(ctypes.byref(context), 0) == 0 +assert driver.cuCtxSetCurrent(context) == 0 +assert driver.cuMemAlloc_v2(ctypes.byref(pointer), 16) == 0 +shape = (ctypes.c_int64 * 1)(4) +managed = ManagedTensor( + DLTensor(pointer.value, Device(2, 0), 1, DataType(2, 32, 1), shape, None, 0) +) +make_capsule = ctypes.pythonapi.PyCapsule_New +make_capsule.restype = ctypes.py_object +make_capsule.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + +case, path = sys.argv[1:] +with open(path, "rb") as source: + method = portable_lib._load_program_from_buffer(source.read()).load_method("forward") +assert "torch" not in sys.modules +inputs = [make_capsule(ctypes.addressof(managed), b"dltensor", None)] +if case == "call": + # It returns nothing, so picking the type of its outputs imports nothing. + assert method(inputs) == [] +else: + method.set_inputs(inputs) + try: + method.execute() + except RuntimeError as error: + # The setting keeps torch out, so nothing can wait for the device. + assert case == "kept out" and "or torch imported" in str(error), error + else: + assert case != "kept out", "the wait must fail when torch is kept out" +assert ("torch" in sys.modules) == (case != "kept out") +""" + with tempfile.TemporaryDirectory() as directory: + for case, name, setting in ( + ("execute", "identity", None), + ("call", "no outputs", None), + ("kept out", "identity", "executorch"), + ): + with self.subTest(case=case): + path = os.path.join(directory, "program.pte") + with open(path, "wb") as handle: + handle.write(programs[name]) + environment = os.environ.copy() + environment.pop("EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT", None) + if setting is not None: + environment["EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"] = setting + completed = subprocess.run( + [sys.executable, "-c", script, case, path], + env=environment, + capture_output=True, + text=True, + timeout=60, + ) + self.assertEqual( + completed.returncode, + 0, + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}", + ) + def test_cuda_torch_tensor_interchange_without_aten(self): if self.kernel_mode != "portable": self.skipTest("only portable bindings use the Python Torch adapter") diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index b7422ada832..bea6d7ccf84 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -1188,6 +1188,23 @@ def test_bfloat16_host_result_moves_to_the_device(self): with self.assertRaisesRegex(ValueError, "already on a device"): self.bindings.ExecuTorchResult.from_host(on_device, self.device.index) + def test_bfloat16_host_bits_move_to_the_device(self): + # Host data without torch has no bfloat16 buffer format, so it comes as + # 16-bit unsigned integers and from_host is told the dtype. + value = torch.arange(6, dtype=torch.bfloat16).reshape(2, 3) + words = array.array("H", value.view(torch.int16).flatten().tolist()) + host = memoryview(words).cast("B").cast("H", value.shape) + uploaded = self.bindings.ExecuTorchResult.from_host( + host, self.device.index, dtype="bfloat16" + ) + self.assertEqual(uploaded.dtype, "bfloat16") + self.assertEqual(uploaded.shape, (2, 3)) + data = self._device_identity(self.device.index, value) + method = self.bindings._load_program_from_buffer(data).load_method("forward") + torch.testing.assert_close( + self._tensor(method(uploaded)[0]), value.to(self.device) + ) + def _device_waits(self, body): from torch.profiler import profile, ProfilerActivity diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index 8488ac01e6c..a8ee126b4ff 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/extension/pybindings/test/test_pybindings_no_aten.py @@ -134,7 +134,7 @@ def test_output_policy_without_torch(self) -> None: import os assert Path(runtime.__file__).resolve() == Path(sys.argv[1]).resolve() - assert runtime._tensor_output == "executorch" + assert runtime._tensor_output is None with open(os.environ["EXECUTORCH_PYBIND_TEST_PTE"], "rb") as source: data = source.read() module = runtime._load_for_executorch_from_buffer(data) @@ -211,6 +211,21 @@ def test_executorch_result_to_host_and_repr(self) -> None: self.assertFalse(np.shares_memory(np.asarray(copy), np.asarray(output))) np.testing.assert_array_equal(np.asarray(copy), np.array([3.0])) + def test_executorch_result_is_not_a_bool_or_number(self) -> None: + if runtime._tensor_output != "executorch": + self.skipTest("ExecuTorchResult output was not selected") + 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) + output = module( + (np.array([1.0], dtype=np.float32), np.array([2.0], dtype=np.float32)) + )[0] + + for convert in (bool, int, float): + with self.subTest(convert=convert.__name__): + with self.assertRaisesRegex(TypeError, r"numpy\.asarray\(result\)"): + convert(output) + def test_executorch_result_from_host_requires_cuda(self) -> None: # A build with CUDA support uploads instead; the CUDA suite covers it. if "CudaBackend" in runtime._get_registered_backend_names(): @@ -220,6 +235,40 @@ def test_executorch_result_from_host_requires_cuda(self) -> None: with self.assertRaisesRegex(RuntimeError, "CUDA support"): runtime.ExecuTorchResult.from_host(np.ones(shape, dtype=np.float32)) + def test_executorch_result_from_host_checks_dtype(self) -> None: + # The checks come before the upload, so they need no CUDA. + words = np.zeros((2, 3), dtype=np.uint16) + for source, dtype, message in ( + (words.view(np.float16), "bfloat16", "cannot read a buffer of float16"), + (words.view(np.uint8), "bfloat16", "cannot read a buffer of uint8"), + (words, "int16", "cannot read a buffer of uint16"), + (words.view(np.uint8), "qint8", "cannot read a buffer of uint8"), + ): + with self.subTest(source=source.dtype, dtype=dtype): + with self.assertRaisesRegex(ValueError, message): + runtime.ExecuTorchResult.from_host(source, dtype=dtype) + # Data that comes the way results of the dtype export it passes the + # checks, so bindings without CUDA support stop at the upload. + if "CudaBackend" not in runtime._get_registered_backend_names(): + for source, dtype in ( + (words, "uint16"), + (words, "bfloat16"), + (words.view(np.int8), "qint8"), + (np.zeros(2, dtype=np.int32), "qint32"), + ): + with self.subTest(source=source.dtype, dtype=dtype): + with self.assertRaisesRegex(RuntimeError, "CUDA support"): + runtime.ExecuTorchResult.from_host(source, dtype=dtype) + 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) + output = module( + (np.array([1.0], dtype=np.float32), np.array([2.0], dtype=np.float32)) + )[0] + if isinstance(output, runtime.ExecuTorchResult): + with self.assertRaisesRegex(ValueError, "does not match the result's"): + runtime.ExecuTorchResult.from_host(output, dtype="bfloat16") + def test_borrowed_result_does_not_create_method_cycle(self) -> None: if runtime._tensor_output != "executorch": self.skipTest("ExecuTorchResult output was not selected") @@ -475,9 +524,9 @@ def test_z_torch_tensor_uses_python_api(self) -> None: output = module((torch.tensor([1.0]), torch.tensor([2.0])))[0] expected_type = ( - torch.Tensor - if runtime._tensor_output == "torch" - else runtime.ExecuTorchResult + runtime.ExecuTorchResult + if runtime._tensor_output == "executorch" + else torch.Tensor ) self.assertIsInstance(output, expected_type) np.testing.assert_array_equal(np.asarray(output), np.array([3.0])) diff --git a/runtime/__init__.py b/runtime/__init__.py index 0aa0d89655c..00b9c5a8b71 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -172,8 +172,20 @@ def execute( a later execution may overwrite them, which avoids a copy. Returns: - A list of output values, typically torch.Tensor objects. Without - PyTorch, tensor outputs are ExecuTorchResult objects. + A list of output values. When PyTorch is installed, tensor outputs + are torch.Tensor objects, and the first call that returns one + imports torch. Without PyTorch, they are ExecuTorchResult objects. + + .. warning:: + + A future release will make tensor outputs follow the tensor inputs: + torch.Tensor objects if any input is one, and ExecuTorchResult + objects if all are other tensors, even when PyTorch is installed. + With no tensor inputs, outputs stay torch.Tensor objects whenever + torch can be imported. Each call whose outputs will change warns + with a FutureWarning. Set ``EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT`` to + ``torch`` or ``executorch`` before importing ExecuTorch to choose + the type now. """ return self._method(inputs, clone_outputs=clone_outputs)