Triton v3.6.x iluvatar backend and 5 TLE primitives support #724
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,128 @@ | ||
| import inspect | ||
| import os | ||
| import shutil | ||
| import sys | ||
| from pathlib import Path | ||
|
|
||
| from setuptools import find_packages | ||
|
|
||
| FLAGTREE_BACKEND = os.environ.get("FLAGTREE_BACKEND", "iluvatar") | ||
| PYTHON_ROOT = f"third_party/{FLAGTREE_BACKEND}/python" | ||
| FLAGTREE_PYTHON_ROOT = "python" | ||
| TLE_PACKAGE = "triton.experimental.tle" | ||
| SKIP_BACKEND_PACKAGES_IN_PACKAGE_DIR = True | ||
|
|
||
|
|
||
| def _is_backend_package(package): | ||
| return package == "triton.backends" or package.startswith("triton.backends.") | ||
|
|
||
|
|
||
| def _is_language_extra_package(package): | ||
| return package == "triton.language.extra" or package.startswith("triton.language.extra.") | ||
|
|
||
|
|
||
| def _build_setup_hooks(backend, python_root, skip_backend_packages_in_package_dir=False): | ||
| patched_attr = f"_{backend}_python_root_patched" | ||
|
|
||
| def merge_packages(existing_packages): | ||
| packages = [] | ||
| seen = set() | ||
|
|
||
| def add(package): | ||
| if package not in seen: | ||
| packages.append(package) | ||
| seen.add(package) | ||
|
|
||
| for package in find_packages(where=python_root, include=["triton", "triton.*"]): | ||
| add(package) | ||
|
|
||
| for package in find_packages(where=FLAGTREE_PYTHON_ROOT, include=[TLE_PACKAGE, f"{TLE_PACKAGE}.*"]): | ||
| add(package) | ||
|
|
||
| for package in existing_packages: | ||
| if (not package.startswith("triton.") or _is_backend_package(package) | ||
| or _is_language_extra_package(package) or package == "triton.profiler" | ||
| or package.startswith("triton.profiler.")): | ||
| add(package) | ||
|
|
||
| return packages | ||
|
|
||
| def merge_package_dir(existing_package_dir): | ||
| package_dir = dict(existing_package_dir or {}) | ||
| package_dir[""] = python_root | ||
|
|
||
| for package in find_packages(where=python_root, include=["triton", "triton.*"]): | ||
| if skip_backend_packages_in_package_dir and package.startswith("triton.backends."): | ||
| continue | ||
| rel_package_path = package.replace(".", "/") | ||
| package_dir[package] = f"{python_root}/{rel_package_path}" | ||
|
|
||
| for package in find_packages(where=FLAGTREE_PYTHON_ROOT, include=[TLE_PACKAGE, f"{TLE_PACKAGE}.*"]): | ||
| rel_package_path = package.replace(".", "/") | ||
| package_dir[package] = f"{FLAGTREE_PYTHON_ROOT}/{rel_package_path}" | ||
|
|
||
| return package_dir | ||
|
|
||
| def patch_cmdclass(existing_cmdclass): | ||
| cmdclass = dict(existing_cmdclass or {}) | ||
| original_build_py = cmdclass.get("build_py") | ||
| if original_build_py is None: | ||
| return cmdclass | ||
|
|
||
| class BackendBuildPy(original_build_py): | ||
|
|
||
| def run(self): | ||
| self.force = True | ||
| build_triton_dir = Path(self.build_lib) / "triton" | ||
| if build_triton_dir.exists(): | ||
| shutil.rmtree(build_triton_dir) | ||
| return super().run() | ||
|
|
||
| cmdclass["build_py"] = BackendBuildPy | ||
| return cmdclass | ||
|
|
||
| def wrap_setup(original_setup): | ||
| if getattr(original_setup, patched_attr, False): | ||
| return original_setup | ||
|
|
||
| def setup_with_backend_python_root(*args, **kwargs): | ||
| kwargs["packages"] = merge_packages(kwargs.get("packages", [])) | ||
| kwargs["package_dir"] = merge_package_dir(kwargs.get("package_dir", {})) | ||
| kwargs["cmdclass"] = patch_cmdclass(kwargs.get("cmdclass", {})) | ||
| return original_setup(*args, **kwargs) | ||
|
|
||
| setattr(setup_with_backend_python_root, patched_attr, True) | ||
| setup_with_backend_python_root._backend_python_root_original_setup = original_setup | ||
| return setup_with_backend_python_root | ||
|
|
||
| return wrap_setup | ||
|
|
||
|
|
||
| def _patch_setup(wrap_setup): | ||
| patched = False | ||
|
|
||
| frame = inspect.currentframe() | ||
| while frame is not None: | ||
| setup_func = frame.f_globals.get("setup") | ||
| if callable(setup_func): | ||
| frame.f_globals["setup"] = wrap_setup(setup_func) | ||
| patched = True | ||
| frame = frame.f_back | ||
|
|
||
| main_module = sys.modules.get("__main__") | ||
| if main_module is not None and hasattr(main_module, "setup"): | ||
| main_module.setup = wrap_setup(main_module.setup) | ||
| patched = True | ||
|
|
||
| if not patched: | ||
| raise RuntimeError( | ||
| f"{FLAGTREE_BACKEND} setup hook could not find setup() to patch " | ||
| f"(python root: {PYTHON_ROOT})") | ||
|
|
||
|
|
||
| _wrap_setup = _build_setup_hooks( | ||
| FLAGTREE_BACKEND, | ||
| PYTHON_ROOT, | ||
| skip_backend_packages_in_package_dir=SKIP_BACKEND_PACKAGES_IN_PACKAGE_DIR, | ||
| ) | ||
| _patch_setup(_wrap_setup) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,36 @@ | ||
| include_directories(${CMAKE_CURRENT_SOURCE_DIR}/backend/include) | ||
| include_directories(${CMAKE_CURRENT_BINARY_DIR}/backend/include) | ||
| if(FLAGTREE_ILUVATAR_TLE) | ||
| include_directories(${CMAKE_CURRENT_SOURCE_DIR}/tle/include) | ||
| include_directories(${CMAKE_CURRENT_BINARY_DIR}/tle/include) | ||
| add_subdirectory(tle) | ||
| endif() | ||
| add_subdirectory(backend/include) | ||
| add_subdirectory(backend/lib) | ||
| if(TRITON_BUILD_PYTHON_MODULE) | ||
| find_package(LLD REQUIRED CONFIG PATHS "${LLD_DIR}" NO_DEFAULT_PATH) | ||
| include_directories(${LLD_INCLUDE_DIRS}) | ||
| message(STATUS "Found LLD distro-package @ ${LLD_DIR} and LLD include dirs @ ${LLD_INCLUDE_DIRS}") | ||
| if(FLAGTREE_ILUVATAR_TLE) | ||
| set(_ILUVATAR_TLE_PLUGIN_SOURCES | ||
| ${CMAKE_CURRENT_SOURCE_DIR}/tle/triton_iluvatar_tle.cc) | ||
| set(_ILUVATAR_TLE_PLUGIN_LIBS IluvatarTleIR IluvatarTleTransforms) | ||
| set(_ILUVATAR_TLE_PLUGIN_DEPS IluvatarTleTableGen IluvatarTleTransformsIncGen) | ||
| else() | ||
| set(_ILUVATAR_TLE_PLUGIN_SOURCES "") | ||
| set(_ILUVATAR_TLE_PLUGIN_LIBS "") | ||
| set(_ILUVATAR_TLE_PLUGIN_DEPS "") | ||
| endif() | ||
| add_triton_plugin(TritonILUVATAR | ||
| ${CMAKE_CURRENT_SOURCE_DIR}/triton_iluvatar.cc | ||
| ${_ILUVATAR_TLE_PLUGIN_SOURCES} | ||
| LINK_LIBS TritonILUVATARGPUToLLVM ${_ILUVATAR_TLE_PLUGIN_LIBS}) | ||
| if(FLAGTREE_ILUVATAR_TLE) | ||
| add_dependencies(TritonILUVATAR ${_ILUVATAR_TLE_PLUGIN_DEPS}) | ||
| endif() | ||
| target_link_libraries(TritonILUVATAR PRIVATE Python3::Module pybind11::headers lldCommon lldELF) | ||
| endif() | ||
| if(TRITON_BUILD_UT) | ||
| add_subdirectory(unittest) | ||
| endif() | ||
| # add_subdirectory(test) |
Empty file.
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
elseif(FLAGTREE_BACKEND STREQUAL "iluvatar")