[BACKEND] [BUILD] Enhance ascend submodule commit pull action

.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

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),
 }
 
 

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"""

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(