Add device autoload support to openreg

can-gaa-hou · can-gaa-hou · commit e1d76f2d8383 · 2025-07-17T11:38:12.000Z
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/README.md b/test/cpp_extensions/open_registration_extension/torch_openreg/README.md
@@ -125,7 +125,7 @@ graph LR
 
 ### Autoload
 
-- (TODO)
+- Device Autoload Mechanism
 
 ...
 
@@ -145,7 +145,6 @@ After installation, you can use the `openreg` device in Python just like any oth
 
 ```python
 import torch
-import torch_openreg
 
 if not torch.openreg.is_available():
     print("OpenReg backend is not available in this build.")
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -95,6 +95,11 @@ def main():
         cmdclass={
             "clean": BuildClean,  # type: ignore[misc]
         },
+        entry_points={
+            "torch.backends": [
+                "device_backend = torch_openreg:_autoload",
+            ],
+        },
     )
 
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
@@ -1,8 +1,13 @@
+import os
+
 import torch
 import torch_openreg._C  # type: ignore[misc]
 import torch_openreg.openreg
 
 
-torch.utils.rename_privateuse1_backend("openreg")
-torch._register_device_module("openreg", torch_openreg.openreg)
-torch.utils.generate_methods_for_privateuse1_backend(for_storage=True)
+def _autoload():
+    os.environ["IS_CUSTOM_DEVICE_BACKEND_IMPORTED"] = "1"
+
+    torch.utils.rename_privateuse1_backend("openreg")
+    torch._register_device_module("openreg", torch_openreg.openreg)
+    torch.utils.generate_methods_for_privateuse1_backend(for_storage=True)
