From e0d219248d7029e693a286887d2509961ce2bbe6 Mon Sep 17 00:00:00 2001 From: "Uwe L. Korn" Date: Tue, 25 Jun 2024 15:29:32 +0200 Subject: [PATCH] jaxlib 0.4.25 --- recipe/meta.yaml | 6 +- .../0001-Allow-for-custom-CUDA-build.patch | 4 +- ...pport-third-party-build-of-boringssl.patch | 6 +- .../patches/0003-xla-Fix-abseil-headers.patch | 68 +++++++++++++------ recipe/patches/0004-Upstream-xla-Patch.patch | 56 --------------- 5 files changed, 54 insertions(+), 86 deletions(-) delete mode 100644 recipe/patches/0004-Upstream-xla-Patch.patch diff --git a/recipe/meta.yaml b/recipe/meta.yaml index 7686ef12..15b19107 100644 --- a/recipe/meta.yaml +++ b/recipe/meta.yaml @@ -1,4 +1,4 @@ -{% set version = "0.4.24" %} +{% set version = "0.4.25" %} {% set number = 0 %} {% if cuda_compiler_version != "None" %} @@ -13,12 +13,11 @@ package: source: # only pull sources after upstream PyPI release... url: https://github.com/google/jax/archive/jaxlib-v{{ version }}.tar.gz - sha256: c4e6963c2c36f634a9a1765e476a1ed4e6c4a7954465ebf72e29f344c28ddc28 + sha256: fc1197c401924942eb14185a61688d0c476e3e81ff71f9dc95e620b57c06eec8 patches: - patches/0001-Allow-for-custom-CUDA-build.patch - patches/0002-xla-Support-third-party-build-of-boringssl.patch - patches/0003-xla-Fix-abseil-headers.patch - - patches/0004-Upstream-xla-Patch.patch build: number: {{ number }} @@ -81,7 +80,6 @@ requirements: - zlib run: - python - - {{ pin_compatible('numpy') }} - scipy >=1.9 - ml_dtypes >=0.2.0 - __cuda # [cuda_compiler_version != "None"] diff --git a/recipe/patches/0001-Allow-for-custom-CUDA-build.patch b/recipe/patches/0001-Allow-for-custom-CUDA-build.patch index 077e2da4..abf2b614 100644 --- a/recipe/patches/0001-Allow-for-custom-CUDA-build.patch +++ b/recipe/patches/0001-Allow-for-custom-CUDA-build.patch @@ -1,7 +1,7 @@ -From 7e9afda0f5ad8a0767f0500eb9933b27dde0414e Mon Sep 17 00:00:00 2001 +From 39bed7f08f495d3a544a11c951869fcfd1616003 Mon Sep 17 00:00:00 2001 From: "Uwe L. Korn" Date: Sun, 8 Oct 2023 19:34:34 +0200 -Subject: [PATCH 1/4] Allow for custom CUDA build +Subject: [PATCH 1/3] Allow for custom CUDA build --- build/build.py | 5 +++++ diff --git a/recipe/patches/0002-xla-Support-third-party-build-of-boringssl.patch b/recipe/patches/0002-xla-Support-third-party-build-of-boringssl.patch index 3cc4cad8..9ff99c4f 100644 --- a/recipe/patches/0002-xla-Support-third-party-build-of-boringssl.patch +++ b/recipe/patches/0002-xla-Support-third-party-build-of-boringssl.patch @@ -1,7 +1,7 @@ -From 691dd89b6d14b7c71302a46f54b1564c486ea7ba Mon Sep 17 00:00:00 2001 +From a7ae8b8397a8a3f0f1206c08f3461aff9b686670 Mon Sep 17 00:00:00 2001 From: "Uwe L. Korn" Date: Thu, 14 Dec 2023 17:06:15 +0100 -Subject: [PATCH 2/4] xla: Support third-party build of boringssl +Subject: [PATCH 2/3] xla: Support third-party build of boringssl --- ...pport-third-party-build-of-boringssl.patch | 54 +++++++++++++++++++ @@ -70,7 +70,7 @@ index 0000000..1f11444 +2.43.0 + diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl -index 0b09fa9..b81e13a 100644 +index ebc8d98..ebb0137 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -29,6 +29,9 @@ def repo(): diff --git a/recipe/patches/0003-xla-Fix-abseil-headers.patch b/recipe/patches/0003-xla-Fix-abseil-headers.patch index b192592a..60b0726a 100644 --- a/recipe/patches/0003-xla-Fix-abseil-headers.patch +++ b/recipe/patches/0003-xla-Fix-abseil-headers.patch @@ -1,35 +1,37 @@ -From 730f54ecf5bd630de69368e0a1108737289ecb38 Mon Sep 17 00:00:00 2001 +From 67075c53bebacbeeda10271e2d5bcb31d4050964 Mon Sep 17 00:00:00 2001 From: "Uwe L. Korn" Date: Sun, 11 Feb 2024 20:18:13 +0100 -Subject: [PATCH 3/4] xla: Fix abseil headers +Subject: [PATCH 3/3] xla: Fix abseil headers --- - third_party/xla/0001-Fix-abseil-headers.patch | 47 +++++++++++++++++++ + third_party/xla/0001-Fix-abseil-headers.patch | 73 +++++++++++++++++++ third_party/xla/workspace.bzl | 1 + - 2 files changed, 48 insertions(+) + 2 files changed, 74 insertions(+) create mode 100644 third_party/xla/0001-Fix-abseil-headers.patch diff --git a/third_party/xla/0001-Fix-abseil-headers.patch b/third_party/xla/0001-Fix-abseil-headers.patch new file mode 100644 -index 0000000..41f411e +index 0000000..ad2cbaa --- /dev/null +++ b/third_party/xla/0001-Fix-abseil-headers.patch -@@ -0,0 +1,47 @@ -+From 97ff4aad5cbb51103ff2bd0880af454b324ff6af Mon Sep 17 00:00:00 2001 -+From: "Uwe L. Korn" -+Date: Sun, 11 Feb 2024 19:35:02 +0100 +@@ -0,0 +1,73 @@ ++From 138a28794618469fbcab7634c5febca42e10fef6 Mon Sep 17 00:00:00 2001 ++From: "Uwe L. Korn" ++Date: Thu, 23 May 2024 15:45:52 +0200 +Subject: [PATCH] Fix abseil headers + +--- + third_party/tsl/tsl/distributed_runtime/rpc/BUILD | 1 + ++ third_party/tsl/tsl/platform/default/BUILD | 2 ++ + third_party/tsl/tsl/profiler/rpc/client/BUILD | 4 ++++ -+ 2 files changed, 5 insertions(+) ++ xla/python/ifrt_proxy/common/BUILD | 3 +++ ++ 4 files changed, 10 insertions(+) + +diff --git a/third_party/tsl/tsl/distributed_runtime/rpc/BUILD b/third_party/tsl/tsl/distributed_runtime/rpc/BUILD -+index 9b52ee1cb..875b575b4 100644 ++index ae18059b63..8b2f356a08 100644 +--- a/third_party/tsl/tsl/distributed_runtime/rpc/BUILD ++++ b/third_party/tsl/tsl/distributed_runtime/rpc/BUILD -+@@ -43,6 +43,7 @@ cc_library( ++@@ -42,6 +42,7 @@ cc_library( + "//tsl/platform:stringprintf", + "//tsl/protobuf:distributed_runtime_payloads_proto_cc", + "@com_google_absl//absl/status", @@ -37,8 +39,21 @@ index 0000000..41f411e + "@com_google_absl//absl/strings:cord", + ] + tsl_grpc_cc_dependencies(), + ) ++diff --git a/third_party/tsl/tsl/platform/default/BUILD b/third_party/tsl/tsl/platform/default/BUILD ++index 6b0e9a8c8c..d3c406fb7b 100644 ++--- a/third_party/tsl/tsl/platform/default/BUILD +++++ b/third_party/tsl/tsl/platform/default/BUILD ++@@ -221,6 +221,8 @@ cc_library( ++ deps = [ ++ "//tsl/platform:logging", ++ "@com_google_absl//absl/log:check", +++ "@com_google_absl//absl/status", +++ "@com_google_absl//absl/status:statusor", ++ ] + tsl_grpc_cc_dependencies(), ++ ) ++ +diff --git a/third_party/tsl/tsl/profiler/rpc/client/BUILD b/third_party/tsl/tsl/profiler/rpc/client/BUILD -+index fb3a39cd1..2b93965e8 100644 ++index 1dd7d8a227..ab28e814e5 100644 +--- a/third_party/tsl/tsl/profiler/rpc/client/BUILD ++++ b/third_party/tsl/tsl/profiler/rpc/client/BUILD +@@ -100,6 +100,8 @@ cc_library( @@ -50,20 +65,31 @@ index 0000000..41f411e + "@com_google_absl//absl/strings", + "@com_google_absl//absl/time", + ], -+@@ -128,6 +130,8 @@ cc_library( ++@@ -126,6 +128,8 @@ cc_library( ++ "//tsl/platform:types", ++ "//tsl/profiler/protobuf:profiler_analysis_cc_grpc_proto", + "//tsl/profiler/protobuf:profiler_service_cc_grpc_proto", +++ "@com_google_absl//absl/status", +++ "@com_google_absl//absl/status:statusor", + "//tsl/protobuf:error_codes_proto_impl_cc", + "@com_google_absl//absl/memory", ++ "@com_google_absl//absl/strings", ++diff --git a/xla/python/ifrt_proxy/common/BUILD b/xla/python/ifrt_proxy/common/BUILD ++index 8e750cb525..aede7facd0 100644 ++--- a/xla/python/ifrt_proxy/common/BUILD +++++ b/xla/python/ifrt_proxy/common/BUILD ++@@ -51,6 +51,9 @@ cc_library( ++ "@com_github_grpc_grpc//:grpc++", ++ "@com_google_absl//absl/log", ++ "@com_google_absl//absl/log:check", ++ "@com_google_absl//absl/status", ++ "@com_google_absl//absl/status:statusor", -+ "@com_google_absl//absl/strings", -+ "@com_google_absl//absl/time", -+ ] + tsl_grpc_cc_dependencies(), -+-- -+2.43.0 -+ +++ "@com_google_absl//absl/strings", ++ "@tsl//tsl/platform", ++ ], ++ alwayslink = True, diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl -index b81e13a..9671b9b 100644 +index ebb0137..1a55893 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -31,6 +31,7 @@ def repo(): diff --git a/recipe/patches/0004-Upstream-xla-Patch.patch b/recipe/patches/0004-Upstream-xla-Patch.patch deleted file mode 100644 index a5a222a9..00000000 --- a/recipe/patches/0004-Upstream-xla-Patch.patch +++ /dev/null @@ -1,56 +0,0 @@ -From 8e28aa73a45c296914c3f4e19c748418303de190 Mon Sep 17 00:00:00 2001 -From: "Uwe L. Korn" -Date: Mon, 24 Jun 2024 16:42:05 +0200 -Subject: [PATCH 4/4] Upstream xla Patch - ---- - ...4cd346594fc7ea2fe75570c9c53a4a444f60.patch | 27 +++++++++++++++++++ - third_party/xla/workspace.bzl | 1 + - 2 files changed, 28 insertions(+) - create mode 100644 third_party/xla/7a614cd346594fc7ea2fe75570c9c53a4a444f60.patch - -diff --git a/third_party/xla/7a614cd346594fc7ea2fe75570c9c53a4a444f60.patch b/third_party/xla/7a614cd346594fc7ea2fe75570c9c53a4a444f60.patch -new file mode 100644 -index 0000000..44abc07 ---- /dev/null -+++ b/third_party/xla/7a614cd346594fc7ea2fe75570c9c53a4a444f60.patch -@@ -0,0 +1,27 @@ -+From 7a614cd346594fc7ea2fe75570c9c53a4a444f60 Mon Sep 17 00:00:00 2001 -+From: David Majnemer -+Date: Mon, 5 Feb 2024 11:12:46 -0800 -+Subject: [PATCH] [XLA:GPU] Fix macOS build -+ -+DelinearizeInBoundsIndex returns a `llvm::SmallVector` but -+`result` has type `llvm::SmallVector` which can cause compilers to -+complain that there is no known conversion available. -+ -+PiperOrigin-RevId: 604374810 -+--- -+ xla/service/gpu/model/indexing_analysis.cc | 2 +- -+ 1 file changed, 1 insertion(+), 1 deletion(-) -+ -+diff --git a/xla/service/gpu/model/indexing_analysis.cc b/xla/service/gpu/model/indexing_analysis.cc -+index 4f9c36faf422d..e859b8c98d02d 100644 -+--- a/xla/service/gpu/model/indexing_analysis.cc -++++ b/xla/service/gpu/model/indexing_analysis.cc -+@@ -687,7 +687,7 @@ std::vector ToTransposeDimensions(const Layout& l) { -+ llvm::SmallVector DelinearizeInBoundsIndex( -+ mlir::AffineExpr linear, absl::Span sizes, -+ absl::Span strides) { -+- llvm::SmallVector result; -++ llvm::SmallVector result; -+ result.reserve(sizes.size()); -+ for (auto [size, stride] : llvm::zip(sizes, strides)) { -+ result.push_back(linear.floorDiv(stride) % size); -diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl -index 9671b9b..76f050d 100644 ---- a/third_party/xla/workspace.bzl -+++ b/third_party/xla/workspace.bzl -@@ -32,6 +32,7 @@ def repo(): - patch_file = [ - "//third_party/xla:0001-Support-third-party-build-of-boringssl.patch", - "//third_party/xla:0001-Fix-abseil-headers.patch", -+ "//third_party/xla:7a614cd346594fc7ea2fe75570c9c53a4a444f60.patch", - ], - ) -