build_tools/bazel/amdgpu_device_toolchain_repo.bzl

# Copyright 2026 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

"""Repository rule exposing optional AMDGPU builtin device compiler tools."""

_STUB_BUILD = """\
# AUTO-GENERATED by amdgpu_device_toolchain_repo.bzl - do not edit.

package(default_visibility = ["//visibility:public"])

filegroup(
    name = "unused",
    srcs = [],
)
"""

_STUB_PATHS = """\
# AUTO-GENERATED by amdgpu_device_toolchain_repo.bzl - do not edit.

AMDGPU_DEVICE_TOOLCHAIN_AVAILABLE = False
AMDGPU_DEVICE_TOOLCHAIN_ERROR = "AMDGPU device binary source mode requires --repo_env=IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN=rocm, llvm-tools, llvm-project, or auto."
AMDGPU_CLANG_TOOL = ""
AMDGPU_LLVM_LINK_TOOL = ""
AMDGPU_LLD_TOOL = ""
AMDGPU_LLVM_OBJCOPY_TOOL = ""
AMDGPU_CLANG_RESOURCE_HEADERS = ""
AMDGPU_CLANG_RESOURCE_MARKER = ""
AMDGPU_CLANG_RESOURCE_INCLUDE_ARGS = []
"""

_PREBUILT_TOOL_BZL = """\
\"\"\"Rule for exposing pre-built binaries as executable Bazel targets.\"\"\"

def _prebuilt_tool_impl(ctx):
    out = ctx.actions.declare_file(ctx.label.name)
    ctx.actions.symlink(
        output = out,
        target_file = ctx.file.src,
        is_executable = True,
    )
    runfiles = ctx.runfiles(files = ctx.files.data + [ctx.file.src])
    return [DefaultInfo(
        executable = out,
        files = depset([out]),
        runfiles = runfiles,
    )]

prebuilt_tool = rule(
    implementation = _prebuilt_tool_impl,
    executable = True,
    attrs = {
        "src": attr.label(
            allow_single_file = True,
            mandatory = True,
            doc = "The pre-built binary file.",
        ),
        "data": attr.label_list(
            allow_files = True,
            doc = "Additional files the tool needs at runtime.",
        ),
    },
    doc = "Wraps a pre-built binary as an executable Bazel target.",
)
"""

_LOCAL_BUILD_TEMPLATE = """\
# AUTO-GENERATED by amdgpu_device_toolchain_repo.bzl - do not edit.

load(":prebuilt_tool.bzl", "prebuilt_tool")

package(default_visibility = ["//visibility:public"])

prebuilt_tool(
    name = "clang",
    src = "bin/clang",
)

prebuilt_tool(
    name = "llvm-link",
    src = "bin/llvm-link",
)

prebuilt_tool(
    name = "lld",
    src = "bin/lld",
)

prebuilt_tool(
    name = "llvm-objcopy",
    src = "bin/llvm-objcopy",
)

filegroup(
    name = "clang_resource_headers",
    srcs = glob(["clang_resource_include/**"]),
)

filegroup(
    name = "clang_resource_include_marker",
    srcs = ["clang_resource_include/stddef.h"],
)
"""

_LOCAL_PATHS_TEMPLATE = """\
# AUTO-GENERATED by amdgpu_device_toolchain_repo.bzl - do not edit.

AMDGPU_DEVICE_TOOLCHAIN_AVAILABLE = True
AMDGPU_DEVICE_TOOLCHAIN_ERROR = ""
AMDGPU_CLANG_TOOL = "@iree_amdgpu_device_toolchain//:clang"
AMDGPU_LLVM_LINK_TOOL = "@iree_amdgpu_device_toolchain//:llvm-link"
AMDGPU_LLD_TOOL = "@iree_amdgpu_device_toolchain//:lld"
AMDGPU_LLVM_OBJCOPY_TOOL = "@iree_amdgpu_device_toolchain//:llvm-objcopy"
AMDGPU_CLANG_RESOURCE_HEADERS = "@iree_amdgpu_device_toolchain//:clang_resource_headers"
AMDGPU_CLANG_RESOURCE_MARKER = "@iree_amdgpu_device_toolchain//:clang_resource_include_marker"
AMDGPU_CLANG_RESOURCE_INCLUDE_ARGS = [
    "--clang-resource-include",
    "$$(dirname $(location @iree_amdgpu_device_toolchain//:clang_resource_include_marker))",
]
"""

_LLVM_PROJECT_BUILD = """\
# AUTO-GENERATED by amdgpu_device_toolchain_repo.bzl - do not edit.

package(default_visibility = ["//visibility:public"])

alias(
    name = "clang",
    actual = "@llvm-project//clang:clang",
)

alias(
    name = "llvm-link",
    actual = "@llvm-project//llvm:llvm-link",
)

alias(
    name = "lld",
    actual = "@llvm-project//lld:lld",
)

alias(
    name = "llvm-objcopy",
    actual = "@llvm-project//llvm:llvm-objcopy",
)
"""

_LLVM_PROJECT_PATHS = """\
# AUTO-GENERATED by amdgpu_device_toolchain_repo.bzl - do not edit.

AMDGPU_DEVICE_TOOLCHAIN_AVAILABLE = True
AMDGPU_DEVICE_TOOLCHAIN_ERROR = ""
AMDGPU_CLANG_TOOL = "@iree_amdgpu_device_toolchain//:clang"
AMDGPU_LLVM_LINK_TOOL = "@iree_amdgpu_device_toolchain//:llvm-link"
AMDGPU_LLD_TOOL = "@iree_amdgpu_device_toolchain//:lld"
AMDGPU_LLVM_OBJCOPY_TOOL = "@iree_amdgpu_device_toolchain//:llvm-objcopy"
AMDGPU_CLANG_RESOURCE_HEADERS = "@llvm-project//clang:builtin_headers_gen"
AMDGPU_CLANG_RESOURCE_MARKER = ""
AMDGPU_CLANG_RESOURCE_INCLUDE_ARGS = [
    "--clang-resource-include",
    "$(BINDIR)/external/+_repo_rules+llvm-project/clang/staging/include/",
]
"""

def _split_env_list(value):
    if not value:
        return []
    normalized = value
    for separator in [",", ";", "\n", "\t", " "]:
        normalized = normalized.replace(separator, ":")
    return [part for part in normalized.split(":") if part]

def _join_path(left, right):
    if left.endswith("/"):
        return left + right
    return left + "/" + right

def _dirname(path):
    if "/" not in path:
        return "."
    return path.rsplit("/", 1)[0]

def _basename(path):
    if "/" not in path:
        return path
    return path.rsplit("/", 1)[1]

def _append_unique(values, new_values):
    for value in new_values:
        if value not in values:
            values.append(value)

def _is_executable(repository_ctx, path):
    result = repository_ctx.execute(["test", "-x", path], quiet = True)
    return result.return_code == 0

def _find_tool(
        repository_ctx,
        names,
        directories,
        explicit_env_vars,
        description,
        additional_candidates = ()):
    candidates = []
    for env_var in explicit_env_vars:
        value = repository_ctx.getenv(env_var)
        if value:
            candidates.append(value)
    candidates.extend(additional_candidates)
    for directory in directories:
        for name in names:
            candidates.append(_join_path(directory, name))

    for candidate in candidates:
        if _is_executable(repository_ctx, candidate):
            return str(repository_ctx.path(candidate))

    fail("Could not find {}. Tried env vars [{}] and candidate directories [{}].".format(
        description,
        ", ".join(explicit_env_vars),
        ", ".join(directories),
    ))

def _rocm_tool_dirs(root):
    return [
        _join_path(root, "llvm/bin"),
        _join_path(root, "lib/llvm/bin"),
        _join_path(root, "bin"),
    ]

def _candidate_rocm_roots(repository_ctx):
    roots = []
    for env_var in [
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_ROCM_PATH",
        "IREE_ROCM_PATH",
        "ROCM_PATH",
        "ROCM_ROOT",
        "ROCM_HOME",
        "HIP_PATH",
    ]:
        for value in _split_env_list(repository_ctx.getenv(env_var)):
            _append_unique(roots, [value])
            if env_var == "HIP_PATH" and _basename(value) == "hip":
                _append_unique(roots, [_dirname(value)])

    hipcc = repository_ctx.which("hipcc")
    if hipcc:
        _append_unique(roots, [_dirname(_dirname(str(hipcc)))])

    if repository_ctx.path("/opt/rocm").exists:
        _append_unique(roots, ["/opt/rocm"])
    return roots

def _candidate_tool_dirs(repository_ctx, mode):
    dirs = []
    for env_var in [
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLVM_TOOLS_DIR",
        "IREE_HOST_BIN_DIR",
        "IREE_HOST_TOOLS",
        "IREE_HOST_TOOLS_DIR",
        "IREE_BINARY_DIR",
        "IREE_LLVM_TOOLS_DIR",
        "IREE_LLVM_TOOL_DIR",
        "LLVM_TOOLS_BINARY_DIR",
        "LLVM_BINARY_DIR",
    ]:
        for value in _split_env_list(repository_ctx.getenv(env_var)):
            _append_unique(dirs, [value])

    expanded_dirs = []
    for directory in dirs:
        expanded_dirs.extend([
            directory,
            _join_path(directory, "bin"),
            _join_path(directory, "llvm-project/bin"),
            _join_path(directory, "llvm/bin"),
            _join_path(directory, "lib/llvm/bin"),
        ])
    dirs = []
    _append_unique(dirs, expanded_dirs)

    if mode in ["auto", "rocm"]:
        for root in _candidate_rocm_roots(repository_ctx):
            _append_unique(dirs, _rocm_tool_dirs(root))
    return dirs

def _clang_tool_candidates(repository_ctx, clang, names):
    candidates = []
    clang_dir = _dirname(clang)
    for name in names:
        result = repository_ctx.execute(
            [clang, "--print-prog-name=" + name],
            quiet = True,
        )
        if result.return_code != 0:
            continue
        printed_name = result.stdout.strip()
        if not printed_name:
            continue
        if printed_name.startswith("/"):
            _append_unique(candidates, [printed_name])
        else:
            _append_unique(candidates, [
                _join_path(clang_dir, printed_name),
                printed_name,
            ])
    return candidates

def _detect_clang_resource_include(repository_ctx, clang):
    explicit = repository_ctx.getenv("IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_CLANG_RESOURCE_INCLUDE")
    if not explicit:
        explicit = repository_ctx.getenv("IREE_CLANG_BUILTIN_HEADERS_PATH")
    if explicit:
        path = repository_ctx.path(explicit)
        if not path.exists:
            fail("Configured clang resource include directory does not exist: {}".format(explicit))
        return str(path)

    result = repository_ctx.execute([clang, "-print-resource-dir"], quiet = True)
    if result.return_code != 0:
        fail("Failed to query clang resource directory from {}:\n{}".format(
            clang,
            result.stderr,
        ))
    include_dir = _join_path(result.stdout.strip(), "include")
    path = repository_ctx.path(include_dir)
    if not path.exists:
        fail("clang resource include directory does not exist: {}".format(include_dir))
    return str(path)

def _write_local_toolchain_repo(repository_ctx, tools, clang_resource_include):
    repository_ctx.file("prebuilt_tool.bzl", _PREBUILT_TOOL_BZL)
    repository_ctx.file("bin/.keep", "")
    for name, path in tools.items():
        repository_ctx.symlink(repository_ctx.path(path), "bin/" + name)
    repository_ctx.symlink(
        repository_ctx.path(clang_resource_include),
        "clang_resource_include",
    )
    repository_ctx.file("BUILD.bazel", _LOCAL_BUILD_TEMPLATE)
    repository_ctx.file("paths.bzl", _LOCAL_PATHS_TEMPLATE)

def _amdgpu_device_toolchain_repo_impl(repository_ctx):
    mode = repository_ctx.getenv("IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN") or "none"
    if mode not in ["none", "auto", "rocm", "llvm-tools", "llvm-project"]:
        fail("Unknown IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN '{}'. Available: none, auto, rocm, llvm-tools, llvm-project.".format(mode))

    if mode == "none":
        repository_ctx.file("BUILD.bazel", _STUB_BUILD)
        repository_ctx.file("paths.bzl", _STUB_PATHS)
        return

    if mode == "llvm-project":
        repository_ctx.file("BUILD.bazel", _LLVM_PROJECT_BUILD)
        repository_ctx.file("paths.bzl", _LLVM_PROJECT_PATHS)
        return

    tool_dirs = _candidate_tool_dirs(repository_ctx, mode)
    if mode == "rocm" and not tool_dirs:
        fail("IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN=rocm requires a ROCm/TheRock root such as IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_ROCM_PATH, IREE_ROCM_PATH, ROCM_PATH, ROCM_ROOT, ROCM_HOME, HIP_PATH, hipcc on PATH, or /opt/rocm.")

    clang = _find_tool(
        repository_ctx,
        ["clang", "amdclang", "clang-22"],
        tool_dirs,
        ["IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_CLANG_BINARY", "IREE_CLANG_BINARY", "CLANG"],
        "clang with AMDGPU support",
    )
    llvm_link = _find_tool(
        repository_ctx,
        ["llvm-link"],
        tool_dirs,
        ["IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLVM_LINK_BINARY", "IREE_LLVM_LINK_BINARY", "LLVM_LINK"],
        "llvm-link",
        additional_candidates = _clang_tool_candidates(repository_ctx, clang, ["llvm-link"]),
    )
    lld = _find_tool(
        repository_ctx,
        ["lld", "ld.lld"],
        tool_dirs,
        ["IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLD_BINARY", "IREE_LLD_BINARY", "LLD", "LD_LLD"],
        "lld",
        additional_candidates = _clang_tool_candidates(repository_ctx, clang, ["lld", "ld.lld"]),
    )
    llvm_objcopy = _find_tool(
        repository_ctx,
        ["llvm-objcopy"],
        tool_dirs,
        ["IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLVM_OBJCOPY_BINARY", "IREE_LLVM_OBJCOPY_BINARY", "LLVM_OBJCOPY"],
        "llvm-objcopy",
        additional_candidates = _clang_tool_candidates(repository_ctx, clang, ["llvm-objcopy"]),
    )
    clang_resource_include = _detect_clang_resource_include(repository_ctx, clang)

    _write_local_toolchain_repo(
        repository_ctx,
        {
            "clang": clang,
            "llvm-link": llvm_link,
            "lld": lld,
            "llvm-objcopy": llvm_objcopy,
        },
        clang_resource_include,
    )

amdgpu_device_toolchain_repo = repository_rule(
    implementation = _amdgpu_device_toolchain_repo_impl,
    environ = [
        "CLANG",
        "HIP_PATH",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_CLANG_BINARY",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_CLANG_RESOURCE_INCLUDE",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLD_BINARY",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLVM_LINK_BINARY",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLVM_OBJCOPY_BINARY",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_LLVM_TOOLS_DIR",
        "IREE_HAL_AMDGPU_DEVICE_TOOLCHAIN_ROCM_PATH",
        "IREE_CLANG_BINARY",
        "IREE_CLANG_BUILTIN_HEADERS_PATH",
        "IREE_HOST_BIN_DIR",
        "IREE_HOST_TOOLS",
        "IREE_HOST_TOOLS_DIR",
        "IREE_BINARY_DIR",
        "IREE_LLD_BINARY",
        "IREE_LLVM_LINK_BINARY",
        "IREE_LLVM_OBJCOPY_BINARY",
        "IREE_LLVM_TOOL_DIR",
        "IREE_LLVM_TOOLS_DIR",
        "IREE_ROCM_PATH",
        "LD_LLD",
        "LLD",
        "LLVM_BINARY_DIR",
        "LLVM_LINK",
        "LLVM_OBJCOPY",
        "LLVM_TOOLS_BINARY_DIR",
        "ROCM_HOME",
        "ROCM_PATH",
        "ROCM_ROOT",
    ],
    doc = "Exposes optional AMDGPU builtin device compiler tools.",
)