Radix/runtimes/rocm/libpjrt_rocm.BUILD.bazel

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

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

copy_to_directory(
    name = "sandbox",
    srcs = [
        "@rocm-device-libs//:runfiles",
        "@rocm-llvm//:lld",
    ],
    include_external_repositories = ["*"],
)

cc_library(
    name = "zmlxrocm_lib",
    data = ["@rocblas//:runfiles"],
    srcs = ["@zml//runtimes/rocm:zmlxrocm_srcs"],
    linkopts = [
        "-lc",
        "-ldl",
    ],
)

cc_shared_library(
    name = "zmlxrocm_so",
    shared_lib_name = "libzmlxrocm.so.0",
    deps = [":zmlxrocm_lib"],
)

cc_import(
    name = "zmlxrocm",
    shared_library = ":zmlxrocm_so",
)

cc_import(
    name = "libpjrt_rocm",
    data = [
        ":sandbox",
        "@rocblas//:runfiles",
    ] + select({
        ":_hipblaslt": ["@hipblaslt//:runfiles"],
        "//conditions:default": [],
    }),
    add_needed = ["libzmlxrocm.so.0"],
    rename_dynamic_symbols = {
        "dlopen": "zmlxrocm_dlopen",
    },
    shared_library = "libpjrt_rocm.so",
    soname = "libpjrt_rocm.so",
    visibility = ["//visibility:public"],
    deps = [
        "@comgr//:amd_comgr",
        "@hip-runtime-amd//:amdhip",
        "@hipblaslt",
        "@hsa-amd-aqlprofile//:hsa-amd-aqlprofile",
        "@hsa-rocr//:hsa-runtime",
        "@miopen-hip//:MIOpen",
        "@rccl",
        "@rocblas",
        "@rocm-core",
        "@rocm-smi-lib//:rocm_smi",
        "@rocprofiler-register",
        "@roctracer",
        "@libelf1",
        "@libdrm2",
        "@libnuma1",
        "@libzstd1",
        "@libdrm-amdgpu1",
        "@libtinfo6",
        "@zlib1g",
        "@zml//runtimes/rocm:zmlxrocm",
    ],
)