Radix/runtimes/cuda/libpjrt_cuda.BUILD.bazel

load("@aspect_bazel_lib//lib:copy_to_directory.bzl", "copy_to_directory")
load("@zml//bazel:patchelf.bzl", "patchelf")

patchelf(
    name = "libpjrt_cuda_so",
    src = "libpjrt_cuda.so",
    add_needed = [
        "libzmlxcuda.so.0",
    ],
    rename_dynamic_symbols = {
        "dlopen": "zmlxcuda_dlopen",
    },
    set_rpath = "$ORIGIN",
)

copy_to_directory(
    name = "sandbox",
    srcs = [
        ":libpjrt_cuda_so",
        "@cuda_cudart",
        "@cuda_cupti",
        "@cuda_nvcc",
        "@cuda_nvrtc",
        "@cuda_nvtx",
        "@cudnn",
        "@libnvshmem",
        "@libcublas",
        "@libcufft",
        "@libcusolver",
        "@libcusparse",
        "@libnvjitlink",
        "@nccl",
        "@zlib1g",
        "@zml//runtimes/cuda:zmlxcuda",
    ],
    replace_prefixes = {
        "nvidia/nccl/lib": "lib",
        "nvvm/lib64": "lib",
        "libpjrt_cuda_so": "lib",
        "lib/x86_64-linux-gnu": "lib",
        "runtimes/cuda": "lib",
    },
    add_directory_to_runfiles = False,
    include_external_repositories = ["**"],
)

cc_library(
    name = "libpjrt_cuda",
    data = [":sandbox"],
    deps = [
        "@cuda_cudart//:cuda",
    ],
    linkopts = [
        # Defer function call resolution until the function is called
        # (lazy loading) rather than at load time.
        #
        # This is required because we want to let downstream use weak CUDA symbols.
        #
        # We force it here because -z,now (which resolve all symbols at load time),
        # is the default in most bazel CC toolchains as well as in certain linkers.
        "-Wl,-z,lazy",
    ],
    visibility = ["@zml//runtimes/cuda:__subpackages__"],
)