Import torch only when the Python runtime returns torch tensors, and warn before outputs follow inputs

Importing executorch.runtime imports PyTorch whenever it is installed, even
for a program that only passes NumPy arrays. Only setting
EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch avoids that, and never
loading PyTorch for NumPy callers would change what they get back. So that
change has to be announced first.

Without the setting, importing executorch.runtime no longer imports torch.
The outputs stay as they are today: torch.Tensor whenever PyTorch is
installed, and ExecuTorchResult without it. So a program with only NumPy
inputs still loads torch, at its first call that returns a tensor.

A future release will make tensor outputs follow each call's tensor inputs,
including inputs kept from an earlier call:

- any torch.Tensor input: torch.Tensor outputs;
- only NumPy, other buffer, DLPack or ExecuTorchResult inputs:
  ExecuTorchResult outputs, even when PyTorch is installed;
- no tensor inputs: torch.Tensor outputs only if torch is already imported.

get_attribute() and bundled program verification count as calls without
tensor inputs.

The first call whose outputs will change shows one FutureWarning, which
says how to keep torch.Tensor outputs or get ExecuTorchResult outputs now.
A call that runs the method picks the type first, so with an error filter
for this FutureWarning, such a call raises before it runs the method (its
inputs are already installed). A torch that is installed but fails to
import, even with an ImportError, raises its own error, with a note on how
to skip torch. EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT still fixes the type:
"executorch" never imports torch, and "torch" imports it when ExecuTorch is
imported.

Tested on Linux x86_64 with Python 3.12 and a CUDA build of PyTorch: the
pybindings, runtime, bundled program and CUDA DLPack suites give the same
results as on main, and the standalone no-ATen job passes with and without
PyTorch. Importing executorch.runtime takes 0.04 s and 18 MB of peak memory
instead of 1.9 s and 506 MB. A new test runs each case of the rule in a
fresh process, and each of its checks fails when its part of the change is
removed.

Author: shoumikhin

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, it returns ``ExecuTorchResult`` outputs. If PyTorch is installed but fails to import, every call that returns a tensor raises the import error; 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 will return ``torch.Tensor`` outputs only if ``torch`` is already imported. The first call whose outputs will change shows 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
 ~~~~~~~~~~~~

extension/pybindings/README.md

@@ -41,10 +41,14 @@ 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
-DLPack; the extension does not link against PyTorch.
+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, outputs are `ExecuTorchResult`
+values. When it is installed but fails to import, every call that returns a
+tensor raises the import error; 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
 runtime or input storage and later execution may overwrite that storage.
@@ -91,9 +95,16 @@ 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 will return torch tensors only if `torch` is already imported. The first
+call whose outputs will change shows 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,
+and `executorch` never imports PyTorch.
 
 ```python
 import os
@@ -105,11 +116,9 @@ 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. 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)`. 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.
+`TypeError`, so convert it first, for example with `np.asarray(output)`.
+Explicitly selecting `torch` requires it to import successfully. ATen-mode
+bindings continue to return native PyTorch tensors directly.
 
 ```python
 output = module.forward([numpy_input])[0]
@@ -138,8 +147,9 @@ result = np.asarray(output.to_host())
 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.
+so once torch is imported, by you or by a call that returns torch tensors, 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

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