53 lines
1.8 KiB
Python
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 = "cf5ea44b14a6ddc320c5b1d7bb88328dda099571e106b17a214b6ec586e321b8",
|
|
url = "https://github.com/zml/pjrt-artifacts/releases/download/v12.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 = "1b48c26b0d2709730df6921a86d9ae97ab71d3f2cfa17fd1a370e85235370914",
|
|
url = "https://github.com/zml/pjrt-artifacts/releases/download/v12.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 = "66dc6c65933a6d7985b1c69837d8abe13750cd1f06704427f6989c3f952d3511",
|
|
url = "https://github.com/zml/pjrt-artifacts/releases/download/v12.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,
|
|
)
|