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

patchelf(
    name = "libpjrt_rocm_so",
    src = "libpjrt_rocm.so",
    add_needed = [
        "libzmlxrocm.so.0",
    ],
    rename_dynamic_symbols = {
        "dlopen": "zmlxrocm_dlopen",
    },
    set_rpath = "$ORIGIN",
)

copy_to_directory(
    name = "sandbox",
    srcs = [
        ":libpjrt_rocm_so",
        "@comgr//:amd_comgr",
        "@hip-runtime-amd//:amdhip",
        "@hip-runtime-amd//:hiprtc",
        "@hipblaslt//:hipblaslt",
        "@hipfft",
        "@hipsolver",
        "@hsa-amd-aqlprofile//:hsa-amd-aqlprofile",
        "@hsa-rocr//:hsa-runtime",
        "@libdrm-amdgpu-amdgpu1",
        "@libdrm-amdgpu-common//:amdgpu_ids",
        "@libdrm2-amdgpu",
        "@libelf1",
        "@libnuma1",
        "@libtinfo6",
        "@libzstd1",
        "@miopen-hip//:MIOpen",
        "@rccl",
        "@rocblas//:rocblas",
        "@rocblas//:runfiles",
        "@rocfft",
        "@rocm-core",
        "@rocm-device-libs//:runfiles",
        "@rocm-smi-lib//:rocm_smi",
        "@rocprofiler-register",
        "@rocsolver",
        "@roctracer",
        "@roctracer//:roctx",
        "@zlib1g",
        "@zml//runtimes/rocm:zmlxrocm",
    ] + select({
        ":_hipblaslt": ["@hipblaslt//:runfiles"],
        "//conditions:default": [],
    }),
    replace_prefixes = {
        "amdhip": "lib",
        "hipblaslt": "lib",
        "lib/x86_64-linux-gnu": "lib",
        "libdrm-amdgpu-amdgpu1": "lib",
        "libelf1": "lib",
        "libpjrt_rocm_so": "lib",
        "opt/amdgpu/lib/x86_64-linux-gnu": "lib",
        "opt/amdgpu/share": "share",
        "rocblas": "lib",
        "runtimes/rocm": "lib",
        "usr/lib/x86_64-linux-gnu": "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__"],
)