From 53166efa2494764e11734d57de681cc4d2486981 Mon Sep 17 00:00:00 2001 From: Alfie Edwards Date: Mon, 7 Oct 2024 15:29:13 +0100 Subject: [PATCH] Allow third-party backends to add submodules to `triton.language.extra` (#4503) Add an optional language directory to backends. The contents of the directory is added to `triton.language.extra` when the wheel is built. Update the existing `triton.language.extra.cuda` and `triton.language.extra.hip` modules to use the new mechanism. The core Triton is a small number of people, and we receive many PRs (thank you!). To help us review your code more quickly, **if you are a new contributor (less than 3 PRs merged) we ask that you complete the following tasks and include the filled-out checklist in your PR description.** Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. - [x] I am not making a trivial change, such as fixing a typo in a comment. - [x] I have written a PR description following these [rules](https://cbea.ms/git-commit/#why-not-how). - [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`. - Select one of the following. - [ ] I have added tests. - `/test` for `lit` tests - `/unittest` for C++ tests - `/python/test` for end-to-end tests - [x] This PR does not need a test because `It is already tested by python/test/unit/language/test_core.py::test_math_extern`. - Select one of the following. - [x] I have not added any `lit` tests. - [ ] The `lit` tests I have added follow these [best practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices), including the "tests should be minimal" section. (Usually running Python code and using the instructions it generates is not minimal.) --- python/setup.py | 60 ++++++++++++++++--- python/triton/language/extra/__init__.py | 4 -- .../amd/language}/hip/__init__.py | 0 .../amd/language}/hip/libdevice.py | 0 .../nvidia/language}/cuda/__init__.py | 0 .../language}/cuda/_experimental_tma.py | 0 .../nvidia/language}/cuda/libdevice.py | 0 .../nvidia/language}/cuda/utils.py | 0 8 files changed, 53 insertions(+), 11 deletions(-) rename {python/triton/language/extra => third_party/amd/language}/hip/__init__.py (100%) rename {python/triton/language/extra => third_party/amd/language}/hip/libdevice.py (100%) rename {python/triton/language/extra => third_party/nvidia/language}/cuda/__init__.py (100%) rename {python/triton/language/extra => third_party/nvidia/language}/cuda/_experimental_tma.py (100%) rename {python/triton/language/extra => third_party/nvidia/language}/cuda/libdevice.py (100%) rename {python/triton/language/extra => third_party/nvidia/language}/cuda/utils.py (100%) diff --git a/python/setup.py b/python/setup.py index 01de47d52df0..daff988ed0c6 100644 --- a/python/setup.py +++ b/python/setup.py @@ -32,9 +32,11 @@ @dataclass class Backend: name: str - package_data: dict + package_data: list[str] + language_package_data: list[str] src_dir: str backend_dir: str + language_dir: str install_dir: str is_external: bool @@ -62,12 +64,22 @@ def prepare(backend_name: str, backend_src_dir: str = None, is_external: bool = backend_path = os.path.abspath(os.path.join(backend_src_dir, "backend")) assert os.path.exists(backend_path), f"{backend_path} does not exist!" + language_dir = os.path.abspath(os.path.join(backend_src_dir, "language")) + if not os.path.exists(language_dir): + language_dir = None + for file in ["compiler.py", "driver.py"]: assert os.path.exists(os.path.join(backend_path, file)), f"${file} does not exist in ${backend_path}" install_dir = os.path.join(os.path.dirname(__file__), "triton", "backends", backend_name) package_data = [f"{os.path.relpath(p, backend_path)}/*" for p, _, _, in os.walk(backend_path)] - return Backend(name=backend_name, package_data=package_data, src_dir=backend_src_dir, backend_dir=backend_path, + + language_package_data = [] + if language_dir is not None: + language_package_data = [f"{os.path.relpath(p, language_dir)}/*" for p, _, _, in os.walk(language_dir)] + + return Backend(name=backend_name, package_data=package_data, language_package_data=language_package_data, + src_dir=backend_src_dir, backend_dir=backend_path, language_dir=language_dir, install_dir=install_dir, is_external=is_external) # Copy all in-tree backends under triton/third_party. @@ -556,6 +568,19 @@ def add_link_to_backends(): shutil.rmtree(backend.install_dir) os.symlink(backend.backend_dir, backend.install_dir) + if backend.language_dir: + # Link the contents of each backend's `language` directory into + # `triton.language.extra`. + extra_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "triton", "language", "extra")) + for x in os.listdir(backend.language_dir): + src_dir = os.path.join(backend.language_dir, x) + install_dir = os.path.join(extra_dir, x) + if os.path.islink(install_dir): + os.unlink(install_dir) + if os.path.exists(install_dir): + shutil.rmtree(install_dir) + os.symlink(src_dir, install_dir) + def add_link_to_proton(): proton_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, "third_party", "proton", "proton")) @@ -602,12 +627,33 @@ def run(self): package_data = { - "triton/tools": ["compile.h", "compile.c"], - **{f"triton/backends/{b.name}": b.package_data - for b in backends}, + "triton/tools": ["compile.h", "compile.c"], **{f"triton/backends/{b.name}": b.package_data + for b in backends}, "triton/language/extra": sum( + (b.language_package_data for b in backends), []) } +def get_language_extra_packages(): + packages = [] + for backend in backends: + if backend.language_dir is None: + continue + + # Walk the `language` directory of each backend to enumerate + # any subpackages, which will be added to `triton.language.extra`. + for dir, dirs, files in os.walk(backend.language_dir, followlinks=True): + if not any(f for f in files if f.endswith(".py")) or dir == backend.language_dir: + # Ignore directories with no python files. + # Also ignore the root directory which corresponds to + # "triton/language/extra". + continue + subpackage = os.path.relpath(dir, backend.language_dir) + package = os.path.join("triton/language/extra", subpackage) + packages.append(package) + + return list(packages) + + def get_packages(): packages = [ "triton", @@ -615,15 +661,15 @@ def get_packages(): "triton/compiler", "triton/language", "triton/language/extra", - "triton/language/extra/cuda", - "triton/language/extra/hip", "triton/runtime", "triton/backends", "triton/tools", ] packages += [f'triton/backends/{backend.name}' for backend in backends] + packages += get_language_extra_packages() if check_env_flag("TRITON_BUILD_PROTON", "ON"): # Default ON packages += ["triton/profiler"] + return packages diff --git a/python/triton/language/extra/__init__.py b/python/triton/language/extra/__init__.py index 14e1778d2fdc..e69de29bb2d1 100644 --- a/python/triton/language/extra/__init__.py +++ b/python/triton/language/extra/__init__.py @@ -1,4 +0,0 @@ -from . import cuda -from . import hip - -__all__ = ['cuda', 'hip'] diff --git a/python/triton/language/extra/hip/__init__.py b/third_party/amd/language/hip/__init__.py similarity index 100% rename from python/triton/language/extra/hip/__init__.py rename to third_party/amd/language/hip/__init__.py diff --git a/python/triton/language/extra/hip/libdevice.py b/third_party/amd/language/hip/libdevice.py similarity index 100% rename from python/triton/language/extra/hip/libdevice.py rename to third_party/amd/language/hip/libdevice.py diff --git a/python/triton/language/extra/cuda/__init__.py b/third_party/nvidia/language/cuda/__init__.py similarity index 100% rename from python/triton/language/extra/cuda/__init__.py rename to third_party/nvidia/language/cuda/__init__.py diff --git a/python/triton/language/extra/cuda/_experimental_tma.py b/third_party/nvidia/language/cuda/_experimental_tma.py similarity index 100% rename from python/triton/language/extra/cuda/_experimental_tma.py rename to third_party/nvidia/language/cuda/_experimental_tma.py diff --git a/python/triton/language/extra/cuda/libdevice.py b/third_party/nvidia/language/cuda/libdevice.py similarity index 100% rename from python/triton/language/extra/cuda/libdevice.py rename to third_party/nvidia/language/cuda/libdevice.py diff --git a/python/triton/language/extra/cuda/utils.py b/third_party/nvidia/language/cuda/utils.py similarity index 100% rename from python/triton/language/extra/cuda/utils.py rename to third_party/nvidia/language/cuda/utils.py