update

Author: shink

advanced_source/python_extension_autoload.rst

@@ -11,31 +11,29 @@ experience and enables users to adhere to the familiar PyTorch device
 programming model without needing to explicitly load or import device-specific
 extensions. On the other hand, it facilitates effortless
 adoption of existing PyTorch applications with zero-code changes on
-out-of-tree devices. For more information,
-see `[RFC] Autoload Device Extension <https://github.com/pytorch/pytorch/issues/122468>`_.
+out-of-tree devices. For further details, refer to the
+`[RFC] Autoload Device Extension <https://github.com/pytorch/pytorch/issues/122468>`_.
 
 .. note::
 
     This feature is enabled by default and can be disabled using
     ``export TORCH_DEVICE_BACKEND_AUTOLOAD=0``.
     If you get an error like this: "Failed to load the backend extension",
-    this error has nothing to do with PyTorch, you should disable this feature
+    this error is independent with PyTorch, you should disable this feature
     and ask the out-of-tree extension maintainer for help.
 
 How to apply this mechanism to out-of-tree extensions?
 ------------------------------------------------------
 
-For example, if you have a backend named ``foo`` and a package named
-``torch_foo``. Make sure your package is based on PyTorch 2.5+ and includes
-the following in its ``__init__.py``:
+For instance, suppose you have a backend named ``foo`` and a corresponding package named ``torch_foo``. Ensure that
+your package is compatible with PyTorch 2.5+ and includes the following snippet in its ``__init__.py`` file:
 
 .. code-block:: python
 
     def _autoload():
         print("No need to import torch_foo anymore! Check things are working with `torch.foo.is_available()`.")
 
-Then the only thing you need to do is add an entry point to your Python
-package:
+Then the only thing you need to do is define an entry point within your Python package:
 
 .. code-block:: python
 
@@ -62,8 +60,7 @@ Examples
 ^^^^^^^^
 
 Here we take Intel Gaudi HPU and Huawei Ascend NPU as examples to determine how to
-integrate your out-of-tree extension with PyTorch based on the autoloading
-mechanism.
+integrate your out-of-tree extension with PyTorch using the autoloading mechanism.
 
 `habana_frameworks.torch`_ is a Python package that enables users to run
 PyTorch programs on Intel Gaudi via the PyTorch ``HPU`` device key.
@@ -72,24 +69,58 @@ is applied.
 
 .. _habana_frameworks.torch: https://docs.habana.ai/en/latest/PyTorch/Getting_Started_with_PyTorch_and_Gaudi/Getting_Started_with_PyTorch.html
 
+``habana_frameworks.torch`` is a submodule of ``habana_frameworks``, we add an entry point to
+``__autoload()`` in ``habana_frameworks/setup.py``:
+
 .. code-block:: diff
 
-    import torch
-    import torchvision.models as models
-    - import habana_frameworks.torch # <-- extra import
-    model = models.resnet50().eval().to("hpu")
-    input = torch.rand(128, 3, 224, 224).to("hpu")
-    output = model(input)
+    setup(
+        name="habana_frameworks",
+        version="2.5",
+    +   entry_points={
+    +       'torch.backends': [
+    +           "device_backend = habana_frameworks:__autoload",
+    +       ],
+    +   }
+    )
+
+In ``habana_frameworks/init.py``, we use a global variable to track if our module has been loaded:
+
+.. code-block:: python
+
+    import os
+
+    is_loaded = False  # A member variable of habana_frameworks module to track if our module has been imported
+
+    def __autoload():
+        # This is an entrypoint for pytorch autoload mechanism
+        # If the following condition is true, that means our backend has already been loaded, either explicitly
+        # or by the autoload mechanism and importing it again should be skipped to avoid circular imports
+        global is_loaded
+        if is_loaded:
+            return
+        import habana_frameworks.torch
+
+In ``habana_frameworks/torch/init.py``, We prevent circular imports by updating the state of the global variable:
+
+.. code-block:: python
+
+    import os
+
+    # This is to prevent torch autoload mechanism from causing circular imports
+    import habana_frameworks
+
+    habana_frameworks.is_loaded = True
 
 `torch_npu`_ enables users to run PyTorch programs on Huawei Ascend NPU, it
 leverages the ``PrivateUse1`` device key and exposes the device name
 as ``npu`` to the end users.
 
 .. _torch_npu: https://github.com/Ascend/pytorch
 
-Define an entry point in `torch_npu/setup.py`_:
+We define an entry point in `torch_npu/setup.py`_:
 
-.. _torch_npu/setup.py: https://github.com/Ascend/pytorch/blob/c164fbd5bb74790191ff8496b77d620fddf806d8/setup.py#L618
+.. _torch_npu/setup.py: https://github.com/Ascend/pytorch/blob/master/setup.py#L618
 
 .. code-block:: diff
 
@@ -103,16 +134,14 @@ Define an entry point in `torch_npu/setup.py`_:
     +   }
     )
 
-``import torch_npu`` is also no longer needed after applying this mechanism:
+Unlike ``habana_frameworks``, ``torch_npu`` uses the environment variable ``TORCH_DEVICE_BACKEND_AUTOLOAD``
+to control the autoloading process. For example, we set it to `0` to disable autoloading to prevent circular imports:
 
-.. code-block:: diff
+.. code-block:: python
+    # Disable autoloading before running 'import torch'
+    os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'
 
     import torch
-    import torchvision.models as models
-    - import torch_npu # <-- extra import
-    model = models.resnet50().eval().to("npu")
-    input = torch.rand(128, 3, 224, 224).to("npu")
-    output = model(input)
 
 How it works
 ------------