tarry.singh/Radix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update Bazel workspace and runtime configs: rework sandboxing, bump PJRT to 7.0.0, and upgrade CUDA (12.8), cuDNN (9.8), and ROCm (6.3.4).

Browse Source
This commit is contained in:
Tarry Singh 2024-06-25 11:00:29 +00:00
parent 8613a6da7d
commit 967eeb928f
14 changed files with 1110 additions and 3556 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
MODULE.bazel
View File

@ -85,15 +85,21 @@ pip.parse(
)
use_repo(pip, "neuron_py_deps")
common_apt_packages = use_extension("//runtimes/common:packages.bzl", "common_apt_packages")
use_repo(common_apt_packages, "libdrm-amdgpu1", "libdrm2", "libelf1", "libnuma1", "libtinfo6", "libzstd1", "zlib1g")
cpu = use_extension("//runtimes/cpu:cpu.bzl", "cpu_pjrt_plugin")
use_repo(cpu, "libpjrt_cpu_darwin_arm64", "libpjrt_cpu_linux_amd64")
cuda = use_extension("//runtimes/cuda:cuda.bzl", "cuda_packages")
use_repo(cuda, "libpjrt_cuda")
inject_repo(cuda, "zlib1g")
rocm = use_extension("//runtimes/rocm:rocm.bzl", "rocm_packages")
use_repo(rocm, "hipblaslt", "libpjrt_rocm", "rocblas")
inject_repo(rocm, "libdrm-amdgpu1", "libdrm2", "libelf1", "libnuma1", "libtinfo6", "libzstd1", "zlib1g")
tpu = use_extension("//runtimes/tpu:tpu.bzl", "tpu_packages")
use_repo(tpu, "libpjrt_tpu")
@ -155,3 +161,15 @@ use_repo(crate, "crates")
non_module_deps = use_extension("//:third_party/non_module_deps.bzl", "non_module_deps")
use_repo(non_module_deps, "com_google_sentencepiece", "org_swig_swig")
apt.install(
name = "apt_rocm",
lock = "//runtimes/rocm:packages.lock.json",
manifest = "//runtimes/rocm:packages.yaml",
)
apt.install(
name = "apt_common",
lock = "//runtimes/rocm:packages.lock.json",
manifest = "//runtimes/common:packages.yaml",
)
use_repo(apt, "apt_common", "apt_rocm")

3267
MODULE.bazel.lock
View File

File diff suppressed because it is too large Load Diff

4
runtimes/common/BUILD.bazel Normal file
View File

@ -0,0 +1,4 @@
exports_files(
["packages.lock.json"],
visibility = ["//runtimes:__subpackages__"],
)

64
runtimes/common/packages.bzl Normal file
View File

@ -0,0 +1,64 @@
load("//bazel:http_deb_archive.bzl", "http_deb_archive")
_BUILD_FILE_DEFAULT_VISIBILITY = """\
package(default_visibility = ["//visibility:public"])
"""
def _kwargs(**kwargs):
return repr(struct(**kwargs))[len("struct("):-1]
def _cc_import(**kwargs):
return """cc_import({})""".format(_kwargs(**kwargs))
def _filegroup(**kwargs):
return """filegroup({})""".format(_kwargs(**kwargs))
def _load(bzl, name):
return """load({}, {})""".format(repr(bzl), repr(name))
def _read(mctx, labels):
ret = {}
for label in labels:
data = json.decode(mctx.read(Label(label)))
ret.update({
pkg["name"]: pkg
for pkg in data["packages"]
})
return ret
_DEBIAN_PACKAGES = {
"libdrm2": _cc_import(name = "libdrm2", shared_library = "usr/lib/x86_64-linux-gnu/libdrm.so.2"),
"libelf1": _cc_import(name = "libelf1", shared_library = "usr/lib/x86_64-linux-gnu/libelf.so.1"),
"libnuma1": _cc_import(name = "libnuma1", shared_library = "usr/lib/x86_64-linux-gnu/libnuma.so.1"),
"libzstd1": _cc_import(name = "libzstd1", shared_library = "usr/lib/x86_64-linux-gnu/libzstd.so.1"),
"libdrm-amdgpu1": _cc_import(name = "libdrm-amdgpu1", shared_library = "usr/lib/x86_64-linux-gnu/libdrm_amdgpu.so.1"),
"libtinfo6": _cc_import(name = "libtinfo6", shared_library = "lib/x86_64-linux-gnu/libtinfo.so.6"),
"zlib1g": _cc_import(name = "zlib1g", shared_library = "lib/x86_64-linux-gnu/libz.so.1"),
}
def _common_apt_packages_impl(mctx):
loaded_packages = packages.read(mctx, ["packages.lock.json"])
for pkg_name, build_file_content in _DEBIAN_PACKAGES.items():
pkg = loaded_packages[pkg_name]
http_deb_archive(
name = pkg_name,
urls = pkg["urls"],
sha256 = pkg["sha256"],
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + build_file_content,
)
return mctx.extension_metadata(
reproducible = True,
root_module_direct_deps = "all",
root_module_direct_dev_deps = [],
)
common_apt_packages = module_extension(
implementation = _common_apt_packages_impl,
)
packages = struct(
read = _read,
cc_import = _cc_import,
filegroup = _filegroup,
load_ = _load,
)

258
runtimes/common/packages.lock.json Executable file
View File

@ -0,0 +1,258 @@
{
"packages": [
{
"arch": "amd64",
"dependencies": [
{
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"version": "2.36-9+deb12u9"
},
{
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"version": "12.2.0-14"
},
{
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"version": "12.2.0-14"
},
{
"key": "libdrm-common_2.4.114-1_amd64",
"name": "libdrm-common",
"version": "2.4.114-1"
}
],
"key": "libdrm2_2.4.114-1-p-b1_amd64",
"name": "libdrm2",
"sha256": "be18fb670797ba32da9628cf3e8acd83160d8db8c8dd842501dd8e401c3b5371",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/libd/libdrm/libdrm2_2.4.114-1+b1_amd64.deb"
],
"version": "2.4.114-1+b1"
},
{
"arch": "amd64",
"dependencies": [],
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"sha256": "3a043b9dbf1a8c8b5a9e2268253e6dca8c7d431bd7a202fea96364abf374fea9",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/g/glibc/libc6_2.36-9+deb12u9_amd64.deb"
],
"version": "2.36-9+deb12u9"
},
{
"arch": "amd64",
"dependencies": [],
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"sha256": "f3d1d48c0599aea85b7f2077a01d285badc42998c1a1e7473935d5cf995c8141",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/g/gcc-12/libgcc-s1_12.2.0-14_amd64.deb"
],
"version": "12.2.0-14"
},
{
"arch": "amd64",
"dependencies": [],
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"sha256": "1a03df5a57833d65b5bb08cfa19d50e76f29088dc9e64fb934af42d9023a0807",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/g/gcc-12/gcc-12-base_12.2.0-14_amd64.deb"
],
"version": "12.2.0-14"
},
{
"arch": "amd64",
"dependencies": [],
"key": "libdrm-common_2.4.114-1_amd64",
"name": "libdrm-common",
"sha256": "32f9664138b38b224383c6986457d5ad2ec8efd559b1a0ce7749405f7a451aad",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/libd/libdrm/libdrm-common_2.4.114-1_all.deb"
],
"version": "2.4.114-1"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "zlib1g_1-1.2.13.dfsg-1_amd64",
"name": "zlib1g",
"version": "1:1.2.13.dfsg-1"
},
{
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"version": "2.36-9+deb12u9"
},
{
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"version": "12.2.0-14"
},
{
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"version": "12.2.0-14"
}
],
"key": "libelf1_0.188-2.1_amd64",
"name": "libelf1",
"sha256": "619add379c606b3ac6c1a175853b918e6939598a83d8ebadf3bdfd50d10b3c8c",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/e/elfutils/libelf1_0.188-2.1_amd64.deb"
],
"version": "0.188-2.1"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"version": "2.36-9+deb12u9"
},
{
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"version": "12.2.0-14"
},
{
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"version": "12.2.0-14"
}
],
"key": "zlib1g_1-1.2.13.dfsg-1_amd64",
"name": "zlib1g",
"sha256": "d7dd1d1411fedf27f5e27650a6eff20ef294077b568f4c8c5e51466dc7c08ce4",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/z/zlib/zlib1g_1.2.13.dfsg-1_amd64.deb"
],
"version": "1:1.2.13.dfsg-1"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"version": "2.36-9+deb12u9"
},
{
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"version": "12.2.0-14"
},
{
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"version": "12.2.0-14"
}
],
"key": "libnuma1_2.0.16-1_amd64",
"name": "libnuma1",
"sha256": "639e1ab6bd66ead40db8a22c332d7199679fa22db261cac34444eb8eb4c17dda",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/n/numactl/libnuma1_2.0.16-1_amd64.deb"
],
"version": "2.0.16-1"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"version": "2.36-9+deb12u9"
},
{
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"version": "12.2.0-14"
},
{
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"version": "12.2.0-14"
}
],
"key": "libzstd1_1.5.4-p-dfsg2-5_amd64",
"name": "libzstd1",
"sha256": "6315b5ac38b724a710fb96bf1042019398cb656718b1522279a5185ed39318fa",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/libz/libzstd/libzstd1_1.5.4+dfsg2-5_amd64.deb"
],
"version": "1.5.4+dfsg2-5"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "libdrm2_2.4.114-1-p-b1_amd64",
"name": "libdrm2",
"version": "2.4.114-1+b1"
},
{
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"version": "2.36-9+deb12u9"
},
{
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"version": "12.2.0-14"
},
{
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"version": "12.2.0-14"
},
{
"key": "libdrm-common_2.4.114-1_amd64",
"name": "libdrm-common",
"version": "2.4.114-1"
}
],
"key": "libdrm-amdgpu1_2.4.114-1-p-b1_amd64",
"name": "libdrm-amdgpu1",
"sha256": "b75a71e96f1faac0f131ac657e09efcbe8968eef62cc34b8abfcff2ff9f0cccd",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/libd/libdrm/libdrm-amdgpu1_2.4.114-1+b1_amd64.deb"
],
"version": "2.4.114-1+b1"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "libc6_2.36-9-p-deb12u9_amd64",
"name": "libc6",
"version": "2.36-9+deb12u9"
},
{
"key": "libgcc-s1_12.2.0-14_amd64",
"name": "libgcc-s1",
"version": "12.2.0-14"
},
{
"key": "gcc-12-base_12.2.0-14_amd64",
"name": "gcc-12-base",
"version": "12.2.0-14"
}
],
"key": "libtinfo6_6.4-4_amd64",
"name": "libtinfo6",
"sha256": "072d908f38f51090ca28ca5afa3b46b2957dc61fe35094c0b851426859a49a51",
"urls": [
"https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/n/ncurses/libtinfo6_6.4-4_amd64.deb"
],
"version": "6.4-4"
}
],
"version": 1
}

20
runtimes/common/packages.yaml Normal file
View File

@ -0,0 +1,20 @@
#
# bazel run @rocm_apt//:lock
#
version: 1
sources:
- channel: bookworm main
url: https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/
archs:
- "amd64"
packages:
- "libdrm2"
- "libelf1"
- "libnuma1"
- "libzstd1"
- "libdrm-amdgpu1"
- "libtinfo6"
- "zlib1g"

37
runtimes/cpu/cpu.bzl
View File

@ -1,36 +1,39 @@
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("//runtimes/common:packages.bzl", "packages")
_BUILD_LINUX = """\
load("@zml//bazel:cc_import.bzl", "cc_import")
cc_import(
name = "libpjrt_cpu",
shared_library = "libpjrt_cpu.so",
soname = "libpjrt_cpu.so",
visibility = ["@zml//runtimes/cpu:__subpackages__"],
)
_BUILD_FILE_DEFAULT_VISIBILITY = """\
package(default_visibility = ["//visibility:public"])
"""
_BUILD_DARWIN = """\
cc_import(
_BUILD_LINUX = "\n".join([
packages.load_("@zml//bazel:cc_import.bzl", "cc_import"),
packages.cc_import(
name = "libpjrt_cpu",
shared_library = "libpjrt_cpu.so",
soname = "libpjrt_cpu.so",
visibility = ["@zml//runtimes/cpu:__subpackages__"],
),
])
_BUILD_DARWIN = packages.cc_import(
name = "libpjrt_cpu",
shared_library = "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_LINUX,
sha256 = "e17d15331a3f42b90c8131459c235a9cf1145913a581c8b192845c50d313a7d6",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v6.0.0/pjrt-cpu_linux-amd64.tar.gz",
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + _BUILD_LINUX,
sha256 = "1cda1325095c12bd0019838d28ee92d811ac478d22ed3c08020d5a0cd2d9f34a",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v7.0.0/pjrt-cpu_linux-amd64.tar.gz",
)
http_archive(
name = "libpjrt_cpu_darwin_arm64",
build_file_content = _BUILD_DARWIN,
sha256 = "2e18f8426ec5fa698163d55871fdbaed3616325cab10fad54f9d14f265fbf00d",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v6.0.0/pjrt-cpu_darwin-arm64.tar.gz",
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + _BUILD_DARWIN,
sha256 = "da4deaf850d715997614768b2fc0283595ee8181133ab3243d65635e3439de69",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v7.0.0/pjrt-cpu_darwin-arm64.tar.gz",
)
return mctx.extension_metadata(

139
runtimes/cuda/cuda.bzl
View File

@ -1,105 +1,85 @@
load("@bazel_skylib//lib:paths.bzl", "paths")
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("//bazel:http_deb_archive.bzl", "http_deb_archive")
load("//runtimes/common:packages.bzl", "packages")
_BUILD_FILE_DEFAULT_VISIBILITY = """\
package(default_visibility = ["//visibility:public"])
"""
ARCH = "linux-x86_64"
CUDA_VERSION = "12.6.3"
CUDNN_VERSION = "9.6.0"
CUDA_REDIST_PREFIX = "https://developer.download.nvidia.com/compute/cuda/redist/"
CUDA_VERSION = "12.8.1"
CUDA_REDIST_JSON_SHA256 = "249e28a83008d711d5f72880541c8be6253f6d61608461de4fcb715554a6cf17"
def _filegroup(name, srcs):
return """\
filegroup(
name = {name},
srcs = {srcs},
visibility = ["@libpjrt_cuda//:__subpackages__"],
)
""".format(name = repr(name), srcs = repr(srcs))
def _cc_import(name, shared_library, deps = []):
return """\
cc_import(
name = {name},
shared_library = {shared_library},
deps = {deps},
visibility = ["@libpjrt_cuda//:__subpackages__"],
)
""".format(name = repr(name), shared_library = repr(shared_library), deps = repr(deps))
def _cc_import_static(name, static_library, deps = []):
return """\
cc_import(
name = {name},
static_library = {static_library},
deps = {deps},
visibility = ["@libpjrt_cuda//:__subpackages__"],
)
""".format(name = repr(name), static_library = repr(static_library), deps = repr(deps))
CUDNN_REDIST_PREFIX = "https://developer.download.nvidia.com/compute/cudnn/redist/"
CUDNN_VERSION = "9.8.0"
CUDNN_REDIST_JSON_SHA256 = "a1599fa1f8dcb81235157be5de5ab7d3936e75dfc4e1e442d07970afad3c4843"
CUDA_PACKAGES = {
"cuda_cudart": _cc_import(
"cuda_cudart": packages.cc_import(
name = "cudart",
shared_library = "lib/libcudart.so.12",
),
"cuda_cupti": _cc_import(
"cuda_cupti": packages.cc_import(
name = "cupti",
shared_library = "lib/libcupti.so.12",
),
"libcufft": _cc_import(
"libcufft": packages.cc_import(
name = "cufft",
shared_library = "lib/libcufft.so.11",
),
"libcusolver": _cc_import(
"libcusolver": packages.cc_import(
name = "cusolver",
shared_library = "lib/libcusolver.so.11",
),
"libcusparse": _cc_import(
"libcusparse": packages.cc_import(
name = "cusparse",
shared_library = "lib/libcusparse.so.12",
),
"libnvjitlink": _cc_import(
"libnvjitlink": packages.cc_import(
name = "nvjitlink",
shared_library = "lib/libnvJitLink.so.12",
),
"cuda_nvcc": "\n".join([
_filegroup(
packages.filegroup(
name = "ptxas",
srcs = ["bin/ptxas"],
),
_filegroup(
packages.filegroup(
name = "nvlink",
srcs = ["bin/nvlink"],
),
_filegroup(
packages.filegroup(
name = "libdevice",
srcs = ["nvvm/libdevice/libdevice.10.bc"],
),
_cc_import(
packages.cc_import(
name = "nvvm",
shared_library = "nvvm/lib64/libnvvm.so.4",
),
_cc_import_static(
packages.cc_import(
name = "nvptxcompiler",
static_library = "lib/libnvptxcompiler_static.a",
),
]),
"cuda_nvrtc": "\n".join([
_cc_import(
packages.cc_import(
name = "nvrtc",
shared_library = "lib/libnvrtc.so.12",
deps = [":nvrtc_builtins"],
),
_cc_import(
packages.cc_import(
name = "nvrtc_builtins",
shared_library = "lib/libnvrtc-builtins.so.12.6",
shared_library = "lib/libnvrtc-builtins.so.12.8",
),
]),
"libcublas": "\n".join([
_cc_import(
packages.cc_import(
name = "cublasLt",
shared_library = "lib/libcublasLt.so.12",
),
_cc_import(
packages.cc_import(
name = "cublas",
shared_library = "lib/libcublas.so.12",
deps = [":cublasLt"],
@ -109,7 +89,7 @@ CUDA_PACKAGES = {
CUDNN_PACKAGES = {
"cudnn": "\n".join([
_cc_import(
packages.cc_import(
name = "cudnn",
shared_library = "lib/libcudnn.so.9",
deps = [
@ -122,41 +102,59 @@ CUDNN_PACKAGES = {
":cudnn_heuristic",
],
),
_cc_import(
packages.cc_import(
name = "cudnn_adv",
shared_library = "lib/libcudnn_adv.so.9",
),
_cc_import(
packages.cc_import(
name = "cudnn_ops",
shared_library = "lib/libcudnn_ops.so.9",
),
_cc_import(
packages.cc_import(
name = "cudnn_cnn",
shared_library = "lib/libcudnn_cnn.so.9",
deps = [":cudnn_ops"],
),
_cc_import(
packages.cc_import(
name = "cudnn_graph",
shared_library = "lib/libcudnn_graph.so.9",
),
_cc_import(
packages.cc_import(
name = "cudnn_engines_precompiled",
shared_library = "lib/libcudnn_engines_precompiled.so.9",
),
_cc_import(
packages.cc_import(
name = "cudnn_engines_runtime_compiled",
shared_library = "lib/libcudnn_engines_runtime_compiled.so.9",
),
_cc_import(
packages.cc_import(
name = "cudnn_heuristic",
shared_library = "lib/libcudnn_heuristic.so.9",
),
]),
}
def _read_redist_json(mctx, url, sha256):
fname = ".{}.json".format(sha256)
mctx.download(
url = url,
output = fname,
sha256 = sha256,
)
return json.decode(mctx.read(fname))
def _cuda_impl(mctx):
CUDA_REDIST = json.decode(mctx.read(Label("@zml//runtimes/cuda:cuda.redistrib_{}.json".format(CUDA_VERSION))))
CUDNN_REDIST = json.decode(mctx.read(Label("@zml//runtimes/cuda:cudnn.redistrib_{}.json".format(CUDNN_VERSION))))
CUDA_REDIST = _read_redist_json(
mctx,
url = CUDA_REDIST_PREFIX + "redistrib_{}.json".format(CUDA_VERSION),
sha256 = CUDA_REDIST_JSON_SHA256,
)
CUDNN_REDIST = _read_redist_json(
mctx,
url = CUDNN_REDIST_PREFIX + "redistrib_{}.json".format(CUDNN_VERSION),
sha256 = CUDNN_REDIST_JSON_SHA256,
)
for pkg, build_file_content in CUDA_PACKAGES.items():
pkg_data = CUDA_REDIST[pkg]
@ -165,8 +163,8 @@ def _cuda_impl(mctx):
continue
http_archive(
name = pkg,
build_file_content = build_file_content,
url = "https://developer.download.nvidia.com/compute/cuda/redist/" + arch_data["relative_path"],
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + build_file_content,
url = CUDA_REDIST_PREFIX + arch_data["relative_path"],
sha256 = arch_data["sha256"],
strip_prefix = paths.basename(arch_data["relative_path"]).replace(".tar.xz", ""),
)
@ -179,37 +177,28 @@ def _cuda_impl(mctx):
arch_data = arch_data.get("cuda12", arch_data)
http_archive(
name = pkg,
build_file_content = build_file_content,
url = "https://developer.download.nvidia.com/compute/cudnn/redist/" + arch_data["relative_path"],
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + build_file_content,
url = CUDNN_REDIST_PREFIX + arch_data["relative_path"],
sha256 = arch_data["sha256"],
strip_prefix = paths.basename(arch_data["relative_path"]).replace(".tar.xz", ""),
)
http_archive(
name = "nccl",
urls = ["https://files.pythonhosted.org/packages/ed/1f/6482380ec8dcec4894e7503490fc536d846b0d59694acad9cf99f27d0e7d/nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl"],
urls = ["https://files.pythonhosted.org/packages/11/0c/8c78b7603f4e685624a3ea944940f1e75f36d71bd6504330511f4a0e1557/nvidia_nccl_cu12-2.25.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"],
type = "zip",
sha256 = "b097258d9aab2fa9f686e33c6fe40ae57b27df60cedbd15d139701bb5509e0c1",
build_file_content = _cc_import(
sha256 = "362aed5963fb9ea2ed2f264409baae30143498fd0e5c503aeaa1badd88cdc54a",
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + packages.cc_import(
name = "nccl",
shared_library = "nvidia/nccl/lib/libnccl.so.2",
),
)
http_deb_archive(
name = "zlib",
urls = ["https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/pool/main/z/zlib/zlib1g_1.3.dfsg%2Breally1.3.1-1%2Bb1_amd64.deb"],
sha256 = "015be740d6236ad114582dea500c1d907f29e16d6db00566ca32fb68d71ac90d",
build_file_content = _cc_import(
name = "zlib",
shared_library = "usr/lib/x86_64-linux-gnu/libz.so.1",
),
)
http_archive(
name = "libpjrt_cuda",
build_file = "libpjrt_cuda.BUILD.bazel",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v6.0.0/pjrt-cuda_linux-amd64.tar.gz",
sha256 = "7ee04fdacee1d144bc14a47ab0c453e600f6a821dea2ecb1d5af9b7093038dbf",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v7.0.0/pjrt-cuda_linux-amd64.tar.gz",
sha256 = "64029cd3d68118b166198e3246877ed706ed35eb732b1770a9bf530b5b0a8ab4",
)
return mctx.extension_metadata(

2
runtimes/cuda/libpjrt_cuda.BUILD.bazel
View File

@ -45,6 +45,6 @@ cc_import(
"@libcusparse//:cusparse",
"@libnvjitlink//:nvjitlink",
"@nccl",
"@zlib",
"@zlib1g",
],
)

4
runtimes/rocm/gfx.bzl
View File

@ -26,7 +26,7 @@ def _bytecode_select_impl(ctx):
DefaultInfo(
files = depset([
file
for file in ctx.files.bytecodes
for file in ctx.files.srcs
if _is_file_enabled(file, enabled_gfx)
]),
),
@ -35,7 +35,7 @@ def _bytecode_select_impl(ctx):
bytecode_select = rule(
implementation = _bytecode_select_impl,
attrs = {
"bytecodes": attr.label_list(allow_files = True),
"srcs": attr.label_list(allow_files = True),
"enabled_gfx": attr.label(mandatory = True),
},
)

12
runtimes/rocm/libpjrt_rocm.BUILD.bazel
View File

@ -71,12 +71,6 @@ cc_import(
"@hipblaslt",
"@hsa-amd-aqlprofile//:hsa-amd-aqlprofile",
"@hsa-rocr//:hsa-runtime",
"@libdrm-amdgpu",
"@libdrm",
"@libelf",
"@libnuma",
"@libtinfo",
"@libzstd",
"@miopen-hip//:MIOpen",
"@rccl",
"@rocblas",
@ -84,6 +78,12 @@ cc_import(
"@rocm-smi-lib//:rocm_smi",
"@rocprofiler-register",
"@roctracer",
"@libelf1",
"@libdrm2",
"@libnuma1",
"@libzstd1",
"@libdrm-amdgpu1",
"@libtinfo6",
"@zlib1g",
"@zml//runtimes/rocm:zmlxrocm",
],

490
runtimes/rocm/packages.lock.json Executable file
View File

@ -0,0 +1,490 @@
{
"packages": [
{
"arch": "amd64",
"dependencies": [],
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"sha256": "80ec656c41290335df3609043b4ff52df25de82257d987d2e90a4d86cbc0d080",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocm-core/rocm-core_6.3.4.60304-76~22.04_amd64.deb"
],
"version": "6.3.4.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "rocm-smi-lib_7.4.0.60304-76_22.04_amd64",
"name": "rocm-smi-lib",
"sha256": "0aa35daa0b6336f7ff58720067e16e7b68efcf423d8be355e11c3a59e7b7bf71",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocm-smi-lib/rocm-smi-lib_7.4.0.60304-76~22.04_amd64.deb"
],
"version": "7.4.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocprofiler-register_0.4.0.60304-76_22.04_amd64",
"name": "rocprofiler-register",
"version": "0.4.0.60304-76~22.04"
},
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "hsa-rocr_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr",
"sha256": "c78622adcf567f88a22a2a167baf3675d08ebbfb5d98119a17c5a7ae50e14519",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hsa-rocr/hsa-rocr_1.14.0.60304-76~22.04_amd64.deb"
],
"version": "1.14.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "rocprofiler-register_0.4.0.60304-76_22.04_amd64",
"name": "rocprofiler-register",
"sha256": "6d4474295d47b73d88d1e3e67b3daa5a731ffb466de09cfcefa1746d653f738d",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocprofiler-register/rocprofiler-register_0.4.0.60304-76~22.04_amd64.deb"
],
"version": "0.4.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "hsa-amd-aqlprofile_1.0.0.60304-76_22.04_amd64",
"name": "hsa-amd-aqlprofile",
"sha256": "9d6eefd1b5002eb15f0fe0a05d3a9936b216a5d63c45aefa6483e5bd122ec937",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hsa-amd-aqlprofile/hsa-amd-aqlprofile_1.0.0.60304-76~22.04_amd64.deb"
],
"version": "1.0.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "comgr_2.8.0.60304-76_22.04_amd64",
"name": "comgr",
"sha256": "d5473abde7ab958b6245b7eee4dd502e4b23876d7ac63ad0ff7308e21d774809",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/c/comgr/comgr_2.8.0.60304-76~22.04_amd64.deb"
],
"version": "2.8.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
},
{
"key": "rocrand_3.2.0.60304-76_22.04_amd64",
"name": "rocrand",
"version": "3.2.0.60304-76~22.04"
},
{
"key": "hipblaslt_0.10.0.60304-76_22.04_amd64",
"name": "hipblaslt",
"version": "0.10.0.60304-76~22.04"
},
{
"key": "rocblas_4.3.0.60304-76_22.04_amd64",
"name": "rocblas",
"version": "4.3.0.60304-76~22.04"
},
{
"key": "hip-runtime-amd_6.3.42134.60304-76_22.04_amd64",
"name": "hip-runtime-amd",
"version": "6.3.42134.60304-76~22.04"
},
{
"key": "rocprofiler-register_0.4.0.60304-76_22.04_amd64",
"name": "rocprofiler-register",
"version": "0.4.0.60304-76~22.04"
},
{
"key": "comgr_2.8.0.60304-76_22.04_amd64",
"name": "comgr",
"version": "2.8.0.60304-76~22.04"
},
{
"key": "rocminfo_1.0.0.60304-76_22.04_amd64",
"name": "rocminfo",
"version": "1.0.0.60304-76~22.04"
},
{
"key": "hsa-rocr_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr",
"version": "1.14.0.60304-76~22.04"
},
{
"key": "roctracer_4.1.60304.60304-76_22.04_amd64",
"name": "roctracer",
"version": "4.1.60304.60304-76~22.04"
}
],
"key": "miopen-hip_3.3.0.60304-76_22.04_amd64",
"name": "miopen-hip",
"sha256": "4cd1c953cf456b9166f7eb935099c25300aed75b5d51665172d26776e09a0ab0",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/m/miopen-hip/miopen-hip_3.3.0.60304-76~22.04_amd64.deb"
],
"version": "3.3.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [],
"key": "rocrand_3.2.0.60304-76_22.04_amd64",
"name": "rocrand",
"sha256": "8c29882110df41ba7fcfb95509093083b5956ad62c5d38a21dd053277a5d4ada",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocrand/rocrand_3.2.0.60304-76~22.04_amd64.deb"
],
"version": "3.2.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "hipblaslt_0.10.0.60304-76_22.04_amd64",
"name": "hipblaslt",
"sha256": "a7765f361fa078b4cb9847a4f9de7c0e3fe6dd079b61e02d3da8a817fd8e1e85",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hipblaslt/hipblaslt_0.10.0.60304-76~22.04_amd64.deb"
],
"version": "0.10.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
},
{
"key": "hip-runtime-amd_6.3.42134.60304-76_22.04_amd64",
"name": "hip-runtime-amd",
"version": "6.3.42134.60304-76~22.04"
},
{
"key": "rocprofiler-register_0.4.0.60304-76_22.04_amd64",
"name": "rocprofiler-register",
"version": "0.4.0.60304-76~22.04"
},
{
"key": "comgr_2.8.0.60304-76_22.04_amd64",
"name": "comgr",
"version": "2.8.0.60304-76~22.04"
},
{
"key": "rocminfo_1.0.0.60304-76_22.04_amd64",
"name": "rocminfo",
"version": "1.0.0.60304-76~22.04"
},
{
"key": "hsa-rocr_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr",
"version": "1.14.0.60304-76~22.04"
},
{
"key": "hipblaslt_0.10.0.60304-76_22.04_amd64",
"name": "hipblaslt",
"version": "0.10.0.60304-76~22.04"
}
],
"key": "rocblas_4.3.0.60304-76_22.04_amd64",
"name": "rocblas",
"sha256": "8f3ca6d113ffdc7f478326f243f23b9791b38f2b07fac36a329a34e67b6dc73e",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocblas/rocblas_4.3.0.60304-76~22.04_amd64.deb"
],
"version": "4.3.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocprofiler-register_0.4.0.60304-76_22.04_amd64",
"name": "rocprofiler-register",
"version": "0.4.0.60304-76~22.04"
},
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
},
{
"key": "comgr_2.8.0.60304-76_22.04_amd64",
"name": "comgr",
"version": "2.8.0.60304-76~22.04"
},
{
"key": "rocminfo_1.0.0.60304-76_22.04_amd64",
"name": "rocminfo",
"version": "1.0.0.60304-76~22.04"
},
{
"key": "hsa-rocr_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr",
"version": "1.14.0.60304-76~22.04"
}
],
"key": "hip-runtime-amd_6.3.42134.60304-76_22.04_amd64",
"name": "hip-runtime-amd",
"sha256": "d122495575c47125f6a4b1aa53a61d834ce7ab8b429ff86c39751e13197e82d0",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hip-runtime-amd/hip-runtime-amd_6.3.42134.60304-76~22.04_amd64.deb"
],
"version": "6.3.42134.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [],
"key": "rocminfo_1.0.0.60304-76_22.04_amd64",
"name": "rocminfo",
"sha256": "9a7d77959ec3cf4ae3134c2f5341fa0ea37309b083c1ed179aba932566b5a528",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocminfo/rocminfo_1.0.0.60304-76~22.04_amd64.deb"
],
"version": "1.0.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "roctracer_4.1.60304.60304-76_22.04_amd64",
"name": "roctracer",
"sha256": "b8398f10b5f27fe01fce891a64e27963687bb0e63e85bbc11801fe3fcbe30607",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/roctracer/roctracer_4.1.60304.60304-76~22.04_amd64.deb"
],
"version": "4.1.60304.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocprofiler-register_0.4.0.60304-76_22.04_amd64",
"name": "rocprofiler-register",
"version": "0.4.0.60304-76~22.04"
},
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
},
{
"key": "rocm-smi-lib_7.4.0.60304-76_22.04_amd64",
"name": "rocm-smi-lib",
"version": "7.4.0.60304-76~22.04"
},
{
"key": "hip-runtime-amd_6.3.42134.60304-76_22.04_amd64",
"name": "hip-runtime-amd",
"version": "6.3.42134.60304-76~22.04"
},
{
"key": "comgr_2.8.0.60304-76_22.04_amd64",
"name": "comgr",
"version": "2.8.0.60304-76~22.04"
},
{
"key": "rocminfo_1.0.0.60304-76_22.04_amd64",
"name": "rocminfo",
"version": "1.0.0.60304-76~22.04"
},
{
"key": "hsa-rocr_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr",
"version": "1.14.0.60304-76~22.04"
}
],
"key": "rccl_2.21.5.60304-76_22.04_amd64",
"name": "rccl",
"sha256": "ee344582ef81140ed81508897ff1b1d096f7a82a489491d14ac5f4b3ae384bc8",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rccl/rccl_2.21.5.60304-76~22.04_amd64.deb"
],
"version": "2.21.5.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "rocm-device-libs_1.0.0.60304-76_22.04_amd64",
"name": "rocm-device-libs",
"sha256": "060ef802fa7516fb1d50c971ceead8acaa1651a7119b75db8e6454d4a58b67d3",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocm-device-libs/rocm-device-libs_1.0.0.60304-76~22.04_amd64.deb"
],
"version": "1.0.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "hsa-rocr-dev_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr-dev",
"version": "1.14.0.60304-76~22.04"
},
{
"key": "hsa-rocr_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr",
"version": "1.14.0.60304-76~22.04"
},
{
"key": "rocprofiler-register_0.4.0.60304-76_22.04_amd64",
"name": "rocprofiler-register",
"version": "0.4.0.60304-76~22.04"
},
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
},
{
"key": "rocm-llvm_18.0.0.25012.60304-76_22.04_amd64",
"name": "rocm-llvm",
"version": "18.0.0.25012.60304-76~22.04"
},
{
"key": "hip-runtime-amd_6.3.42134.60304-76_22.04_amd64",
"name": "hip-runtime-amd",
"version": "6.3.42134.60304-76~22.04"
},
{
"key": "comgr_2.8.0.60304-76_22.04_amd64",
"name": "comgr",
"version": "2.8.0.60304-76~22.04"
},
{
"key": "rocminfo_1.0.0.60304-76_22.04_amd64",
"name": "rocminfo",
"version": "1.0.0.60304-76~22.04"
}
],
"key": "hip-dev_6.3.42134.60304-76_22.04_amd64",
"name": "hip-dev",
"sha256": "7a1cb7cbfbc316bbb9dccc101a83ca2ebd6f5fff1a1d86d970f806e777624afd",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hip-dev/hip-dev_6.3.42134.60304-76~22.04_amd64.deb"
],
"version": "6.3.42134.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [],
"key": "hsa-rocr-dev_1.14.0.60304-76_22.04_amd64",
"name": "hsa-rocr-dev",
"sha256": "4c6d435b47e87064a4478eb5b6c19de21a4b0e425a51dbe027a3cce2d1bfd048",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hsa-rocr-dev/hsa-rocr-dev_1.14.0.60304-76~22.04_amd64.deb"
],
"version": "1.14.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
}
],
"key": "rocm-llvm_18.0.0.25012.60304-76_22.04_amd64",
"name": "rocm-llvm",
"sha256": "3de2e3b758b1fbe11d83915e0672f152cb0ecfab6b210b714fdbee8856030ba6",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/r/rocm-llvm/rocm-llvm_18.0.0.25012.60304-76~22.04_amd64.deb"
],
"version": "18.0.0.25012.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [
{
"key": "hipblaslt_0.10.0.60304-76_22.04_amd64",
"name": "hipblaslt",
"version": "0.10.0.60304-76~22.04"
},
{
"key": "rocm-core_6.3.4.60304-76_22.04_amd64",
"name": "rocm-core",
"version": "6.3.4.60304-76~22.04"
},
{
"key": "hipblas-common-dev_1.0.0.60304-76_22.04_amd64",
"name": "hipblas-common-dev",
"version": "1.0.0.60304-76~22.04"
}
],
"key": "hipblaslt-dev_0.10.0.60304-76_22.04_amd64",
"name": "hipblaslt-dev",
"sha256": "5bbd7d19a05b389400732ecec5cc16df25a6c536c54de1d554ea1b74a64772ff",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hipblaslt-dev/hipblaslt-dev_0.10.0.60304-76~22.04_amd64.deb"
],
"version": "0.10.0.60304-76~22.04"
},
{
"arch": "amd64",
"dependencies": [],
"key": "hipblas-common-dev_1.0.0.60304-76_22.04_amd64",
"name": "hipblas-common-dev",
"sha256": "7a03e99f09cefaea9b13a59fc41def4214fff9a90cf05ab19be0f83ef9653b5e",
"urls": [
"https://repo.radeon.com/rocm/apt/6.3.4/pool/main/h/hipblas-common-dev/hipblas-common-dev_1.0.0.60304-76~22.04_amd64.deb"
],
"version": "1.0.0.60304-76~22.04"
}
],
"version": 1
}

38
runtimes/rocm/packages.yaml Normal file
View File

@ -0,0 +1,38 @@
#
# bazel run @rocm_apt//:lock
#
version: 1
sources:
- channel: jammy main
url: https://repo.radeon.com/rocm/apt/6.3.4/
# - channel: bookworm main
# url: https://snapshot-cloudflare.debian.org/archive/debian/20241127T143620Z/
archs:
- "amd64"
packages:
# - "libdrm2"
# - "libelf1"
# - "libnuma1"
# - "libzstd1"
# - "libdrm-amdgpu1"
# - "libtinfo6"
# - "zlib1g"
- "rocm-core"
- "rocm-smi-lib"
- "hsa-rocr"
- "hsa-amd-aqlprofile"
- "comgr"
- "rocprofiler-register"
- "miopen-hip"
- "rccl"
- "rocm-device-libs"
- "hip-dev"
- "rocblas"
- "roctracer"
- "hipblaslt"
- "hipblaslt-dev"
- "hip-runtime-amd"
- "rocm-llvm"

313
runtimes/rocm/rocm.bzl
View File

@ -1,219 +1,134 @@
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("//bazel:dpkg.bzl", "dpkg")
load("//bazel:http_deb_archive.bzl", "http_deb_archive")
load("//runtimes/common:packages.bzl", "packages")
ROCM_VERSION = "6.3.1"
BASE_URL = "https://repo.radeon.com/rocm/apt/{}".format(ROCM_VERSION)
STRIP_PREFIX = "opt/rocm-6.3.1"
def pkg_kwargs(pkg):
return dict(
name = pkg.Package,
urls = [BASE_URL + "/" + pkg.Filename],
sha256 = pkg.SHA256,
strip_prefix = STRIP_PREFIX,
)
def _ubuntu_package(path, deb_path, sha256, name, shared_library):
return {
"urls": ["http://archive.ubuntu.com/ubuntu/pool/main/{}".format(path)],
"sha256": sha256,
"build_file_content": """\
cc_import(
name = {name},
shared_library = "{deb_path}{shared_library}",
visibility = ["//visibility:public"],
)
""".format(name = repr(name), shared_library = shared_library, deb_path = deb_path),
}
_UBUNTU_PACKAGES = {
"libdrm": _ubuntu_package(
path = "libd/libdrm/libdrm2_2.4.107-8ubuntu1~20.04.2_amd64.deb",
deb_path = "usr/lib/x86_64-linux-gnu/",
sha256 = "9b01d73313841abe8e3f24c2715edced675fbe329bbd10be912a5b135cd51fb6",
name = "libdrm",
shared_library = "libdrm.so.2",
),
"libelf": _ubuntu_package(
path = "e/elfutils/libelf1_0.176-1.1build1_amd64.deb",
deb_path = "usr/lib/x86_64-linux-gnu/",
sha256 = "78a8761227efc04a1e37527f2f33ba608c6fb5d6c911616346ada5d7b9b72ee3",
name = "libelf",
shared_library = "libelf.so.1",
),
"libnuma": _ubuntu_package(
path = "n/numactl/libnuma1_2.0.12-1_amd64.deb",
deb_path = "usr/lib/x86_64-linux-gnu/",
sha256 = "0b1edf08cf9befecd21fe94e298ac25e476f87fd876ddd4adf42ef713449e637",
name = "libnuma",
shared_library = "libnuma.so.1",
),
"libzstd": _ubuntu_package(
path = "libz/libzstd/libzstd1_1.4.4+dfsg-3ubuntu0.1_amd64.deb",
deb_path = "usr/lib/x86_64-linux-gnu/",
sha256 = "7a4422dadb90510dc90765c308d65e61a3e244ceb3886394335e48cff7559e69",
name = "libzstd",
shared_library = "libzstd.so.1",
),
"libdrm-amdgpu": _ubuntu_package(
path = "libd/libdrm/libdrm-amdgpu1_2.4.107-8ubuntu1~20.04.2_amd64.deb",
deb_path = "usr/lib/x86_64-linux-gnu/",
sha256 = "0d95779b581f344e3d658e0f21f6e4b57da6eb3606c0bcb8cb874c12f5754bf2",
name = "libdrm-amdgpu",
shared_library = "libdrm_amdgpu.so.1",
),
"libtinfo": _ubuntu_package(
path = "n/ncurses/libtinfo6_6.2-0ubuntu2.1_amd64.deb",
deb_path = "lib/x86_64-linux-gnu/",
sha256 = "711a3a901c3a71561565558865699efa9c07a99fdc810ffe086a5636f89c6431",
name = "libtinfo",
shared_library = "libtinfo.so.6",
),
"zlib1g": _ubuntu_package(
path = "z/zlib/zlib1g_1.2.11.dfsg-2ubuntu1.5_amd64.deb",
deb_path = "lib/x86_64-linux-gnu/",
sha256 = "bf67018f5303466eb468680b637a5d3f3bb17b9d44decf3d82d40b35babcd3e0",
name = "zlib1g",
shared_library = "libz.so.1",
),
}
_CC_IMPORT_TPL = """\
cc_import(
name = "{name}",
shared_library = "lib/{shared_library}",
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
_BUILD_FILE_DEFAULT_VISIBILITY = """\
package(default_visibility = ["//visibility:public"])
"""
_RUNFILES_TPL = """\
filegroup(
name = "{name}",
srcs = glob({glob}),
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
"""
_ROCM_STRIP_PREFIX = "opt/rocm-6.3.4"
_PACKAGES = {
"rocm-core": _CC_IMPORT_TPL.format(name = "rocm-core", shared_library = "librocm-core.so.1"),
"rocm-smi-lib": _CC_IMPORT_TPL.format(name = "rocm_smi", shared_library = "librocm_smi64.so.7"),
"hsa-rocr": _CC_IMPORT_TPL.format(name = "hsa-runtime", shared_library = "libhsa-runtime64.so.1"),
"hsa-amd-aqlprofile": _CC_IMPORT_TPL.format(name = "hsa-amd-aqlprofile", shared_library = "libhsa-amd-aqlprofile64.so.1"),
"comgr": _CC_IMPORT_TPL.format(name = "amd_comgr", shared_library = "libamd_comgr.so.2"),
"rocprofiler-register": _CC_IMPORT_TPL.format(name = "rocprofiler-register", shared_library = "librocprofiler-register.so.0"),
"miopen-hip": "".join([
_CC_IMPORT_TPL.format(name = "MIOpen", shared_library = "libMIOpen.so.1"),
_RUNFILES_TPL.format(name = "runfiles", glob = repr(["share/miopen/**"])),
# def _kwargs(**kwargs):
# return repr(struct(**kwargs))[len("struct("):-1]
# def packages.cc_import(**kwargs):
# return """cc_import({})""".format(_kwargs(**kwargs))
# def packages.filegroup(**kwargs):
# return """filegroup({})""".format(_kwargs(**kwargs))
# def packages.load_(bzl, name):
# return """load({}, {})""".format(repr(bzl), repr(name))
# _UBUNTU_PACKAGES = {
# "libdrm2": packages.cc_import(name = "libdrm2", shared_library = "usr/lib/x86_64-linux-gnu/libdrm.so.2"),
# "libelf1": packages.cc_import(name = "libelf1", shared_library = "usr/lib/x86_64-linux-gnu/libelf.so.1"),
# "libnuma1": packages.cc_import(name = "libnuma1", shared_library = "usr/lib/x86_64-linux-gnu/libnuma.so.1"),
# "libzstd1": packages.cc_import(name = "libzstd1", shared_library = "usr/lib/x86_64-linux-gnu/libzstd.so.1"),
# "libdrm-amdgpu1": packages.cc_import(name = "libdrm-amdgpu1", shared_library = "usr/lib/x86_64-linux-gnu/libdrm_amdgpu.so.1"),
# "libtinfo6": packages.cc_import(name = "libtinfo6", shared_library = "lib/x86_64-linux-gnu/libtinfo.so.6"),
# "zlib1g": packages.cc_import(name = "zlib1g", shared_library = "lib/x86_64-linux-gnu/libz.so.1"),
# }
_ROCM_PACKAGES = {
"rocm-core": packages.cc_import(name = "rocm-core", shared_library = "lib/librocm-core.so.1"),
"rocm-smi-lib": packages.cc_import(name = "rocm_smi", shared_library = "lib/librocm_smi64.so.7"),
"hsa-rocr": packages.cc_import(name = "hsa-runtime", shared_library = "lib/libhsa-runtime64.so.1"),
"hsa-amd-aqlprofile": packages.cc_import(name = "hsa-amd-aqlprofile", shared_library = "lib/libhsa-amd-aqlprofile64.so.1"),
"comgr": packages.cc_import(name = "amd_comgr", shared_library = "lib/libamd_comgr.so.2"),
"rocprofiler-register": packages.cc_import(name = "rocprofiler-register", shared_library = "lib/librocprofiler-register.so.0"),
"miopen-hip": "\n".join([
packages.cc_import(name = "MIOpen", shared_library = "lib/libMIOpen.so.1"),
"""filegroup(name = "runfiles", srcs = glob(["share/miopen/**"]))""",
]),
"rccl": _CC_IMPORT_TPL.format(name = "rccl", shared_library = "librccl.so.1"),
"rocm-device-libs": _RUNFILES_TPL.format(name = "runfiles", glob = repr(["amdgcn/**"])),
"hip-dev": _RUNFILES_TPL.format(name = "runfiles", glob = repr(["share/**"])),
"rocblas": """\
load("@zml//bazel:cc_import.bzl", "cc_import")
load("@zml//runtimes/rocm:gfx.bzl", "bytecode_select")
cc_import(
name = "rocblas",
shared_library = "lib/librocblas.so.4",
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
bytecode_select(
name = "bytecodes",
bytecodes = glob(["lib/rocblas/library/*"]),
enabled_gfx = "@libpjrt_rocm//:gfx",
)
filegroup(
name = "runfiles",
srcs = [
":bytecodes",
"lib/rocblas/library/TensileManifest.txt",
],
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
""",
"roctracer": """\
cc_import(
name = "roctracer",
shared_library = "lib/libroctracer64.so.4",
visibility = ["@libpjrt_rocm//:__subpackages__"],
deps = [":roctx"],
)
cc_import(
name = "roctx",
shared_library = "lib/libroctx64.so.4",
)
""",
"hipblaslt": """\
load("@zml//bazel:cc_import.bzl", "cc_import")
load("@zml//runtimes/rocm:gfx.bzl", "bytecode_select")
cc_import(
name = "hipblaslt",
shared_library = "lib/libhipblaslt.so.0",
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
bytecode_select(
name = "bytecodes",
bytecodes = glob(
include = ["lib/hipblaslt/library/*"],
exclude = ["lib/hipblaslt/library/hipblasltExtOpLibrary.dat"],
),
enabled_gfx = "@libpjrt_rocm//:gfx",
)
filegroup(
name = "runfiles",
srcs = [
":bytecodes",
"lib/hipblaslt/library/hipblasltExtOpLibrary.dat",
"lib/hipblaslt/library/TensileManifest.txt",
],
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
""",
"hip-runtime-amd": """\
cc_import(
name = "amdhip",
shared_library = "lib/libamdhip64.so.6",
visibility = ["@libpjrt_rocm//:__subpackages__"],
deps = [":hiprtc"],
)
cc_import(
name = "hiprtc",
shared_library = "lib/libhiprtc.so.6",
)
""",
"rocm-llvm": """\
filegroup(
name = "lld",
srcs = ["llvm/bin/ld.lld"],
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
""",
"rccl": packages.cc_import(name = "rccl", shared_library = "lib/librccl.so.1"),
"rocm-device-libs": """filegroup(name = "runfiles", srcs = glob(["amdgcn/**"]))""",
"hip-dev": """filegroup(name = "runfiles", srcs = glob(["share/**"]))""",
"rocblas": "\n".join([
packages.load_("@zml//bazel:cc_import.bzl", "cc_import"),
packages.load_("@zml//runtimes/rocm:gfx.bzl", "bytecode_select"),
packages.cc_import(
name = "rocblas",
shared_library = "lib/librocblas.so.4",
add_needed = ["libzmlxrocm.so.0"],
rename_dynamic_symbols = {
"dlopen": "zmlxrocm_dlopen",
},
),
"""bytecode_select(
name = "bytecodes",
srcs = glob(["lib/rocblas/library/*"]),
enabled_gfx = "@libpjrt_rocm//:gfx",
)
""",
packages.filegroup(
name = "runfiles",
srcs = [
"lib/rocblas/library/TensileManifest.txt",
":bytecodes",
],
),
]),
"roctracer": "\n".join([
packages.cc_import(name = "roctracer", shared_library = "lib/libroctracer64.so.4", deps = [":roctx"]),
packages.cc_import(name = "roctx", shared_library = "lib/libroctx64.so.4"),
]),
"hipblaslt": "\n".join([
packages.load_("@zml//bazel:cc_import.bzl", "cc_import"),
packages.load_("@zml//runtimes/rocm:gfx.bzl", "bytecode_select"),
packages.cc_import(
name = "hipblaslt",
shared_library = "lib/libhipblaslt.so.0",
add_needed = ["libzmlxrocm.so.0"],
rename_dynamic_symbols = {
"dlopen": "zmlxrocm_dlopen",
},
),
"""bytecode_select(
name = "bytecodes",
srcs = glob(
include = ["lib/hipblaslt/library/*"],
exclude = ["lib/hipblaslt/library/hipblasltExtOpLibrary.dat"],
),
enabled_gfx = "@libpjrt_rocm//:gfx",
)
""",
packages.filegroup(
name = "runfiles",
srcs = [
"lib/hipblaslt/library/hipblasltExtOpLibrary.dat",
"lib/hipblaslt/library/TensileManifest.txt",
":bytecodes",
],
),
]),
"hip-runtime-amd": "\n".join([
packages.cc_import(name = "amdhip", shared_library = "lib/libamdhip64.so.6", deps = [":hiprtc"]),
packages.cc_import(name = "hiprtc", shared_library = "lib/libhiprtc.so.6"),
]),
"rocm-llvm": packages.filegroup(name = "lld", srcs = ["llvm/bin/ld.lld"], visibility = ["//visibility:public"]),
}
def _rocm_impl(mctx):
all_packages = dpkg.read_packages(mctx, "@zml//runtimes/rocm:packages.amd64.txt")
loaded_packages = packages.read(mctx, [
"@zml//runtimes/rocm:packages.lock.json",
])
for pkg_name, build_file_content in _PACKAGES.items():
pkg = all_packages[pkg_name].values()[0]
for pkg_name, build_file_content in _ROCM_PACKAGES.items():
pkg = loaded_packages[pkg_name]
http_deb_archive(
build_file_content = build_file_content,
**pkg_kwargs(pkg)
name = pkg_name,
urls = pkg["urls"],
sha256 = pkg["sha256"],
strip_prefix = _ROCM_STRIP_PREFIX,
build_file_content = _BUILD_FILE_DEFAULT_VISIBILITY + build_file_content,
)
for repository, kwargs in _UBUNTU_PACKAGES.items():
http_deb_archive(name = repository, **kwargs)
http_archive(
name = "libpjrt_rocm",
build_file = "libpjrt_rocm.BUILD.bazel",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v6.0.0/pjrt-rocm_linux-amd64.tar.gz",
sha256 = "0a41e5b255d0b3284e029c11c2362d777557f87778e1531a2117bc46b835395a",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v7.0.0/pjrt-rocm_linux-amd64.tar.gz",
sha256 = "13a2ced965a44a0e8ad0d752c8ac5aa99107b17b41bd850967d40b82e102ec50",
)
return mctx.extension_metadata(