Radix/runtimes/neuron/BUILD.bazel

load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_python//python:pip.bzl", "compile_pip_requirements")
load("@rules_python//python/entry_points:py_console_script_binary.bzl", "py_console_script_binary")
load("@rules_uv//uv:pip.bzl", "pip_compile")
load("@rules_zig//zig:defs.bzl", "BINARY_KIND", "zig_binary", "zig_library")
load(":neuron.bzl", "py_binary_with_script")
load(":pyenv.bzl", "pyenv_zig")

pyenv_zig(
    name = "libneuronxla_pyenv_zig",
    deps = [":libneuronxla"],
)

zig_library(
    name = "libneuronxla_pyenv",
    main = ":libneuronxla_pyenv_zig",
)

# A proxy PJRT Plugin that loads the Neuron PJRT Plugin
# and returns the instance from nested GetPjrtApi.
# 
# Additionally, it provides a way to load implicit transitive dependencies
# of neuronx-cc (see add_needed of the patchelf target below).
#
# TODO(cerisier): Use a zig_cc_shared_library instead.
zig_binary(
    name = "libpjrt_neuron_proxy",
    kind = BINARY_KIND.shared_lib,
    main = "libpjrt_neuron.zig",
    copts = ["-lc"],
    deps = [
        "//stdx",
        "@rules_zig//zig/runfiles",
    ],
    visibility = ["@libpjrt_neuron//:__subpackages__"],
)

pip_compile(
    name = "update_requirements",
    args = [
        "--generate-hashes",
        "--emit-index-url",
        "--emit-find-links",
        "--no-strip-extras",
        "--index-strategy=unsafe-best-match",
        "--upgrade",
    ],
    python_platform = "x86_64-unknown-linux-gnu",
    requirements_in = "requirements.in",
    requirements_txt = "requirements.lock.txt",
    tags = [
        "manual",
    ],
)

py_console_script_binary(
    name = "neuronx-cc",
    binary_rule = py_binary_with_script,
    pkg = "@neuron_py_deps//neuronx_cc",
    visibility = ["//visibility:public"],
)

py_binary_with_script(
    name = "libneuronxla",
    srcs = ["empty.py"],
    main = "empty.py",
    deps = [
        ":neuronx-cc",
        "@neuron_py_deps//libneuronxla",
    ],
)

cc_library(
    name = "libpython",
    hdrs = ["libpython.h"],
    visibility = ["//visibility:public"],
    deps = [
        "@rules_python//python/cc:current_py_cc_headers",
        "@rules_python//python/cc:current_py_cc_libs",
    ],
)

compile_pip_requirements(
    name = "requirements",
    src = "requirements.in",
    py_binary = py_binary_with_script,
    requirements_txt = "requirements.lock.txt",
    tags = [
        "manual",
    ],
)

cc_library(
    name = "zmlxneuron_lib",
    srcs = ["zmlxneuron.c"],
    linkopts = ["-ldl"],
    visibility = ["@libpjrt_neuron//:__subpackages__"],
)

alias(
    name = "bootstrap",
    actual = "@rules_python//python/config_settings:bootstrap_impl",
)

cc_library(
    name = "empty",
)

cc_library(
    name = "libpjrt_neuron",
    hdrs = ["libpjrt_neuron.h"],
    defines = ["ZML_RUNTIME_NEURON"],
    deps = [
        "@libpjrt_neuron",
    ],
)

zig_library(
    name = "neuron",
    import_name = "runtimes/neuron",
    main = "neuron.zig",
    visibility = ["//visibility:public"],
    data = select({
        "//runtimes:neuron.enabled": [
            ":libneuronxla",
        ],
        "//conditions:default": [],
    }),
    deps = [
        "//pjrt",
    ] + select({
        "//runtimes:neuron.enabled": [
            ":libpjrt_neuron",
            ":libneuronxla_pyenv",
            ":libpython",
            "//async",
            "//stdx",
            "@rules_zig//zig/runfiles",
        ],
        "//conditions:default": [":empty"],
    }),
)

filegroup(
    name = "layers",
    srcs = [
        "@apt_neuron//bash/amd64",
    ],
    visibility = ["//visibility:public"],
)