diff --git a/.github/workflows/ascend-build-and-test.yml b/.github/workflows/ascend-build-and-test.yml index f491930d1..30bf975ab 100644 --- a/.github/workflows/ascend-build-and-test.yml +++ b/.github/workflows/ascend-build-and-test.yml @@ -21,6 +21,7 @@ jobs: shell: bash run: | export FLAGTREE_BACKEND=ascend + export NPUIR=af5499b3b9f3dbab50b2834bcfff5da5c2a1d920 source ~/env.sh cd python MAX_JOBS=32 python3 -m pip install . --no-build-isolation diff --git a/python/setup_tools/utils/__init__.py b/python/setup_tools/utils/__init__.py index 7d3487c96..53b31e1be 100644 --- a/python/setup_tools/utils/__init__.py +++ b/python/setup_tools/utils/__init__.py @@ -11,8 +11,8 @@ dst_path=os.path.join(flagtree_submodule_dir, "triton_shared")), "ascend": tools.Module(name="ascend", url="https://gitcode.com/FlagTree/triton-ascend.git", - dst_path=os.path.join(flagtree_submodule_dir, - "triton_ascend"), spec_submodule="third_party/ascendnpu-ir"), + dst_path=os.path.join(flagtree_submodule_dir, "triton_ascend"), + spec_submodule="third_party/ascendnpu-ir" if ascend.ascend_npuir_commit is None else None), } diff --git a/python/setup_tools/utils/ascend.py b/python/setup_tools/utils/ascend.py index cbc790458..eb4ac14b9 100644 --- a/python/setup_tools/utils/ascend.py +++ b/python/setup_tools/utils/ascend.py @@ -1,10 +1,16 @@ import os import shutil from pathlib import Path -from setup_tools.utils.tools import flagtree_root_dir, flagtree_submodule_dir, DownloadManager +from setup_tools.utils.tools import flagtree_root_dir, flagtree_submodule_dir, DownloadManager, Module downloader = DownloadManager() +ascend_npuir_commit = os.environ.get("NPUIR", "") + +submodules = (Module(name="ascendnpu-ir", url="https://gitcode.com/Ascend/AscendNPU-IR.git", + commit_id=ascend_npuir_commit, dst_path=os.path.join(flagtree_submodule_dir, + "ascend/third_party/ascendnpu-ir")), ) + def get_backend_cmake_args(*args, **kargs): build_ext = kargs['build_ext'] @@ -107,8 +113,6 @@ def precompile_hock(*args, **kargs): ascend_path = Path(third_party_base_dir) / "ascend" patch_path = Path(ascend_path) / "triton_patch" project_path = Path(third_party_base_dir) / "triton_ascend" - project_thirdparty_path = project_path / "third_party/ascendnpu-ir" - ascend_thirdparty_path = ascend_path / "third_party/ascendnpu-ir" if os.path.exists(ascend_path): shutil.rmtree(ascend_path) if not os.path.exists(project_path): @@ -117,8 +121,13 @@ def precompile_hock(*args, **kargs): patch_src_path = Path(project_path) / "triton_patch" shutil.copytree(ascend_src_path, ascend_path, dirs_exist_ok=True) shutil.copytree(patch_src_path, patch_path, dirs_exist_ok=True) - shutil.copytree(project_thirdparty_path, ascend_thirdparty_path, dirs_exist_ok=True) + if not ascend_npuir_commit: + project_thirdparty_path = project_path / "third_party/ascendnpu-ir" + ascend_thirdparty_path = ascend_path / "third_party/ascendnpu-ir" + shutil.copytree(project_thirdparty_path, ascend_thirdparty_path, dirs_exist_ok=True) shutil.rmtree(project_path) + if ascend_npuir_commit: + [downloader.download(module=submodule, required=False) for submodule in submodules] cmake_patch_copy() patched_code = """ set(triton_abs_dir "${TRITON_ROOT_DIR}/include/triton/Dialect/Triton/IR") """ src_code = """set(triton_abs_dir""" diff --git a/python/setup_tools/utils/tools.py b/python/setup_tools/utils/tools.py index 9f5407842..28e3f8f55 100644 --- a/python/setup_tools/utils/tools.py +++ b/python/setup_tools/utils/tools.py @@ -150,7 +150,7 @@ def git_clone(self, module, required=False): if not os.path.exists(module.dst_path): succ = self.clone_module(module) else: - print(f'Found third_party {module.name} at {module.dst_path}\n') + print(f'[INFO] : Found third_party {module.name} at {module.dst_path}\n') return True if not succ and required: raise RuntimeError(