148 lines
3.5 KiB
Python
148 lines
3.5 KiB
Python
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",
|
|
copts = ["-lc"],
|
|
kind = BINARY_KIND.shared_lib,
|
|
main = "libpjrt_neuron.zig",
|
|
visibility = ["@libpjrt_neuron//:__subpackages__"],
|
|
deps = [
|
|
"//stdx",
|
|
"@rules_zig//zig/runfiles",
|
|
],
|
|
)
|
|
|
|
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",
|
|
)
|
|
|
|
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",
|
|
data = select({
|
|
"//runtimes:neuron.enabled": [
|
|
":libneuronxla",
|
|
],
|
|
"//conditions:default": [],
|
|
}),
|
|
import_name = "runtimes/neuron",
|
|
main = "neuron.zig",
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
"//pjrt",
|
|
] + select({
|
|
"//runtimes:neuron.enabled": [
|
|
":libneuronxla_pyenv",
|
|
":libpjrt_neuron",
|
|
":libpython",
|
|
"//async",
|
|
"//stdx",
|
|
"@rules_zig//zig/runfiles",
|
|
],
|
|
"//conditions:default": [":empty"],
|
|
}),
|
|
)
|
|
|
|
filegroup(
|
|
name = "layers",
|
|
srcs = [
|
|
"@apt_neuron//bash/amd64",
|
|
],
|
|
visibility = ["//visibility:public"],
|
|
)
|