venv/Lib/site-packages/triton/runtime/build.py with benefits

build.py

# venv/Lib/site-packages/triton/runtime/build.py
import sysconfig
import os
import shutil
import subprocess
import sys

from .. import knobs

def _build(name: str, src: str, srcdir: str, library_dirs: list[str], include_dirs: list[str],
           libraries: list[str]) -> str:
    if impl := knobs.build.impl:
        return impl(name, src, srcdir, library_dirs, include_dirs, libraries)

    suffix = sysconfig.get_config_var('EXT_SUFFIX')
    so = os.path.join(srcdir, '{name}{suffix}'.format(name=name, suffix=suffix))

    cc = os.environ.get("CC")
    if cc is None:
        rocm_path = os.environ.get("ROCM_PATH") or os.environ.get("HIP_PATH")
        clang = os.path.join(rocm_path, 'bin', 'clang.exe')
        if os.path.exists(clang):
            print("%%% Using HIP SDK Clang: {}".format(clang))
            cc = clang

    if cc is None:
        clang = shutil.which("clang")
        gcc = shutil.which("gcc")
        cc = gcc if gcc is not None else clang
        if cc is None:
            raise RuntimeError(
                "Failed to find C compiler. Please specify via CC environment variable or set triton.knobs.build.impl.")

    if hasattr(sysconfig, 'get_default_scheme'):
        scheme = sysconfig.get_default_scheme()
    else:
        scheme = sysconfig._get_default_scheme()  # type: ignore

    if scheme == 'posix_local':
        scheme = 'posix_prefix'

    # Painful debugging
    py_include_dir = sysconfig.get_paths(scheme=scheme)["include"]
    print("%%% Include path: {}\n".format(scheme, py_include_dir))

    base_prefix = sys.base_prefix  # Use this instead of sys.prefix to escape the venv
    libs_dir = os.path.join(base_prefix, "libs")
    lib_name = f"python{sys.version_info.major}{sys.version_info.minor}.lib"
    full_path = os.path.join(libs_dir, lib_name)

    print("%%% Expected location of python*.lib:", full_path)
    print("%%% Exists:", os.path.exists(full_path))
    # End debugging

    custom_backend_dirs = knobs.build.backend_dirs
    include_dirs = include_dirs + [srcdir, py_include_dir, *custom_backend_dirs]

    import site
    library_dirs += [os.path.join(site.getsitepackages()[0], 'libs')]
    library_dirs += [os.path.join(os.environ['HIP_PATH'], 'lib')]
    # Add library directory manually (probably not needed for build script, but needed when testing manually
    library_dirs += [libs_dir]

    # Added -Wno-error, though clang shouldn't fail on a warning anyway.
    cc_cmd = [cc, src, "-v", "-O3", "-shared", "-Wno-psabi", "-Wno-error", "-march=native", "-mtune=native", "-o", so]
    cc_cmd += [f'-l{lib}' for lib in libraries]
    cc_cmd += [f"-L{dir}" for dir in library_dirs]
    cc_cmd += [f"-I{dir}" for dir in include_dirs if dir is not None]

    ## -target-cpu x86-64 -tune-cpu generic
    ###
    # Make a copy of all the temporary files so we can run clang manually to test
    ###
    # Copy anything from temp dir into c:\tmp\amd and substitute in printed command
    temp_root = r"C:\Users\sfink\AppData\Local\Temp\\"
    tmp_dest_root = r"C:\tmp\amd\\"
    os.makedirs(tmp_dest_root, exist_ok=True)

    def map_path(path: str) -> str:
        rel = os.path.relpath(path, temp_root)
        return os.path.join(tmp_dest_root, rel)

    def handle_temp_arg(arg: str) -> str:
        norm_temp_root = os.path.normcase(os.path.normpath(temp_root))

        def is_under_temp(path: str) -> bool:
            norm_path = os.path.normcase(os.path.normpath(path))
            return norm_path.startswith(norm_temp_root)

        # Direct path
        if is_under_temp(arg):
            dst = map_path(arg)
            if os.path.exists(arg):
                os.makedirs(os.path.dirname(dst), exist_ok=True)
                if os.path.isdir(arg):
                    shutil.copytree(arg, dst, dirs_exist_ok=True)
                else:
                    shutil.copyfile(arg, dst)
            return dst  # return dst even if the file didn't exist

        # -I or -L prefixes
        for prefix in ("-I", "-L"):
            if arg.startswith(prefix):
                value = arg[len(prefix):]
                if is_under_temp(value):
                    dst = map_path(value)
                    if os.path.exists(value):
                        os.makedirs(os.path.dirname(dst), exist_ok=True)
                        if os.path.isdir(value):
                            shutil.copytree(value, dst, dirs_exist_ok=True)
                        else:
                            shutil.copyfile(value, dst)
                    return prefix + dst  # return substituted dst either way

        return arg

    cmd_for_print = []
    for arg in cc_cmd:
        fixed = handle_temp_arg(f"{arg}")
        quoted = f'"{fixed}"' if ' ' in fixed or fixed.endswith('.exe') else fixed
        cmd_for_print.append(quoted)

    print("\n%%% CMD.exe-compatible build command with temp file copies:")
    print(" ".join(cmd_for_print))
    print()
    ###
    # END DEBUG
    ###

    subprocess.check_call(cc_cmd, stdout=subprocess.DEVNULL)
    # This point never reached (signal):
    #     clang: error: linker command failed due to signal (use -v to see invocation)
    print("%%% subprocess.check_call finished\n")
    return so