Use entry points to locate backends

Use entry points to locate Triton backends instead of iterating over
the `triton.backends` package directory.  This fixes support for
nontrivial install layouts, particularly editable installs made using
`package_dir`.  Now `setup.py` keeps record of all installed backends
in entry points, and `triton.backends` uses them to locate and load
the backends.  This also technically permits installing third-party
backends without altering the Triton installation, by adding additional
entry points.

Author: mgorny

python/triton/backends/__init__.py

@@ -1,10 +1,14 @@
 import importlib
-import inspect
-import os
+import sys
 from dataclasses import dataclass
 from .driver import DriverBase
 from .compiler import BaseBackend
 
+if sys.version_info >= (3, 10):
+    from importlib.metadata import entry_points
+else:
+    from importlib_metadata import entry_points
+
 
 def _find_concrete_subclasses(module, base_class):
     ret = []
@@ -27,16 +31,11 @@ class Backend:
 
 def _discover_backends():
     backends = dict()
-    root = os.path.dirname(__file__)
-    for name in os.listdir(root):
-        if not os.path.isdir(os.path.join(root, name)):
-            continue
-        if name.startswith('__'):
-            continue
-        compiler = importlib.import_module(f"triton.backends.{name}.compiler")
-        driver = importlib.import_module(f"triton.backends.{name}.driver")
-        backends[name] = Backend(_find_concrete_subclasses(compiler, BaseBackend),
-                                 _find_concrete_subclasses(driver, DriverBase))
+    for ep in entry_points().select(group="triton.backends"):
+        compiler = importlib.import_module(f"{ep.value}.compiler")
+        driver = importlib.import_module(f"{ep.value}.driver")
+        backends[ep.name] = Backend(_find_concrete_subclasses(compiler, BaseBackend),
+                                    _find_concrete_subclasses(driver, DriverBase))
     return backends
 
 

setup.py

@@ -598,6 +598,7 @@ def get_entry_points():
             "proton-viewer = triton.profiler.viewer:main",
             "proton = triton.profiler.proton:main",
         ]
+    entry_points["triton.backends"] = [f"{b.name} = triton.backends.{b.name}" for b in backends]
     return entry_points
 
 
@@ -638,7 +639,10 @@ def get_git_version_suffix():
     author_email="[email protected]",
     description="A language and compiler for custom Deep Learning operations",
     long_description="",
-    install_requires=["setuptools>=40.8.0"],
+    install_requires=[
+        "setuptools>=40.8.0",
+        "importlib-metadata; python_version < '3.10'",
+    ],
     packages=find_packages(where="python") + extra_packages,
     package_dir=package_dirs,
     entry_points=get_entry_points(),