[BUILD] Support building out of tree plugins

CMakeLists.txt

@@ -145,6 +145,13 @@ if(TRITON_BUILD_PYTHON_MODULE)
     add_link_options(${Python3_LINK_OPTIONS})
   endif()
 
+  if (DEFINED TRITON_EXTERNAL_CODEGEN_BACKEND_DIRS)
+    foreach(BACKEND_DIR ${TRITON_EXTERNAL_CODEGEN_BACKEND_DIRS})
+      cmake_path(APPEND BACKEND_DIR "build" OUTPUT_VARIABLE BACKEND_DIR_BUILD_OUTPUT)
+      add_subdirectory(${BACKEND_DIR} ${BACKEND_DIR_BUILD_OUTPUT})
+    endforeach()
+  endif()
+
   foreach(CODEGEN_BACKEND ${TRITON_CODEGEN_BACKENDS})
     add_subdirectory(third_party/${CODEGEN_BACKEND})
   endforeach()
@@ -191,6 +198,12 @@ if(TRITON_BUILD_PYTHON_MODULE)
 
   # Define triton library
   string(JOIN "," TRITON_BACKENDS_TUPLE ${TRITON_CODEGEN_BACKENDS})
+
+  if (DEFINED TRITON_EXTERNAL_CODEGEN_BACKEND_DIRS)
+    string(JOIN "," TRITON_BACKENDS_TUPLE ${TRITON_BACKENDS_TUPLE} ${TRITON_EXTERNAL_CODEGEN_BACKENDS})
+  endif()
+
+
   set(TRITON_BACKENDS_TUPLE "(${TRITON_BACKENDS_TUPLE})")
   add_compile_definitions(TRITON_BACKENDS_TUPLE=${TRITON_BACKENDS_TUPLE})
   add_library(triton SHARED ${PYTHON_SRC_PATH}/main.cc

python/setup.py

@@ -28,6 +28,7 @@ class Backend:
     src_dir: str
     backend_dir: str
     install_dir: str
+    is_external: bool
 
 
 def _copy_backends(active):
@@ -54,7 +55,42 @@ def _copy_backends(active):
         package_data = [f"{os.path.relpath(p, backend_path)}/*" for p, _, _, in os.walk(backend_path)]
         ret.append(
             Backend(name=backend, package_data=package_data, src_dir=curr_path, backend_dir=backend_path,
-                    install_dir=install_dir))
+                    install_dir=install_dir, is_external=False))
+    return ret
+
+def _copy_external_backends():
+    # TRITON_EXTERNAL_CODEGEN_BACKENDS is a semicolon-separated list of backend names
+    # TRITON_EXTERNAL_CODEGEN_BACKEND_DIRS is a semicolon-separated list of paths to the backends
+    # Note that the paths have to include the provided backend names as a subfolder
+    backends = os.getenv('TRITON_EXTERNAL_CODEGEN_BACKENDS')
+    backend_dirs = os.getenv('TRITON_EXTERNAL_CODEGEN_BACKEND_DIRS')
+
+    if backends is None or backend_dirs is None:
+        return []
+
+    backends = backends.strip().split(';')
+    backend_dirs = backend_dirs.strip().split(';')
+
+    if len(backends) != len(backend_dirs):
+        raise Exception(f"External backend names and directories mismatch")
+
+    ret = []
+
+    for backend, root_dir in zip(backends, backend_dirs):
+        curr_path = os.path.join(root_dir, backend)
+        backend_path = os.path.abspath(os.path.join(curr_path, "backend"))
+        install_dir = os.path.join(os.path.dirname(__file__), "triton", "backends", backend)
+
+        # check conditions
+        assert backend in os.listdir(root_dir), f"{backend} is requested for install but not present in {root_dir}"
+        assert os.path.exists(backend_path), f"{backend_path} does not exist!"
+        for file in ["compiler.py", "driver.py"]:
+            assert os.path.exists(os.path.join(backend_path, file))
+        # update
+        package_data = [f"{os.path.relpath(p, backend_path)}/*" for p, _, _, in os.walk(backend_path)]
+        ret.append(
+            Backend(name=backend, package_data=package_data, src_dir=curr_path, backend_dir=backend_path,
+                    install_dir=install_dir, is_external=True))
     return ret
 
 
@@ -297,7 +333,9 @@ def build_extension(self, ext):
             "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir, "-DTRITON_BUILD_TUTORIALS=OFF",
             "-DTRITON_BUILD_PYTHON_MODULE=ON", "-DPython3_EXECUTABLE:FILEPATH=" + sys.executable,
             "-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON", "-DPYTHON_INCLUDE_DIRS=" + python_include_dir,
-            "-DTRITON_CODEGEN_BACKENDS=" + ';'.join([b.name for b in backends])
+            "-DTRITON_CODEGEN_BACKENDS=" + ';'.join([b.name for b in backends if not b.is_external]),
+            "-DTRITON_EXTERNAL_CODEGEN_BACKENDS=" + ';'.join([b.name for b in backends if b.is_external]),
+            "-DTRITON_EXTERNAL_CODEGEN_BACKEND_DIRS=" + ';'.join([b.src_dir for b in backends if b.is_external])
         ]
         if lit_dir is not None:
             cmake_args.append("-DLLVM_EXTERNAL_LIT=" + lit_dir)
@@ -380,8 +418,8 @@ def build_extension(self, ext):
     url_func=lambda arch, version:
     f"https://anaconda.org/nvidia/cuda-nvdisasm/12.3.52/download/linux-{arch}/cuda-nvdisasm-{version}-0.tar.bz2",
 )
-backends = _copy_backends(["nvidia", "amd"])
 
+backends = [*_copy_backends(["nvidia", "amd"]), *_copy_external_backends()]
 
 def add_link_to_backends():
     for backend in backends: