Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

jaxlib 0.4.35 #291

Merged
merged 1 commit into from
Dec 11, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions recipe/build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ fi
source gen-bazel-toolchain

cat >> .bazelrc <<EOF

build --crosstool_top=//bazel_toolchain:toolchain
build --logging=6
build --verbose_failures
Expand Down
6 changes: 3 additions & 3 deletions recipe/meta.yaml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
{% set version = "0.4.34" %}
{% set build = 2 %}
{% set version = "0.4.35" %}
{% set build = 0 %}

{% if cuda_compiler_version != "None" %}
{% set build = build + 200 %}
Expand All @@ -13,7 +13,7 @@ package:
source:
# only pull sources after upstream PyPI release...
url: https://github.com/google/jax/archive/jax-v{{ version }}.tar.gz
sha256: d3a75ad667772309ade81350fa70c4a78028a920028800282e46d8383c0ee6bb
sha256: 65e086708ae56670676b7b2340ad82b901d8c9993d1241a839c8990bdb8d6212
patches:
- patches/0001-Allow-for-custom-CUDA-build.patch
- patches/0002-Consolidated-build-fixes-for-XLA.patch
Expand Down
4 changes: 2 additions & 2 deletions recipe/patches/0001-Allow-for-custom-CUDA-build.patch
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
From 868e01e0e6672f0b85d63e5564889131a8052e5f Mon Sep 17 00:00:00 2001
From 66310dedabfce69aea6814da770a95755ecc3ef2 Mon Sep 17 00:00:00 2001
From: "Uwe L. Korn" <[email protected]>
Date: Sun, 8 Oct 2023 19:34:34 +0200
Subject: [PATCH 1/3] Allow for custom CUDA build
Expand All @@ -8,7 +8,7 @@ Subject: [PATCH 1/3] Allow for custom CUDA build
1 file changed, 5 insertions(+)

diff --git a/build/build.py b/build/build.py
index de0d5a9..8b7a21f 100755
index 44343eb..0a7a003 100755
--- a/build/build.py
+++ b/build/build.py
@@ -292,6 +292,11 @@ def write_bazelrc(*, remote_build,
Expand Down
76 changes: 44 additions & 32 deletions recipe/patches/0002-Consolidated-build-fixes-for-XLA.patch
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
From e1ce0773ec91aeb5cad98b7ba7332e287151eb2a Mon Sep 17 00:00:00 2001
From 0bfc707ab06b3f1cdcb451fb52a4759895938f54 Mon Sep 17 00:00:00 2001
From: "Uwe L. Korn" <[email protected]>
Date: Thu, 14 Dec 2023 17:06:15 +0100
Subject: [PATCH 2/3] Consolidated build fixes for XLA
Expand All @@ -16,27 +16,29 @@ Co-Authored-By: H. Vetinari <[email protected]>
---
...pport-third-party-build-of-boringssl.patch | 51 ++++
third_party/xla/0002-Fix-abseil-headers.patch | 73 ++++++
.../xla/0003-Omit-usage-of-StrFormat.patch | 24 ++
.../xla/0003-Omit-usage-of-StrFormat.patch | 25 ++
...0004-Add-missing-bits-absl-systemlib.patch | 231 ++++++++++++++++++
...ther-absl-log-is-already-initialized.patch | 58 +++++
third_party/xla/workspace.bzl | 7 +
6 files changed, 444 insertions(+)
.../xla/0006-Add-conda-cuda-path.patch | 31 +++
third_party/xla/workspace.bzl | 8 +
7 files changed, 477 insertions(+)
create mode 100644 third_party/xla/0001-Support-third-party-build-of-boringssl.patch
create mode 100644 third_party/xla/0002-Fix-abseil-headers.patch
create mode 100644 third_party/xla/0003-Omit-usage-of-StrFormat.patch
create mode 100644 third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
create mode 100644 third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch
create mode 100644 third_party/xla/0006-Add-conda-cuda-path.patch

diff --git a/third_party/xla/0001-Support-third-party-build-of-boringssl.patch b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch
new file mode 100644
index 0000000..514ebe1
index 0000000..dd9de48
--- /dev/null
+++ b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch
@@ -0,0 +1,51 @@
+From 12f526179896176c1726d5fa9eab2f0332e8ee21 Mon Sep 17 00:00:00 2001
+From a9bb7b6fc4e14025c68596a609d64e4da8960996 Mon Sep 17 00:00:00 2001
+From: "Uwe L. Korn" <[email protected]>
+Date: Thu, 14 Dec 2023 15:04:51 +0100
+Subject: [PATCH 1/5] Support third-party build of boringssl
+Subject: [PATCH 1/6] Support third-party build of boringssl
+
+---
+ third_party/boringssl.BUILD | 21 +++++++++++++++++++++
Expand Down Expand Up @@ -86,14 +88,14 @@ index 0000000..514ebe1
+
diff --git a/third_party/xla/0002-Fix-abseil-headers.patch b/third_party/xla/0002-Fix-abseil-headers.patch
new file mode 100644
index 0000000..e47e8a6
index 0000000..d3aafc8
--- /dev/null
+++ b/third_party/xla/0002-Fix-abseil-headers.patch
@@ -0,0 +1,73 @@
+From 87e126882738dd5440430f03ff87fdb4ca31516b Mon Sep 17 00:00:00 2001
+From 3f962611a6b5bbbb9c7d1b5f72d0b6e3feae09cb Mon Sep 17 00:00:00 2001
+From: "Uwe L. Korn" <[email protected]>
+Date: Thu, 23 May 2024 15:45:52 +0200
+Subject: [PATCH 2/5] Fix abseil headers
+Subject: [PATCH 2/6] Fix abseil headers
+
+---
+ third_party/tsl/tsl/platform/default/BUILD | 2 ++
Expand All @@ -103,7 +105,7 @@ index 0000000..e47e8a6
+ 4 files changed, 10 insertions(+)
+
+diff --git a/third_party/tsl/tsl/platform/default/BUILD b/third_party/tsl/tsl/platform/default/BUILD
+index 8cbb72cfc9..c99667c050 100644
+index 7f18f8108f..94cdfe8540 100644
+--- a/third_party/tsl/tsl/platform/default/BUILD
++++ b/third_party/tsl/tsl/platform/default/BUILD
+@@ -225,6 +225,8 @@ cc_library(
Expand Down Expand Up @@ -142,7 +144,7 @@ index 0000000..e47e8a6
+ "@tsl//tsl/platform:protobuf",
+ "@tsl//tsl/platform:status",
+diff --git a/xla/tsl/profiler/rpc/client/BUILD b/xla/tsl/profiler/rpc/client/BUILD
+index e135bd82dc..81b79f4600 100644
+index ac61088fd3..2861f77d9a 100644
+--- a/xla/tsl/profiler/rpc/client/BUILD
++++ b/xla/tsl/profiler/rpc/client/BUILD
+@@ -98,6 +98,8 @@ cc_library(
Expand All @@ -154,38 +156,39 @@ index 0000000..e47e8a6
+ "@com_google_absl//absl/strings",
+ "@com_google_absl//absl/time",
+ "@tsl//tsl/platform:status",
+@@ -121,6 +123,8 @@ cc_library(
+ "//tensorflow/python/profiler/internal:__pkg__",
+@@ -122,6 +124,8 @@ cc_library(
+ ]),
+ deps = [
+ "//xla/tsl/protobuf:error_codes_proto_impl_cc",
++ "@com_google_absl//absl/status",
++ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/memory",
+ "@com_google_absl//absl/strings",
+ "@com_google_absl//absl/time",
diff --git a/third_party/xla/0003-Omit-usage-of-StrFormat.patch b/third_party/xla/0003-Omit-usage-of-StrFormat.patch
new file mode 100644
index 0000000..ec44acf
index 0000000..01fd4df
--- /dev/null
+++ b/third_party/xla/0003-Omit-usage-of-StrFormat.patch
@@ -0,0 +1,24 @@
+From 62ff28675279b0d72cfe477a94dc1d133d680630 Mon Sep 17 00:00:00 2001
@@ -0,0 +1,25 @@
+From ac263373f90084807fe80665538ca5018bc3a951 Mon Sep 17 00:00:00 2001
+From: "Uwe L. Korn" <[email protected]>
+Date: Thu, 4 Jul 2024 10:36:03 +0200
+Subject: [PATCH 3/5] Omit usage of StrFormat
+Subject: [PATCH 3/6] Omit usage of StrFormat
+
+---
+ xla/stream_executor/gpu/gpu_executor.h | 4 ++++
+ 1 file changed, 4 insertions(+)
+ xla/stream_executor/gpu/gpu_executor.h | 5 +++++
+ 1 file changed, 5 insertions(+)
+
+diff --git a/xla/stream_executor/gpu/gpu_executor.h b/xla/stream_executor/gpu/gpu_executor.h
+index 0fdde6096e..6d0825301f 100644
+index 0d90dadf84..25afa7ac9b 100644
+--- a/xla/stream_executor/gpu/gpu_executor.h
++++ b/xla/stream_executor/gpu/gpu_executor.h
+@@ -18,6 +18,10 @@ limitations under the License.
+@@ -17,6 +17,11 @@ limitations under the License.
+ #define XLA_STREAM_EXECUTOR_GPU_GPU_EXECUTOR_H_
+
+ #include <cstdint>
+ #include <memory>
++#include <memory>
++#include <optional>
++#include <string>
++#include <sstream>
Expand All @@ -195,14 +198,14 @@ index 0000000..ec44acf
+ #include <vector>
diff --git a/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
new file mode 100644
index 0000000..2230d05
index 0000000..cca52f1
--- /dev/null
+++ b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
@@ -0,0 +1,231 @@
+From 0e97115cbe5c1068834f1aa2e4e4956215edca82 Mon Sep 17 00:00:00 2001
+From 1bd70877037d67945c0834aca417a66480416902 Mon Sep 17 00:00:00 2001
+From: "Uwe L. Korn" <[email protected]>
+Date: Thu, 4 Jul 2024 15:58:32 +0200
+Subject: [PATCH 4/5] Add missing bits absl systemlib
+Subject: [PATCH 4/6] Add missing bits absl systemlib
+
+---
+ .../third_party/absl/system.absl.base.BUILD | 16 +++++
Expand Down Expand Up @@ -432,14 +435,14 @@ index 0000000..2230d05
+ "numeric",
diff --git a/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch b/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch
new file mode 100644
index 0000000..a7f7f0b
index 0000000..b648cf9
--- /dev/null
+++ b/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch
@@ -0,0 +1,58 @@
+From 4c362c9a88ff5ca06a9a075be12e94f79b9affef Mon Sep 17 00:00:00 2001
+From a4fa7b3d14e3d29ee50a84f4e74e9f76c301ea92 Mon Sep 17 00:00:00 2001
+From: "Uwe L. Korn" <[email protected]>
+Date: Fri, 22 Nov 2024 10:51:18 +0100
+Subject: [PATCH 5/5] Check whether absl log is already initialized
+Subject: [PATCH 5/6] Check whether absl log is already initialized
+
+---
+ xla/pjrt/c/pjrt_c_api_gpu.cc | 7 ++++++-
Expand Down Expand Up @@ -496,10 +499,19 @@ index 0000000..a7f7f0b
+ } // namespace xla
diff --git a/third_party/xla/0006-Add-conda-cuda-path.patch b/third_party/xla/0006-Add-conda-cuda-path.patch
new file mode 100644
index 0000000..a7f7f0b
index 0000000..dcd6a1a
--- /dev/null
+++ b/third_party/xla/0006-Add-conda-cuda-path.patch
@@ -0,0 +1,22 @@
@@ -0,0 +1,31 @@
+From 79444089903f92a228425a000e47a41b4d9dd810 Mon Sep 17 00:00:00 2001
+From: Silvio Traversaro <[email protected]>
+Date: Tue, 10 Dec 2024 09:42:51 +0100
+Subject: [PATCH 6/6] Add conda cuda path
+
+---
+ .../tsl/tsl/platform/default/cuda_root_path.cc | 11 +++++++++++
+ 1 file changed, 11 insertions(+)
+
+diff --git a/third_party/tsl/tsl/platform/default/cuda_root_path.cc b/third_party/tsl/tsl/platform/default/cuda_root_path.cc
+index ca6da0e553..1d8a9450c0 100644
+--- a/third_party/tsl/tsl/platform/default/cuda_root_path.cc
Expand All @@ -523,7 +535,7 @@ index 0000000..a7f7f0b
+ #endif // defined(PLATFORM_POSIX) && !defined(__APPLE__)
+
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index 37861d9..45fbcde 100644
index 3dd8cbd..1ec0ff4 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -30,6 +30,14 @@ def repo():
Expand Down
6 changes: 3 additions & 3 deletions recipe/patches/0003-Add-missing-typename.patch
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
From fd329cb4cd76ed6221f9ef4f9edddd13fb77d2cf Mon Sep 17 00:00:00 2001
From 317afc9ff9e2cb021bc81a278710a6a3f147c536 Mon Sep 17 00:00:00 2001
From: "Uwe L. Korn" <[email protected]>
Date: Wed, 9 Oct 2024 09:39:46 +0000
Subject: [PATCH 3/3] Add missing typename
Expand All @@ -8,10 +8,10 @@ Subject: [PATCH 3/3] Add missing typename
1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jaxlib/gpu/solver_kernels_ffi.cc b/jaxlib/gpu/solver_kernels_ffi.cc
index 32cd975..c90dfab 100644
index 7852da4..b5742b5 100644
--- a/jaxlib/gpu/solver_kernels_ffi.cc
+++ b/jaxlib/gpu/solver_kernels_ffi.cc
@@ -815,7 +815,7 @@ ffi::Error GesvdjImpl(int64_t batch, int64_t rows, int64_t cols,
@@ -915,7 +915,7 @@ ffi::Error GesvdjImpl(int64_t batch, int64_t rows, int64_t cols,

auto a_data = static_cast<T*>(a.untyped_data());
auto out_data = static_cast<T*>(out->untyped_data());
Expand Down