Radix/runtimes/neuron/libpjrt_neuron.BUILD.bazel

load("@aspect_bazel_lib//lib:copy_to_directory.bzl", "copy_to_directory")
load("@zml//bazel:patchelf.bzl", "patchelf")
load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
load("@bazel_skylib//rules:select_file.bzl", "select_file")
load("@rules_cc//cc:cc_shared_library.bzl", "cc_shared_library")
load("@zml//bazel:runfiles.bzl", "runfiles_to_default")
load("@rules_python//python:py_binary.bzl", "py_binary")

runfiles_to_default(
    name = "libneuronxla_files",
    deps = ["@neuron_py_deps//libneuronxla:pkg"],
)

select_file(
    name = "libneuronpjrt_so_orig",
    srcs = ":libneuronxla_files",
    subpath = "site-packages/libneuronxla/libneuronpjrt.so",
)

patchelf(
    name = "libneuronpjrt_so",
    set_rpath = '$ORIGIN',
    src = ":libneuronpjrt_so_orig",
    soname = "libneuronpjrt.so",
)

patchelf(
    name = "libpjrt_neuron_so",
    set_rpath = '$ORIGIN',
    src = "@zml//runtimes/neuron:libpjrt_neuron",
    soname = "libpjrt_neuron.so",
)

copy_to_directory(
    name = "sandbox",
    srcs = [
        ":libneuronpjrt_so",
        ":libpjrt_neuron_so",
        "@zml//runtimes/neuron:neuronx-cc_files",
        "@zml//runtimes/neuron:libneuronxla",
        "@rules_python//python:current_py_toolchain",
        "@aws-neuronx-runtime-lib//:libnrt.patchelf",
        "@aws-neuronx-runtime-lib//:libncfw.patchelf",
        "@aws-neuronx-collectives//:libnccom",
        "@zlib1g",
        "@libgomp1",
    ],
    replace_prefixes = {
        "runtimes/neuron/neuronx-cc": "bin/neuronx-cc",
        "runtimes/neuron/libneuronxla.so": "site-packages/libneuronxla.so",
        "libneuronpjrt_so": "lib",
        "libpjrt_neuron_so": "lib",
        "libnrt.patchelf": "lib",
        "libncfw.patchelf": "lib",
        "lib/x86_64-linux-gnu": "lib",
        "usr/lib/x86_64-linux-gnu": "lib",
        "opt/neuron": "lib",
    },
    exclude_srcs_patterns = [
        "**/tests/**",
        "**/include/**",
        "**/_solib_*/**",
    ],
    include_srcs_patterns = [
        "runtimes/neuron/neuronx-cc",
        "**/lib**",
        "lib/python*/**",
        "site-packages/**",
    ],
    allow_overwrites = True,
    add_directory_to_runfiles = False,
    include_external_repositories = ["**"],
    hardlink = "on",
)

cc_library(
    name = "libnrt_headers",
    deps = [
        "@aws-neuronx-runtime-lib//:libnrt_headers",
    ],
    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 NRT 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/neuron:__subpackages__"],
)

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