setup_helper.py

import os
import shutil
import sys
import sysconfig
import functools
from pathlib import Path
import hashlib
from distutils.sysconfig import get_python_lib
from . import utils
import importlib.util
import importlib.metadata
from typing import List, Tuple
from .utils.tools import flagtree_configs as configs

downloader = utils.tools.DownloadManager()
configs = configs
flagtree_backend = configs.flagtree_backend


def _patch_llvm_exports():
    """Remove the cmake import-file existence check from LLVMExports.cmake.

    Some trust/xtdk-llvm22 packages export targets whose build-only binaries
    (llvm-tblgen/opt/llvm-link/...) are not shipped, which makes the cmake
    "verify imported files exist" loop raise FATAL_ERROR during
    find_package(MLIR). This patch disables that check.

    Notes:
    - Idempotent: a `[patched]` marker is written in place of the check block,
      so re-running on an already-patched file is a no-op.
    - Harmless when the check block is absent (e.g. the 20260615 package does
      not register those binaries for checking, so nothing matches).
    """
    import re
    import glob
    syspath = os.environ.get('LLVM_SYSPATH', '')
    if not syspath:
        return
    marker = "# [patched] import file check disabled - binaries not shipped in trust package"
    for f in glob.glob(f"{syspath}/lib/cmake/llvm/LLVMExports*.cmake"):
        with open(f) as fh:
            content = fh.read()
        if marker in content:
            continue  # already patched; stay idempotent
        # Non-greedy match of a single check block. The block always ends at
        # the first `unset(_cmake_import_check_targets)`, so `.*?` cannot span
        # into a following block.
        patched = re.sub(r'# Loop over all imported files.*?unset\(_cmake_import_check_targets\)', marker, content,
                         flags=re.DOTALL)
        if patched != content:
            with open(f, 'w') as fh:
                fh.write(patched)
            print(f"[XPU] patched LLVMExports: {f}")


set_llvm_env = lambda path: set_env(
    {
        'LLVM_INCLUDE_DIRS': Path(path) / "include",
        'LLVM_LIBRARY_DIR': Path(path) / "lib",
        'LLVM_SYSPATH': path,
        'PYTHONPATH': os.pathsep.join([str(Path(path) / "python_packages" / "mlir_core"),
                                       os.getenv("PYTHONPATH", "")]),
    })


def install_extension(*args, **kargs):
    backend_spec_install_extension_fn = get_hook_instance("install_extension")
    if backend_spec_install_extension_fn:
        backend_spec_install_extension_fn(*args, **kargs)


def get_backend_cmake_args(*args, **kargs):
    if "editable_wheel" in sys.argv:
        editable = True
    else:
        editable = False
    handle_plugin_backend(editable)
    try:
        cmake_args = configs.activated_module.get_backend_cmake_args(*args, **kargs)
    except Exception:
        cmake_args = []
    if editable:
        cmake_args += ["-DEDITABLE_MODE=ON"]
    return cmake_args


def get_device_name():
    return configs.device_alias_map[flagtree_backend]


def get_extra_packages():
    packages = []
    try:
        packages = configs.activated_module.get_extra_install_packages()
    except Exception:
        packages = []
    return packages


def get_package_data_tools():
    package_data = ["compile.h", "compile.c"]
    try:
        package_data += configs.activated_module.get_package_data_tools()
    except Exception:
        package_data
    return package_data


def dir_rollback(deep, base_path):
    while (deep):
        base_path = os.path.dirname(base_path)
        deep -= 1
    return Path(base_path)


def get_hook_instance(hook_name):
    if not configs.activated_module or not hook_name:
        return None
    hook_instance = getattr(configs.activated_module, hook_name, None)
    return hook_instance if callable(hook_instance) else None


def enable_flagtree_third_party(name):
    if name in ["triton_shared", "flagcx"]:
        return os.environ.get(f"USE_{name.upper()}", 'OFF') == 'ON'
    else:
        return os.environ.get(f"USE_{name.upper()}", 'ON') == 'ON'


def download_flagtree_third_party(name, condition, required=False, hook=None):
    if condition:
        if enable_flagtree_third_party(name):
            submodule = utils.flagtree_submodules[name]
            downloader.download(module=submodule, required=required)
            hook_call = get_hook_instance(hook)
            if hook_call:
                hook_call(configs=configs, backend=submodule, cache=cache)

        else:
            print(f"\033[1;33m[Note] Skip downloading {name} since USE_{name.upper()} is set to OFF\033[0m")


def post_install():
    backend_spec_post_install_fn = get_hook_instance("post_install")
    if backend_spec_post_install_fn:
        backend_spec_post_install_fn()


def write_flagtree_backend_file(triton_pkg_dir=None):
    if triton_pkg_dir is None:
        triton_pkg_dir = Path(__file__).resolve().parents[1] / "triton"
    backend_value = os.environ.get("FLAGTREE_BACKEND", "")
    dest_file = Path(triton_pkg_dir) / "FLAGTREE_BACKEND"
    dest_file.write_text(backend_value)


class FlagTreeCache:

    def __init__(self):
        self.flagtree_dir = str(Path(__file__).resolve().parents[2])
        self.dir_name = ".flagtree"
        self.sub_dirs = {}
        self.cache_files = {}
        self.dir_path = self._get_cache_dir_path()
        self._create_cache_dir()
        if flagtree_backend:
            self._create_subdir(subdir_name=flagtree_backend)

    @functools.lru_cache(maxsize=None)
    def _get_cache_dir_path(self) -> Path:
        _cache_dir = os.environ.get("FLAGTREE_CACHE_DIR")
        if _cache_dir is None:
            _cache_dir = Path.home() / self.dir_name
        else:
            _cache_dir = Path(_cache_dir)
        return _cache_dir

    def _create_cache_dir(self) -> Path:
        if not os.path.exists(self.dir_path):
            os.makedirs(self.dir_path, exist_ok=True)

    def _create_subdir(self, subdir_name, path=None):
        if path is None:
            subdir_path = Path(self.dir_path) / subdir_name
        else:
            subdir_path = Path(path) / subdir_name

        if not os.path.exists(subdir_path):
            os.makedirs(subdir_path, exist_ok=True)
        self.sub_dirs[subdir_name] = subdir_path

    def _md5(self, file_path):
        md5_hash = hashlib.md5()
        with open(file_path, "rb") as file:
            while chunk := file.read(4096):
                md5_hash.update(chunk)
        return md5_hash.hexdigest()

    def check_file(self, file_name=None, url=None, path=None, md5_digest=None):
        origin_file_path = None
        if url is not None:
            origin_file_name = url.split("/")[-1].split('.')[0]
            origin_file_path = self.cache_files.get(origin_file_name, "")
        if path is not None:
            _path = path
        else:
            _path = self.cache_files.get(file_name, "")
        empty = (not os.path.exists(_path)) or (origin_file_path and not os.path.exists(origin_file_path))
        if empty:
            return False
        if md5_digest is None:
            return True
        else:
            cur_md5 = self._md5(_path)
            return cur_md5[:8] == md5_digest

    def clear(self):
        shutil.rmtree(self.dir_path)

    def reverse_copy(self, src_path, cache_file_path, md5_digest):
        if src_path is None or not os.path.exists(src_path):
            return False
        if os.path.exists(cache_file_path):
            return False
        copy_needed = True
        if md5_digest is None or self._md5(src_path) == md5_digest:
            copy_needed = False
        if copy_needed:
            print(f"copying {src_path} to {cache_file_path}")
            if os.path.isdir(src_path):
                shutil.copytree(src_path, cache_file_path, dirs_exist_ok=True)
            else:
                shutil.copy(src_path, cache_file_path)
            return True
        return False

    def store(self, file=None, condition=None, url=None, copy_src_path=None, copy_dst_path=None, files=None,
              md5_digest=None, pre_hook=None, post_hook=None, version=None):

        if not condition or (pre_hook and pre_hook()):
            return
        is_url = False if url is None else True
        path = self.sub_dirs[flagtree_backend] if flagtree_backend else self.dir_path

        if files is not None:
            for single_files in files:
                self.cache_files[single_files] = Path(path) / single_files
        else:
            self.cache_files[file] = Path(path) / file
            if url is not None:
                origin_file_name = url.split("/")[-1].split('.')[0]
                self.cache_files[origin_file_name] = Path(path) / file
            if copy_dst_path is not None:
                dst_path_root = Path(self.flagtree_dir) / copy_dst_path
                dst_path = Path(dst_path_root) / file
                if self.reverse_copy(dst_path, self.cache_files[file], md5_digest):
                    return

        if is_url:
            cache_path = self.cache_files[file]
            need_download = not self.check_file(file_name=file, url=url, md5_digest=md5_digest)
            # Version check: re-download if cached version doesn't match expected
            if not need_download and version is not None:
                version_file = Path(cache_path) / "version.txt"
                if version_file.exists():
                    cached_ver = version_file.read_text().strip()
                    if cached_ver != version:
                        print(
                            f"[cache] version mismatch for '{file}': cached='{cached_ver}', expected='{version}', re-downloading..."
                        )
                        shutil.rmtree(cache_path)
                        need_download = True
                # If no version.txt (legacy cache), keep using it
            if need_download:
                downloader.download(url=url, path=path, file_name=file)
                if version is not None:
                    cache_path = self.cache_files[file]
                    version_file = Path(cache_path) / "version.txt"
                    if os.path.isdir(cache_path):
                        version_file.write_text(version)

        if copy_dst_path is not None:
            file_lists = [file] if files is None else list(files)
            for single_file in file_lists:
                dst_path_root = Path(self.flagtree_dir) / copy_dst_path
                os.makedirs(dst_path_root, exist_ok=True)
                dst_path = Path(dst_path_root) / single_file
                if not self.check_file(path=dst_path, md5_digest=md5_digest):
                    if copy_src_path:
                        src_path = Path(copy_src_path) / single_file
                    else:
                        src_path = self.cache_files[single_file]
                    print(f"copying {src_path} to {dst_path}")
                    if os.path.isdir(src_path):
                        shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
                    else:
                        shutil.copy(src_path, dst_path)
        post_hook(self.cache_files[file]) if post_hook else False

    def get(self, file_name) -> Path:
        return self.cache_files[file_name]


cache = FlagTreeCache()

# -----flagtree-tle-raw-----flagtree-mlir---


class LLVMDetector:
    ENV_VARS = [
        "LLVM_INCLUDE_DIRS",
        "LLVM_LIBRARY_DIR",
        "LLVM_SYSPATH",
    ]

    @classmethod
    def has_env_vars(cls) -> List[str]:
        return [k for k in cls.ENV_VARS if k in os.environ]

    @staticmethod
    def is_wheel_installed(pkg_name: str) -> bool:
        try:
            importlib.metadata.version(pkg_name)
            return True
        except importlib.metadata.PackageNotFoundError:
            return False

    @staticmethod
    def get_paths_from_wheel(pkg_name: str) -> Tuple[str, str, str]:
        spec = importlib.util.find_spec(pkg_name)
        if spec is None:
            raise RuntimeError(f"LLVM wheel '{pkg_name}' found via metadata but import failed.")

        if spec.origin:
            pkg_root = os.path.dirname(spec.origin)
        elif spec.submodule_search_locations:
            pkg_root = spec.submodule_search_locations[0]
        else:
            raise RuntimeError(f"LLVM wheel '{pkg_name}' is found but has no filesystem location")

        # New wheel structure: LLVM artifacts are under mlir/llvm_artifact/
        llvm_artifact_dir = os.path.join(pkg_root, "llvm_artifact")
        if os.path.isdir(llvm_artifact_dir):
            llvm_root = llvm_artifact_dir
        else:
            # Fallback: legacy structure where artifacts are directly under the package root
            llvm_root = pkg_root

        include_dir = os.path.join(llvm_root, "include")
        lib_dir = os.path.join(llvm_root, "lib")
        return include_dir, lib_dir, llvm_root


def try_setup_flagtree_mlir(pkg_name: str = "mlir") -> bool:
    is_installed = LLVMDetector.is_wheel_installed(pkg_name)
    has_envs = LLVMDetector.has_env_vars()
    # rule1 : if both exist, fail
    if is_installed and has_envs and not os.environ.get("USE_FLAGTREE_MLIR_BUILD"):
        raise RuntimeError("ERROR: LLVM wheel is installed, but LLVM-related environment variables are set:\n"
                           f"  {has_envs}\n"
                           "Please unset them to avoid conflicts.")

    # rule2：wheel installed & no env → use wheel
    if is_installed:
        include_dir, lib_dir, llvm_root = LLVMDetector.get_paths_from_wheel(pkg_name)
        # env variables will not appear out of python process
        os.environ["USE_FLAGTREE_MLIR_BUILD"] = "1"
        os.environ["LLVM_SYSPATH"] = llvm_root
        os.environ["LLVM_INCLUDE_DIRS"] = include_dir
        os.environ["LLVM_LIBRARY_DIR"] = lib_dir
        return True

    # Rule 3: fallback to legacy
    return False


# --------------------------


class CommonUtils:

    @staticmethod
    def unlink():
        cur_path = dir_rollback(2, __file__)
        if "editable_wheel" in sys.argv:
            installation_dir = cur_path
        else:
            installation_dir = get_python_lib()
        backends_dir_path = Path(installation_dir) / "triton" / "backends"
        # raise RuntimeError(backends_dir_path)
        if not os.path.exists(backends_dir_path):
            return
        for name in os.listdir(backends_dir_path):
            exist_backend_path = os.path.join(backends_dir_path, name)
            if not os.path.isdir(exist_backend_path):
                continue
            if name.startswith('__'):
                continue
            if os.path.islink(exist_backend_path):
                os.unlink(exist_backend_path)
            if os.path.exists(exist_backend_path):
                shutil.rmtree(exist_backend_path)

    @staticmethod
    def skip_package_dir(package):
        if 'backends' in package or 'profiler' in package:
            return True
        try:
            return configs.activated_module.skip_package_dir(package)
        except Exception:
            return False

    @staticmethod
    def get_package_dir(packages):
        package_dict = {}
        if flagtree_backend and flagtree_backend not in configs.plugin_backends:
            connection = []
            backend_triton_path = f"./third_party/{flagtree_backend}/python/"
            for package in packages:
                if CommonUtils.skip_package_dir(package):
                    continue
                pair = (package, f"{backend_triton_path}{package}")
                connection.append(pair)
            package_dict.update(connection)
        try:
            package_dict.update(configs.activated_module.get_package_dir())
        except Exception:
            pass
        return package_dict


def handle_flagtree_backend():
    global ext_sourcedir
    if flagtree_backend:
        print(f"\033[1;32m[INFO] FlagtreeBackend is {flagtree_backend}\033[0m")
        configs.extend_backends.append(flagtree_backend)
        if "editable_wheel" in sys.argv and flagtree_backend not in configs.plugin_backends:
            ext_sourcedir = os.path.abspath(f"./third_party/{flagtree_backend}/python/{configs.ext_sourcedir}") + "/"


def handle_plugin_backend(editable):
    plugin_mode = os.getenv("FLAGTREE_PLUGIN")
    if (plugin_mode and plugin_mode.upper() not in ["0", "OFF"]) or not flagtree_backend:
        return
    flagtree_backend_dir = cache.sub_dirs[flagtree_backend]
    flagtree_plugin_so = flagtree_backend + "TritonPlugin.so"
    src_build_plugin_path = flagtree_backend_dir / flagtree_plugin_so
    if not src_build_plugin_path.exists():
        return
    if flagtree_backend in ["iluvatar", "mthreads", "sunrise"]:
        if editable is False:
            dst_build_plugin_dir = Path(sysconfig.get_path("purelib")) / "triton" / "_C"
            if not os.path.exists(dst_build_plugin_dir):
                os.makedirs(dst_build_plugin_dir)
            dst_build_plugin_path = dst_build_plugin_dir / flagtree_plugin_so
            shutil.copy(src_build_plugin_path, dst_build_plugin_path)
        if flagtree_backend in ("mthreads", ):
            dst_install_plugin_dir = Path(
                __file__).resolve().parent.parent.parent / "third_party" / flagtree_backend / "python" / "triton" / "_C"
        else:
            dst_install_plugin_dir = Path(__file__).resolve().parent.parent / "triton" / "_C"
        if not os.path.exists(dst_install_plugin_dir):
            os.makedirs(dst_install_plugin_dir)
        shutil.copy(src_build_plugin_path, dst_install_plugin_dir)


def set_env(env_dict: dict):
    for env_k, env_v in env_dict.items():
        os.environ[env_k] = str(env_v)


def check_env(env_val):
    return os.environ.get(env_val, '') != ''


def uninstall_triton():
    is_bdist_wheel = any(cmd in sys.argv for cmd in ['bdist_wheel', 'egg_info', 'sdist'])
    if is_bdist_wheel:
        return
    try:
        import pkg_resources
        import subprocess
        try:
            pkg_resources.get_distribution('triton')
            print("Detected existing 'triton' package. Uninstalling to avoid conflicts...")
            subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", "triton"])
            print("Successfully uninstalled 'triton'.")
        except pkg_resources.DistributionNotFound:
            print("'triton' package not found, no need to uninstall.")
    except Exception as e:
        print(f"Warning: Failed to check/uninstall triton: {e}")


offline_handler = utils.OfflineBuildManager()
if offline_handler.is_offline:
    print("[INFO] FlagTree Offline Build: Use offline build for triton origin toolkits")
    offline_handler.handle_triton_origin_toolkits()
    offline_build = True
else:
    print('[INFO] FlagTree Offline Build: No offline build for triton origin toolkits')
    offline_build = False

cache = FlagTreeCache()
'''
   FlagCX is a third-party library adopted by the tle distributed system,
   refer to https://github.com/flagos-ai/FlagCX
'''

download_flagtree_third_party("flagcx", condition=(not flagtree_backend), hook="handle_flagcx", required=True)

handle_flagtree_backend()

# iluvatar
cache.store(
    file="iluvatar-llvm22-x86_64",
    condition=("iluvatar" == flagtree_backend),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/iluvatar-llvm22-x86_64_v0.6.0.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

cache.store(
    file="iluvatarTritonPlugin.so", condition=("iluvatar" == flagtree_backend) and (not configs.flagtree_plugin), url=
    "https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/iluvatarTritonPlugin-cpython3.10-glibc2.30-glibcxx3.4.28-cxxabi1.3.12-ubuntu-x86_64_v0.3.0.tar.gz",
    copy_dst_path=f"third_party/{flagtree_backend}", md5_digest="015b9af8")

# klx xpu - upgraded to trust LLVM (llvm22) for Triton 3.6
# 20260615 xtdk-llvm22 ships the full xpu3 device-codegen toolchain
# (clang/ld.lld/xpu3-elfconv/xpu3-elfconv-triton/xpu-kernel.t/llvm-readelf/objdump/objcopy),
# so kernel launch works. It omits build-only tools (llvm-tblgen/opt/llvm-link/llvm-as/
# llvm-dis) that LLVMExports.cmake references, so _patch_llvm_exports() disables that check.
# decompress() renames the extracted top-level dir to the `file=` name, so LLVM_SYSPATH
# resolves to <cache>/xpu/llvm_trust regardless of the tarball's top-level name.
cache.store(
    file="llvm_trust",
    condition=("xpu" == flagtree_backend),
    url="https://klx-sdk-release-public.su.bcebos.com/XTriton/llvm22/20260615/xtdk-llvm22-ubuntu2004_x86_64.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=lambda path: (set_llvm_env(path), _patch_llvm_exports()),
    version="20260615",
)

cache.store(file="xre-Linux-x86_64", condition=("xpu" == flagtree_backend),
            url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/xre-Linux-x86_64_v0.3.0.tar.gz",
            copy_dst_path='python/_deps/xre3', version="v0.3.0")

# xpu device libs (liblaunch_shared.so, libxpujitc.so, LLVM-15, clang-cpp)
cache.store(file="xpu-device-libs", condition=("xpu" == flagtree_backend),
            url="https://klx-sdk-release-public.su.bcebos.com/XTriton/xpu-device-libs-ubuntu-x64_v0.3.6.1.1.tar.gz",
            version="v0.3.6.1.1")

cache.store(files=("liblaunch_shared.so", "libLLVM-15.so", "libclang-cpp.so.15", "libxpujitc.so"),
            condition=("xpu" == flagtree_backend), copy_src_path=f"{cache.dir_path}/{flagtree_backend}/xpu-device-libs",
            copy_dst_path=f"third_party/{flagtree_backend}/device")


# xpu SDNN prebuilt objects + libTritonSharedForXPU.a (344MB compressed)
def _install_xpu_sdnn_objects(cached_path):
    """Copy prebuilt SDNN objects from cache to third_party/xpu/."""
    flagtree_dir = str(Path(__file__).resolve().parents[2])
    dst_root = os.path.join(flagtree_dir, "third_party", "xpu")
    for item in os.listdir(cached_path):
        src = os.path.join(str(cached_path), item)
        dst = os.path.join(dst_root, item)
        if os.path.isdir(src):
            shutil.copytree(src, dst, dirs_exist_ok=True)
        else:
            shutil.copy(src, dst)
    print(f"[XPU] SDNN prebuilt objects installed to {dst_root}")


cache.store(file="xpu-sdnn-objects", condition=("xpu" == flagtree_backend),
            url="https://klx-sdk-release-public.su.bcebos.com/XTriton/xpu-sdnn-objects_v0.3.6.1.1.tar.gz",
            post_hook=_install_xpu_sdnn_objects)

# xpu3 bin tools from LLVM (xtdk-llvm22 20260615 ships xpu3-elfconv-triton directly)
cache.store(
    files=("clang", "xpu-xxd", "xpu3-elfconv", "xpu3-elfconv-triton", "xpu-kernel.t", "ld.lld", "llvm-readelf",
           "llvm-objdump", "llvm-objcopy"), condition=("xpu" == flagtree_backend),
    copy_src_path=f"{os.environ.get('LLVM_SYSPATH','')}/bin", copy_dst_path="third_party/xpu/backend/xpu3/bin")


# xpu3-elfconv-triton resolves llvm-readelf/objdump/objcopy from xpu3/ (parent of bin/),
# see the elfconv script line ~110. Symlink them so the resolver finds them.
def _link_elfconv_triton():
    import os
    xpu3_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "third_party",
                            "xpu", "backend", "xpu3")
    bin_dir = os.path.join(xpu3_dir, "bin")
    for tool in ("llvm-readelf", "llvm-objdump", "llvm-objcopy"):
        tool_src = os.path.join(bin_dir, tool)
        tool_dst = os.path.join(xpu3_dir, tool)
        if os.path.exists(tool_src) and not os.path.exists(tool_dst):
            os.symlink(os.path.join("bin", tool), tool_dst)
            print(f"Created symlink: {tool_dst} -> bin/{tool}")


if "xpu" == flagtree_backend:
    _link_elfconv_triton()

# xpu3 builtins + new crt*.o / xpuprintf (trust LLVM additions)
cache.store(
    files=("libclang_rt.builtins-xpu3.a", "libclang_rt.builtins-xpu3s.a", "clang_rt.crtbegin-xpu3.o",
           "clang_rt.crtend-xpu3.o", "libclang_rt.xpuprintf-xpu3.a", "libclang_rt.xpuprintfs-xpu3.a"),
    condition=("xpu" == flagtree_backend), copy_src_path=f"{os.environ.get('LLVM_SYSPATH','')}/lib/linux",
    copy_dst_path="third_party/xpu/backend/xpu3/lib/linux")

cache.store(files=("include", "so"), condition=("xpu" == flagtree_backend),
            copy_src_path=f"{cache.dir_path}/xpu/xre-Linux-x86_64", copy_dst_path="third_party/xpu/backend/xpu3")

# mthreads
cache.store(
    file="mthreads-llvm22",
    condition=("mthreads" == flagtree_backend),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/mthreads-llvm22-x64_v0.5.1.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

cache.store(file="mthreads_local_binary", condition=("mthreads" == flagtree_backend),
            url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/mthreads_local_binary_v0.6.0.tar.gz")

cache.store(files=("ld.lld", "llc"), condition=("mthreads" == flagtree_backend),
            copy_src_path=f"{cache.dir_path}/{flagtree_backend}/mthreads_local_binary",
            copy_dst_path=f"third_party/{flagtree_backend}/bin")

# ascend
cache.store(
    file="llvm-b5cc222d-ubuntu-arm64",
    condition=("ascend" == flagtree_backend),
    url="https://oaitriton.blob.core.windows.net/public/llvm-builds/llvm-b5cc222d-ubuntu-arm64.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

# aipu
cache.store(
    file="llvm-a66376b0-ubuntu-x64-clang16-lld16",
    condition=("aipu" == flagtree_backend),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-a66376b0-ubuntu-x64-clang16-lld16_v0.4.0.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

# enflame
cache.store(
    file="llvm-fc83c68-gcc9-x64",
    condition=("enflame" == flagtree_backend),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/enflame-llvm23-fc83c68-gcc9-x64_v0.4.0.tar.gz",
    pre_hook=lambda: check_env('KURAMA_LLVM_DIR'),
    post_hook=lambda path: set_env({
        'KURAMA_LLVM_DIR': path,
        'LLVM_INCLUDE_DIRS': Path(path) / "include",
        'LLVM_LIBRARY_DIR': Path(path) / "lib",
        'LLVM_SYSPATH': path,
    }),
)

# tsingmicro
cache.store(
    file="tsingmicro-llvm21-glibc2.30-glibcxx3.4.28-python3.11-x64",
    condition=("tsingmicro" == flagtree_backend),
    url=
    "https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/tsingmicro-llvm21-glibc2.30-glibcxx3.4.28-python3.11-x64_v0.2.0.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

cache.store(
    file="tx8_deps",
    condition=("tsingmicro" == flagtree_backend),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/tx8_depends_release_20250814_195126_v0.2.0.tar.gz",
    pre_hook=lambda: check_env('TX8_DEPS_ROOT'),
    post_hook=lambda path: set_env({
        'LLVM_SYSPATH': path,
    }),
)

# hcu
cache.store(
    file="hcu-llvm22-b0ca808-glibc2.35-glibcxx3.4.30-ubuntu-x86_64",
    condition=("hcu" == flagtree_backend),
    url=
    "https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/hcu-llvm22-b0ca808-glibc2.35-glibcxx3.4.30-ubuntu-x86_64_v0.5.0.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

# metax
cache.store(
    file="metax-llvm19",
    condition=("metax" == flagtree_backend),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/metax-llvm19-3.8.0.6-x86_64_v0.6.0.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

cache.store(
    file="metaxTritonPlugin.so",
    condition=("metax" == flagtree_backend) and (not configs.flagtree_plugin),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/metaxTritonPlugin-cpython3.12-x86_64_v0.6.0.tar.gz",
    copy_dst_path=f"third_party/{flagtree_backend}",
    md5_digest="415a08bd",
)

# thrive
cache.store(
    file="llvm-f6ded0be-ubuntu-x64",
    condition=("thrive" == flagtree_backend),
    url="https://oaitriton.blob.core.windows.net/public/llvm-builds/llvm-f6ded0be-ubuntu-x64.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=set_llvm_env,
)

cache.store(
    file="sunrise_llvm22_dev_release",
    condition=("sunrise" == flagtree_backend),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-34b694004c-triton-v3.6.x.tar.gz",
    pre_hook=lambda: check_env('LLVM_SYSPATH'),
    post_hook=lambda path: [f(path) for f in (set_llvm_env, utils.activate("sunrise").sunrise_cp_bc_files)],
)

cache.store(
    file="sunriseTritonPlugin.so",
    condition=("sunrise" == flagtree_backend) and (not configs.flagtree_plugin),
    url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/sunrise-plugin-triton-v3.6.x.tar.gz",
    md5_digest="3526d699",
)