Radix/runtimes/rocm/libpjrt_rocm.BUILD.bazel

load("@aspect_bazel_lib//lib:copy_to_directory.bzl", "copy_to_directory")
load("@aspect_bazel_lib//lib:tar.bzl", "tar")
load("@bazel_skylib//rules:common_settings.bzl", "bool_flag", "string_list_flag")
load("@zml//bazel:patchelf.bzl", "patchelf")

string_list_flag(
    name = "gfx",
    build_setting_default = ["all"],
    visibility = ["//visibility:public"],
)

bool_flag(
    name = "hipblaslt",
    build_setting_default = True,
    visibility = ["//visibility:public"],
)

config_setting(
    name = "_hipblaslt",
    flag_values = {":hipblaslt": "True"},
)

cc_shared_library(
    name = "zmlxrocm_so",
    shared_lib_name = "lib/libzmlxrocm.so.0",
    deps = ["@zml//runtimes/rocm:zmlxrocm_lib"],
)

patchelf(
    name = "libpjrt_rocm.patchelf",
    shared_library = "libpjrt_rocm.so",
    add_needed = [
        "libzmlxrocm.so.0",
        # So that RPATH is taken into account.
        "librocblas.so.4",
        "libMIOpen.so.1",
    ] + select({
        "_hipblaslt": [
            "libhipblaslt.so.0",
        ],
        "//conditions:default": [],
    }),
    rename_dynamic_symbols = {
        "dlopen": "zmlxrocm_dlopen",
    },
    set_rpath = "$ORIGIN",
)

copy_to_directory(
    name = "sandbox",
    srcs = [
        ":zmlxrocm_so",
        ":libpjrt_rocm.patchelf",
        "@comgr//:amd_comgr",
        "@hip-runtime-amd//:amdhip_patched",
        "@hip-runtime-amd//:hiprtc",
        "@hipblaslt//:hipblaslt",
        "@hipfft",
        "@hipsolver",
        "@hsa-amd-aqlprofile//:hsa-amd-aqlprofile",
        "@hsa-rocr//:hsa-runtime",
        "@miopen-hip//:MIOpen",
        "@rccl",
        "@rocblas//:rocblas",
        "@rocblas//:runfiles",
        "@rocm-core",
        "@rocm-device-libs//:runfiles",
        "@rocm-smi-lib//:rocm_smi",
        "@rocprofiler-register",
        "@rocfft",
        "@rocsolver",
        "@roctracer",
        "@roctracer//:roctx",
        "@libelf1",
        "@libdrm2-amdgpu",
        "@libnuma1",
        "@libzstd1",
        "@libdrm-amdgpu-amdgpu1",
        "@libtinfo6",
        "@zlib1g",
    ] + select({
        ":_hipblaslt": ["@hipblaslt//:runfiles"],
        "//conditions:default": [],
    }),
    replace_prefixes = {
        "libpjrt_rocm.patchelf": "lib",
        "lib/x86_64-linux-gnu": "lib",
        "usr/lib/x86_64-linux-gnu": "lib",
        "libelf1": "lib",
        "hipblaslt": "lib",
        "rocblas": "lib",
        "opt/amdgpu/lib/x86_64-linux-gnu": "lib",
        "libdrm-amdgpu-amdgpu1": "lib",
        "amdhip_patched": "lib",
    },
    add_directory_to_runfiles = True,
    include_external_repositories = ["**"],
)

tar(
    name = "amdgpu_ids_layer",
    srcs = ["@libdrm-amdgpu-common//:amdgpu_ids"],
    visibility = ["//visibility:public"],
)

cc_library(
    name = "libpjrt_rocm",
    data = [":sandbox"],
    visibility = ["@zml//runtimes/rocm:__subpackages__"],
)