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

workspace: Update LLVM, XLA, StableHLO, and PJRT plugins to latest versions.

Browse Source
This commit is contained in:
Tarry Singh 2023-12-13 10:10:32 +00:00
parent 6a4a7fb9a1
commit 145e60b4dd
24 changed files with 4370 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
MODULE.bazel
View File

@ -85,7 +85,7 @@ use_repo(zls, "zls_aarch64-macos", "zls_x86_64-linux")
register_toolchains("//third_party/zls:all")
bazel_dep(name = "libxev", version = "20241119.0-6afcde9")
bazel_dep(name = "llvm-raw", version = "20241022.0-6c4267f")
bazel_dep(name = "llvm-raw", version = "20250102.0-f739aa4")
llvm = use_extension("@llvm-raw//utils/bazel:extension.bzl", "llvm")
llvm.configure(
@ -97,8 +97,8 @@ llvm.configure(
)
use_repo(llvm, "llvm-project")
bazel_dep(name = "stablehlo", version = "20241021.0-1c0b606")
bazel_dep(name = "xla", version = "20241025.0-4663f04")
bazel_dep(name = "stablehlo", version = "20241220.0-38bb2f9")
bazel_dep(name = "xla", version = "20250103.0-5f1fe6a")
tsl = use_extension("@xla//:tsl.bzl", "tsl")
use_repo(tsl, "tsl")

4076
MODULE.bazel.lock
View File

File diff suppressed because one or more lines are too long

8
runtimes/cpu/cpu.bzl
View File

@ -12,15 +12,15 @@ def _cpu_pjrt_plugin_impl(mctx):
http_archive(
name = "libpjrt_cpu_linux_amd64",
build_file_content = _BUILD.format(ext = "so"),
sha256 = "646b8ea61e690af0e4133637343674fb072e7d5e3a29694e6f84bb66ea75a6f0",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v3.0.0/pjrt-cpu_linux-amd64.tar.gz",
sha256 = "35b6aefa0359317ae2892f846d6da892bee2116d8c6722e397ef0120cf572183",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v4.0.0/pjrt-cpu_linux-amd64.tar.gz",
)
http_archive(
name = "libpjrt_cpu_darwin_arm64",
build_file_content = _BUILD.format(ext = "dylib"),
sha256 = "f166ee5ba1d50383731aa79831d4bd2ef3338c5948ae92c2442105d20280506c",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v3.0.0/pjrt-cpu_darwin-arm64.tar.gz",
sha256 = "a532a2e1511f91ec6d6adc60290f6bc4d88d2521508661e90b9824061ebabb3a",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v4.0.0/pjrt-cpu_darwin-arm64.tar.gz",
)
return mctx.extension_metadata(

4
runtimes/cuda/cuda.bzl
View File

@ -190,9 +190,9 @@ def _cuda_impl(mctx):
http_archive(
name = "libpjrt_cuda",
build_file = "libpjrt_cuda.BUILD.bazel",
url = "https://files.pythonhosted.org/packages/d7/aa/f15ea857ad9bcff7a0c942dc570ca718b026cc0cc5c513525bb08cacf3c0/jax_cuda12_pjrt-0.4.35-py3-none-manylinux2014_x86_64.whl",
url = "https://files.pythonhosted.org/packages/90/43/ac2c369e202e3e3e7e5aa7929b197801ba02eaf11868437adaa5341704e4/jax_cuda12_pjrt-0.4.38-py3-none-manylinux2014_x86_64.whl",
type = "zip",
sha256 = "0ffe7e1ba65659bd5738c2cc5addaf0a56205d2188eec5da194b63c068e1fdd2",
sha256 = "83be4c59fbcf30077a60085d98e7d59dc738b1c91e0d628e4ac1779fde15ac2b",
)
return mctx.extension_metadata(

4
runtimes/rocm/rocm.bzl
View File

@ -216,8 +216,8 @@ def _rocm_impl(mctx):
http_archive(
name = "libpjrt_rocm",
build_file = "libpjrt_rocm.BUILD.bazel",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v3.0.0/pjrt-rocm_linux-amd64.tar.gz",
sha256 = "a7da45dfca820d3defa6de8e782cc334a3f6bdffe65fa972c048994923c2e110",
url = "https://github.com/zml/pjrt-artifacts/releases/download/v4.0.0/pjrt-rocm_linux-amd64.tar.gz",
sha256 = "75c2baf2efba0b2c6fe2513d06e542ed3f3a966e43498cc1d932465f646ca34d",
)
return mctx.extension_metadata(

4
runtimes/tpu/tpu.bzl
View File

@ -4,9 +4,9 @@ def _tpu_impl(mctx):
# https://storage.googleapis.com/jax-releases/libtpu_releases.html
http_archive(
name = "libpjrt_tpu",
url = "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20240915+nightly-py3-none-any.whl",
url = "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20250102+nightly-py3-none-linux_x86_64.whl",
type = "zip",
sha256 = "4eee6c7dc92e7e60334c5b0261308a0a07c2f5b6235c7c60f465263576c602bf",
sha256 = "df8339b4f852bd19ad4ed380facc08f28c04e214e9dabb88863e70907b08817e",
build_file = "libpjrt_tpu.BUILD.bazel",
)
return mctx.extension_metadata(

11
third_party/modules/llvm-raw/20250102.0-f739aa4/MODULE.bazel vendored Normal file
View File

@ -0,0 +1,11 @@
module(
name = "llvm-raw",
version = "20250102.0-f739aa4",
compatibility_level = 1,
)
bazel_dep(name = "bazel_skylib", version = "1.7.1")
bazel_dep(name = "platforms", version = "0.0.10")
bazel_dep(name = "zstd", version = "1.5.6", repo_name = "llvm_zstd")
bazel_dep(name = "zlib", version = "1.3.1.bcr.3", repo_name = "llvm_zlib")
bazel_dep(name = "rules_python", version = "0.29.0")

0
third_party/modules/llvm-raw/20250102.0-f739aa4/overlay/BUILD.bazel vendored Normal file
View File

11
third_party/modules/llvm-raw/20250102.0-f739aa4/overlay/MODULE.bazel vendored Normal file
View File

@ -0,0 +1,11 @@
module(
name = "llvm-raw",
version = "20250102.0-f739aa4",
compatibility_level = 1,
)
bazel_dep(name = "bazel_skylib", version = "1.7.1")
bazel_dep(name = "platforms", version = "0.0.10")
bazel_dep(name = "zstd", version = "1.5.6", repo_name = "llvm_zstd")
bazel_dep(name = "zlib", version = "1.3.1.bcr.3", repo_name = "llvm_zlib")
bazel_dep(name = "rules_python", version = "0.29.0")

28
third_party/modules/llvm-raw/20250102.0-f739aa4/overlay/utils/bazel/extension.bzl vendored Normal file
View File

@ -0,0 +1,28 @@
load("//utils/bazel:configure.bzl", _llvm_configure = "llvm_configure")
def _llvm_impl(mctx):
_targets = {}
for mod in mctx.modules:
for conf in mod.tags.configure:
for target in conf.targets:
_targets[target] = True
_llvm_configure(
name = "llvm-project",
targets = _targets.keys(),
)
return mctx.extension_metadata(
reproducible = True,
root_module_direct_deps = "all",
root_module_direct_dev_deps = [],
)
llvm = module_extension(
implementation = _llvm_impl,
tag_classes = {
"configure": tag_class(
attrs = {
"targets": attr.string_list(mandatory = True),
},
),
},
)

10
third_party/modules/llvm-raw/20250102.0-f739aa4/source.json vendored Normal file
View File

@ -0,0 +1,10 @@
{
"strip_prefix": "llvm-project-f739aa4004165dc64d3a1f418d5ad3c84886f01a",
"url": "https://github.com/llvm/llvm-project/archive/f739aa4004165dc64d3a1f418d5ad3c84886f01a.tar.gz",
"integrity": "sha256-hdoTTnugRO9Q68AJ0aV6h/sOLbIIoEZQ7y/lZOlWSqc=",
"overlay": {
"BUILD.bazel": "",
"MODULE.bazel": "",
"utils/bazel/extension.bzl": ""
}
}

1
third_party/modules/llvm-raw/metadata.json vendored
View File

@ -14,6 +14,7 @@
"20240823.0-f142f8a",
"20240919.0-94c024a",
"20241022.0-6c4267f",
"20250102.0-f739aa4",
],
"yanked_versions": {}
}

15
third_party/modules/stablehlo/20241220.0-38bb2f9/MODULE.bazel vendored Normal file
View File

@ -0,0 +1,15 @@
module(
name = "stablehlo",
version = "20241220.0-38bb2f9",
compatibility_level = 1,
)
bazel_dep(name = "bazel_skylib", version = "1.7.1")
bazel_dep(name = "rules_cc", version = "0.0.9")
bazel_dep(name = "llvm-raw", version = "20250102.0-f739aa4")
llvm = use_extension("@llvm-raw//utils/bazel:extension.bzl", "llvm")
llvm.configure(
targets = ["AArch64", "X86", "NVPTX"],
)
use_repo(llvm, "llvm-project")

15
third_party/modules/stablehlo/20241220.0-38bb2f9/overlay/MODULE.bazel vendored Normal file
View File

@ -0,0 +1,15 @@
module(
name = "stablehlo",
version = "20241220.0-38bb2f9",
compatibility_level = 1,
)
bazel_dep(name = "bazel_skylib", version = "1.7.1")
bazel_dep(name = "rules_cc", version = "0.0.9")
bazel_dep(name = "llvm-raw", version = "20250102.0-f739aa4")
llvm = use_extension("@llvm-raw//utils/bazel:extension.bzl", "llvm")
llvm.configure(
targets = ["AArch64", "X86", "NVPTX"],
)
use_repo(llvm, "llvm-project")

8
third_party/modules/stablehlo/20241220.0-38bb2f9/source.json vendored Normal file
View File

@ -0,0 +1,8 @@
{
"strip_prefix": "stablehlo-38bb2f9bf63b714e8a49fe34a478139058ee1660",
"url": "https://github.com/openxla/stablehlo/archive/38bb2f9bf63b714e8a49fe34a478139058ee1660.tar.gz",
"integrity": "sha256-TuaG8q6WCdJGl4ECpu8YO8ORYcRo0NSff3xHKVj8Ixw=",
"overlay": {
"MODULE.bazel": ""
}
}

1
third_party/modules/stablehlo/metadata.json vendored
View File

@ -14,6 +14,7 @@
"20240829.0-54aa1a5",
"20240917.0-78c753a",
"20241021.0-1c0b606",
"20241220.0-38bb2f9",
],
"yanked_versions": {}
}

34
third_party/modules/xla/20250103.0-5f1fe6a/MODULE.bazel vendored Normal file
View File

@ -0,0 +1,34 @@
module(
name = "xla",
version = "20250103.0-5f1fe6a",
compatibility_level = 1,
)
bazel_dep(name = "platforms", version = "0.0.8")
bazel_dep(name = "bazel_skylib", version = "1.5.0")
bazel_dep(name = "rules_cc", version = "0.0.9")
bazel_dep(name = "rules_apple", version = "3.2.1", repo_name = "build_bazel_rules_apple")
bazel_dep(name = "abseil-cpp", version = "20240116.0", repo_name = "com_google_absl")
bazel_dep(name = "rules_python", version = "0.29.0")
bazel_dep(name = "rules_proto", version = "6.0.0-rc1")
bazel_dep(name = "rules_java", version = "7.3.2")
bazel_dep(name = "rules_pkg", version = "0.9.1")
bazel_dep(name = "zlib", version = "1.2.13")
bazel_dep(name = "re2", version = "2024-02-01", repo_name = "com_googlesource_code_re2")
bazel_dep(name = "rules_license", version = "0.0.8")
bazel_dep(name = "stablehlo", version = "20241220.0-38bb2f9")
tsl = use_extension("//:tsl.bzl", "tsl")
use_repo(tsl, "tsl")
xla_workspace = use_extension("//:workspace.bzl", "xla_workspace")
use_repo(
xla_workspace,
"com_github_grpc_grpc",
"com_google_protobuf",
"local_config_cuda",
"local_config_remote_execution",
"local_config_rocm",
"local_config_tensorrt",
)

34
third_party/modules/xla/20250103.0-5f1fe6a/overlay/MODULE.bazel vendored Normal file
View File

@ -0,0 +1,34 @@
module(
name = "xla",
version = "20250103.0-5f1fe6a",
compatibility_level = 1,
)
bazel_dep(name = "platforms", version = "0.0.8")
bazel_dep(name = "bazel_skylib", version = "1.5.0")
bazel_dep(name = "rules_cc", version = "0.0.9")
bazel_dep(name = "rules_apple", version = "3.2.1", repo_name = "build_bazel_rules_apple")
bazel_dep(name = "abseil-cpp", version = "20240116.0", repo_name = "com_google_absl")
bazel_dep(name = "rules_python", version = "0.29.0")
bazel_dep(name = "rules_proto", version = "6.0.0-rc1")
bazel_dep(name = "rules_java", version = "7.3.2")
bazel_dep(name = "rules_pkg", version = "0.9.1")
bazel_dep(name = "zlib", version = "1.2.13")
bazel_dep(name = "re2", version = "2024-02-01", repo_name = "com_googlesource_code_re2")
bazel_dep(name = "rules_license", version = "0.0.8")
bazel_dep(name = "stablehlo", version = "20241220.0-38bb2f9")
tsl = use_extension("//:tsl.bzl", "tsl")
use_repo(tsl, "tsl")
xla_workspace = use_extension("//:workspace.bzl", "xla_workspace")
use_repo(
xla_workspace,
"com_github_grpc_grpc",
"com_google_protobuf",
"local_config_cuda",
"local_config_remote_execution",
"local_config_rocm",
"local_config_tensorrt",
)

19
third_party/modules/xla/20250103.0-5f1fe6a/overlay/tsl.bzl vendored Normal file
View File

@ -0,0 +1,19 @@
load("//third_party:repo.bzl", "tf_vendored")
load("//third_party/py:python_init_repositories.bzl", "python_init_repositories")
def _tsl_impl(mctx):
python_init_repositories(
requirements = {
"3.11": "//:requirements_lock_3_11.txt",
},
)
tf_vendored(name = "tsl", relpath = "third_party/tsl")
return mctx.extension_metadata(
reproducible = True,
root_module_direct_deps = ["tsl"],
root_module_direct_dev_deps = [],
)
tsl = module_extension(
implementation = _tsl_impl,
)

52
third_party/modules/xla/20250103.0-5f1fe6a/overlay/workspace.bzl vendored Normal file
View File

@ -0,0 +1,52 @@
load("@tsl//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
load("@tsl//third_party/gpus:cuda_configure.bzl", "cuda_configure")
load("@tsl//third_party/gpus:rocm_configure.bzl", "rocm_configure")
load("@tsl//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure")
load("@tsl//tools/toolchains/remote:configure.bzl", "remote_execution_configure")
def _xla_workspace_impl(mctx):
cuda_configure(name = "local_config_cuda")
remote_execution_configure(name = "local_config_remote_execution")
rocm_configure(name = "local_config_rocm")
tensorrt_configure(name = "local_config_tensorrt")
tf_http_archive(
name = "com_github_grpc_grpc",
sha256 = "b956598d8cbe168b5ee717b5dafa56563eb5201a947856a6688bbeac9cac4e1f",
strip_prefix = "grpc-b54a5b338637f92bfcf4b0bc05e0f57a5fd8fadd",
system_build_file = "@tsl//third_party/systemlibs:grpc.BUILD",
patch_file = [
"@tsl//third_party/grpc:generate_cc_env_fix.patch",
"@tsl//third_party/grpc:register_go_toolchain.patch",
],
system_link_files = {
"@tsl//third_party/systemlibs:BUILD": "bazel/BUILD",
"@tsl//third_party/systemlibs:grpc.BUILD": "src/compiler/BUILD",
"@tsl//third_party/systemlibs:grpc.bazel.grpc_deps.bzl": "bazel/grpc_deps.bzl",
"@tsl//third_party/systemlibs:grpc.bazel.grpc_extra_deps.bzl": "bazel/grpc_extra_deps.bzl",
"@tsl//third_party/systemlibs:grpc.bazel.cc_grpc_library.bzl": "bazel/cc_grpc_library.bzl",
"@tsl//third_party/systemlibs:grpc.bazel.generate_cc.bzl": "bazel/generate_cc.bzl",
"@tsl//third_party/systemlibs:grpc.bazel.protobuf.bzl": "bazel/protobuf.bzl",
},
urls = tf_mirror_urls("https://github.com/grpc/grpc/archive/b54a5b338637f92bfcf4b0bc05e0f57a5fd8fadd.tar.gz"),
)
tf_http_archive(
name = "com_google_protobuf",
patch_file = ["@tsl//third_party/protobuf:protobuf.patch"],
sha256 = "f66073dee0bc159157b0bd7f502d7d1ee0bc76b3c1eac9836927511bdc4b3fc1",
strip_prefix = "protobuf-3.21.9",
system_build_file = "@tsl//third_party/systemlibs:protobuf.BUILD",
system_link_files = {
"@tsl//third_party/systemlibs:protobuf.bzl": "protobuf.bzl",
"@tsl//third_party/systemlibs:protobuf_deps.bzl": "protobuf_deps.bzl",
},
urls = tf_mirror_urls("https://github.com/protocolbuffers/protobuf/archive/v3.21.9.zip"),
)
return mctx.extension_metadata(
reproducible = True,
root_module_direct_deps = "all",
root_module_direct_dev_deps = [],
)
xla_workspace = module_extension(
implementation = _xla_workspace_impl,
)

27
third_party/modules/xla/20250103.0-5f1fe6a/patches/0003-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extensio.patch vendored Normal file
View File

@ -0,0 +1,27 @@
From 4db5de34f70d991fedbe28915c8239b97ba7a064 Mon Sep 17 00:00:00 2001
From: Steeve Morin <steeve.morin@gmail.com>
Date: Mon, 18 Mar 2024 17:17:34 +0100
Subject: [PATCH 3/3] [PJRT C API] Ensure C compliance for Profiler Extension
---
xla/pjrt/c/pjrt_c_api_profiler_extension.h | 2 ++
1 file changed, 2 insertions(+)
diff --git a/xla/pjrt/c/pjrt_c_api_profiler_extension.h b/xla/pjrt/c/pjrt_c_api_profiler_extension.h
index c821916ad..89a596123 100644
--- a/xla/pjrt/c/pjrt_c_api_profiler_extension.h
+++ b/xla/pjrt/c/pjrt_c_api_profiler_extension.h
@@ -16,8 +16,10 @@ limitations under the License.
#ifndef XLA_PJRT_C_PJRT_C_API_PROFILER_EXTENSION_H_
#define XLA_PJRT_C_PJRT_C_API_PROFILER_EXTENSION_H_
+#ifdef __cplusplus
#include <cstddef>
#include <cstdint>
+#endif
#include "xla/backends/profiler/plugin/profiler_c_api.h"
#include "xla/pjrt/c/pjrt_c_api.h"
--
2.39.3 (Apple Git-146)

14
third_party/modules/xla/20250103.0-5f1fe6a/source.json vendored Normal file
View File

@ -0,0 +1,14 @@
{
"strip_prefix": "xla-5f1fe6a2f35756238e4a0a0c4e754d6454296efe",
"url": "https://github.com/openxla/xla/archive/5f1fe6a2f35756238e4a0a0c4e754d6454296efe.tar.gz",
"integrity": "sha256-0+kdYIla4F63oPfACtw6LmmMLdKTFlBBrgvc9PmxFoI=",
"overlay": {
"tsl.bzl": "",
"workspace.bzl": "",
"MODULE.bazel": ""
},
"patch_strip": 1,
"patches": {
"0003-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extensio.patch": ""
}
}

1
third_party/modules/xla/metadata.json vendored
View File

@ -14,6 +14,7 @@
"20240902.0-d18cd64",
"20240919.0-1b18dd6",
"20241025.0-4663f04",
"20250103.0-5f1fe6a",
],
"yanked_versions": {}
}

1
zml/testing.zig
View File

@ -14,7 +14,6 @@ pub fn env() zml.Platform {
if (!builtin.is_test) @compileError("Cannot use zml.testing.env outside of a test block");
if (_platform == null) {
_test_compile_opts = .{
.xla_dump_to = "/tmp/zml/tests/",
.sharding_enabled = true,
};