Radix/runtimes/cpu/cpu.bzl

53 lines
1.8 KiB
Python

load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("//runtimes/common:packages.bzl", "packages")
_BUILD_FILE_DEFAULT_VISIBILITY = """\
package(default_visibility = ["//visibility:public"])
"""
_BUILD_LINUX = "\n".join([
packages.filegroup(
name = "libpjrt_cpu",
srcs = ["libpjrt_cpu.so"],
visibility = ["@zml//runtimes/cpu:__subpackages__"],
),
])
_BUILD_DARWIN = packages.filegroup(
name = "libpjrt_cpu",
srcs = ["libpjrt_cpu.dylib"],
visibility = ["@zml//runtimes/cpu:__subpackages__"],
)
def _cpu_pjrt_plugin_impl(mctx):
http_archive(
name = "libpjrt_cpu_linux_amd64",
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + _BUILD_LINUX,
sha256 = "124dc500291a5930f910ca23533520e22c90797110b29fd2c0d8274475f4a220",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v13.0.0/pjrt-cpu_linux-amd64.tar.gz",
)
http_archive(
name = "libpjrt_cpu_darwin_amd64",
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + _BUILD_DARWIN,
sha256 = "6e5b59874880f4db37c53fb1d52520d410b0078f9d2606a90762c6c622693c26",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v13.0.0/pjrt-cpu_darwin-amd64.tar.gz",
)
http_archive(
name = "libpjrt_cpu_darwin_arm64",
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + _BUILD_DARWIN,
sha256 = "a6354bfed828a011e6d809eda2230e10c40c80044c67fe618b2a9615c047f092",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v13.0.0/pjrt-cpu_darwin-arm64.tar.gz",
)
return mctx.extension_metadata(
reproducible = True,
root_module_direct_deps = "all",
root_module_direct_dev_deps = [],
)
cpu_pjrt_plugin = module_extension(
implementation = _cpu_pjrt_plugin_impl,
)