mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
95191522e0
# Implement OpenReg device autoload mechanism ## Overview The **Autoload** mechanism in PyTorch simplifies the integration of third-party device backends by enabling automatic discovery and initialization at runtime. Traditionally, integrating a new backend required explicit imports or manual initialization, which could be cumbersome and error-prone. With Autoload, PyTorch dynamically detects and initializes device backends, providing a seamless user experience. This mechanism leverages Python entry points (e.g., `torch.backends`) and dynamic module loading. When PyTorch starts, it scans for registered entry points and invokes their initialization hooks, ensuring that all available backends are ready for use without requiring explicit imports. ## Motivation This PR aims to apply [device autoload mechanism](https://github.com/pytorch/pytorch/issues/122468) to the OpenReg module with some simple changes. ## Change ### Before ```python import torch import torch_openreg x = torch.tensor([1, 2, 3], device="openreg") print(x) ``` ### After ```python import torch # No need to import torch_openreg manually! x = torch.tensor([1, 2, 3], device="openreg") print(x) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/158555 Approved by: https://github.com/FFFrog, https://github.com/albanD Co-authored-by: Jiawei Li <ljw1101.vip@gmail.com>
87 lines
3.5 KiB
Markdown
87 lines
3.5 KiB
Markdown
# Autoload Mechanism
|
||
|
||
The **Autoload** mechanism in PyTorch simplifies the integration of a custom backend by enabling automatic discovery and initialization at runtime. This eliminates the need for explicit imports or manual initialization, allowing developers to seamlessly integrate a new accelerator or backend into PyTorch.
|
||
|
||
## Background
|
||
|
||
The **Autoload Device Extension** proposal in PyTorch is centered on improving support for various hardware backend devices, especially those implemented as out-of-the-tree extensions (not part of PyTorch’s main codebase). Currently, users must manually import or load these device-specific extensions to use them, which complicates the experience and increases cognitive overhead.
|
||
|
||
In contrast, in-tree devices (devices officially supported within PyTorch) are seamlessly integrated—users don’t need extra imports or steps. The goal of autoloading is to make out-of-the-tree devices just as easy to use, so users can follow the standard PyTorch device programming model without explicit loading or code changes. This would allow existing PyTorch applications to run on new devices without any modification, making hardware support more user-friendly and reducing barriers to adoption.
|
||
|
||
For more information about the background of **Autoload**, please refer to its [RFC](https://github.com/pytorch/pytorch/issues/122468).
|
||
|
||
## Design
|
||
|
||
The core idea of **Autoload** is to Use Python’s plugin discovery (entry points) so PyTorch automatically loads out-of-tree device extensions when torch is imported—no explicit user import needed.
|
||
|
||
For more instructions of the design of **Autoload**, please refer to [**How it works**](https://docs.pytorch.org/tutorials/unstable/python_extension_autoload.html#how-it-works).
|
||
|
||
## Implementation
|
||
|
||
This tutorial will take **OpenReg** as a new out-of-the-tree device and guide you through the steps to enable and use the **Autoload** mechanism.
|
||
|
||
### Entry Point Setup
|
||
|
||
To enable **Autoload**, register the `_autoload` function as an entry point in `setup.py` file.
|
||
|
||
::::{tab-set}
|
||
|
||
:::{tab-item} Python
|
||
|
||
```{eval-rst}
|
||
.. literalinclude:: ../../../test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
|
||
:language: python
|
||
:start-after: LITERALINCLUDE START: SETUP
|
||
:end-before: LITERALINCLUDE END: SETUP
|
||
:linenos:
|
||
:emphasize-lines: 9-13
|
||
```
|
||
|
||
:::
|
||
|
||
::::
|
||
|
||
### Backend Setup
|
||
|
||
Define the initialization hook `_autoload` for backend initialization. This hook will be automatically invoked by PyTorch during startup.
|
||
|
||
::::{tab-set-code}
|
||
```{eval-rst}
|
||
.. literalinclude:: ../../../test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
|
||
:language: python
|
||
:start-after: LITERALINCLUDE START: AUTOLOAD
|
||
:end-before: LITERALINCLUDE END: AUTOLOAD
|
||
:linenos:
|
||
:emphasize-lines: 10-12
|
||
```
|
||
|
||
|
||
::::
|
||
|
||
## Result
|
||
|
||
After setting up the entry point and backend, build and install your backend. Now, we can use the new accelerator without explicitly importing it.
|
||
|
||
```{eval-rst}
|
||
.. grid:: 2
|
||
|
||
.. grid-item-card:: :octicon:`terminal;1em;` Without Autoload
|
||
:class-card: card-prerequisites
|
||
|
||
::
|
||
|
||
>>> import torch
|
||
>>> import torch_openreg
|
||
>>> torch.tensor(1, device="openreg")
|
||
tensor(1, device='openreg:0')
|
||
|
||
.. grid-item-card:: :octicon:`terminal;1em;` With Autoload
|
||
:class-card: card-prerequisites
|
||
|
||
::
|
||
|
||
>>> import torch # Automatically import torch_openreg
|
||
>>> torch.tensor(1, device="openreg")
|
||
tensor(1, device='openreg:0')
|
||
```
|