diff --git a/.ci_support/osx_64_numpy2.0python3.10.____cpython.yaml b/.ci_support/osx_64_numpy2.0python3.10.____cpython.yaml index 8ddb54b4..ccb62c11 100644 --- a/.ci_support/osx_64_numpy2.0python3.10.____cpython.yaml +++ b/.ci_support/osx_64_numpy2.0python3.10.____cpython.yaml @@ -1,15 +1,15 @@ MACOSX_DEPLOYMENT_TARGET: -- '10.14' +- '10.15' MACOSX_SDK_VERSION: -- '10.14' +- '10.15' c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: -- '10.14' +- '10.15' channel_sources: - conda-forge channel_targets: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/.ci_support/osx_64_numpy2.0python3.11.____cpython.yaml b/.ci_support/osx_64_numpy2.0python3.11.____cpython.yaml index e7ead7df..8f061b00 100644 --- a/.ci_support/osx_64_numpy2.0python3.11.____cpython.yaml +++ b/.ci_support/osx_64_numpy2.0python3.11.____cpython.yaml @@ -1,15 +1,15 @@ MACOSX_DEPLOYMENT_TARGET: -- '10.14' +- '10.15' MACOSX_SDK_VERSION: -- '10.14' +- '10.15' c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: -- '10.14' +- '10.15' channel_sources: - conda-forge channel_targets: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/.ci_support/osx_64_numpy2.0python3.12.____cpython.yaml b/.ci_support/osx_64_numpy2.0python3.12.____cpython.yaml index 33092e95..fbc2715f 100644 --- a/.ci_support/osx_64_numpy2.0python3.12.____cpython.yaml +++ b/.ci_support/osx_64_numpy2.0python3.12.____cpython.yaml @@ -1,15 +1,15 @@ MACOSX_DEPLOYMENT_TARGET: -- '10.14' +- '10.15' MACOSX_SDK_VERSION: -- '10.14' +- '10.15' c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: -- '10.14' +- '10.15' channel_sources: - conda-forge channel_targets: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/.ci_support/osx_64_numpy2python3.13.____cp313.yaml b/.ci_support/osx_64_numpy2python3.13.____cp313.yaml index 021aa02e..fb4f1d0d 100644 --- a/.ci_support/osx_64_numpy2python3.13.____cp313.yaml +++ b/.ci_support/osx_64_numpy2python3.13.____cp313.yaml @@ -1,15 +1,15 @@ MACOSX_DEPLOYMENT_TARGET: -- '10.14' +- '10.15' MACOSX_SDK_VERSION: -- '10.14' +- '10.15' c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: -- '10.14' +- '10.15' channel_sources: - conda-forge channel_targets: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/.ci_support/osx_arm64_numpy2.0python3.10.____cpython.yaml b/.ci_support/osx_arm64_numpy2.0python3.10.____cpython.yaml index f3f5abe9..0e4d5c15 100644 --- a/.ci_support/osx_arm64_numpy2.0python3.10.____cpython.yaml +++ b/.ci_support/osx_arm64_numpy2.0python3.10.____cpython.yaml @@ -5,7 +5,7 @@ MACOSX_SDK_VERSION: c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/.ci_support/osx_arm64_numpy2.0python3.11.____cpython.yaml b/.ci_support/osx_arm64_numpy2.0python3.11.____cpython.yaml index 1feab52b..94febae7 100644 --- a/.ci_support/osx_arm64_numpy2.0python3.11.____cpython.yaml +++ b/.ci_support/osx_arm64_numpy2.0python3.11.____cpython.yaml @@ -5,7 +5,7 @@ MACOSX_SDK_VERSION: c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/.ci_support/osx_arm64_numpy2.0python3.12.____cpython.yaml b/.ci_support/osx_arm64_numpy2.0python3.12.____cpython.yaml index 2e4ac60b..472f508b 100644 --- a/.ci_support/osx_arm64_numpy2.0python3.12.____cpython.yaml +++ b/.ci_support/osx_arm64_numpy2.0python3.12.____cpython.yaml @@ -5,7 +5,7 @@ MACOSX_SDK_VERSION: c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/.ci_support/osx_arm64_numpy2python3.13.____cp313.yaml b/.ci_support/osx_arm64_numpy2python3.13.____cp313.yaml index f758ce42..a7ce9eb7 100644 --- a/.ci_support/osx_arm64_numpy2python3.13.____cp313.yaml +++ b/.ci_support/osx_arm64_numpy2python3.13.____cp313.yaml @@ -5,7 +5,7 @@ MACOSX_SDK_VERSION: c_compiler: - clang c_compiler_version: -- '17' +- '18' c_stdlib: - macosx_deployment_target c_stdlib_version: @@ -23,7 +23,7 @@ cudnn: cxx_compiler: - clangxx cxx_compiler_version: -- '17' +- '18' github_actions_labels: - cirun-openstack-cpu-large libabseil: diff --git a/recipe/build.sh b/recipe/build.sh index 47164dab..196f0bfb 100644 --- a/recipe/build.sh +++ b/recipe/build.sh @@ -1,9 +1,6 @@ #!/bin/bash set -euxo pipefail -# see comment in meta.yaml -cp $RECIPE_DIR/patches/0007-xla-cpu-Fix-build-errors-from-ACL.patch $SRC_DIR/third_party/xla/ - export JAX_RELEASE=$PKG_VERSION $RECIPE_DIR/add_py_toolchain.sh diff --git a/recipe/conda_build_config.yaml b/recipe/conda_build_config.yaml index 3d073802..92a54d49 100644 --- a/recipe/conda_build_config.yaml +++ b/recipe/conda_build_config.yaml @@ -1,12 +1,4 @@ -# Otherwise __ZN4absl12lts_2024072212log_internal10LogMessagelsIPKvTnNSt3__19enable_ifIXntsr4absl16HasAbslStringifyIT_EE5valueEiE4typeELi0EEERS2_RKS8_ cannot be found -# demangled: absl::lts_20240722::log_internal::LogMessage& absl::lts_20240722::log_internal::LogMessage::operator<<(void const* const&) -c_compiler_version: # [osx] -- '17' # [osx] -cxx_compiler_version: # [osx] -- '17' # [osx] -MACOSX_SDK_VERSION: # [osx and x86_64] -- '10.14' # [osx and x86_64] c_stdlib_version: # [osx and x86_64] -- '10.14' # [osx and x86_64] +- '10.15' # [osx and x86_64] github_actions_labels: - cirun-openstack-cpu-large diff --git a/recipe/meta.yaml b/recipe/meta.yaml index 50394d89..dc0c9ca1 100644 --- a/recipe/meta.yaml +++ b/recipe/meta.yaml @@ -1,4 +1,4 @@ -{% set version = "0.5.0" %} +{% set version = "0.5.1" %} {% set build = 0 %} {% if cuda_compiler_version != "None" %} @@ -13,19 +13,14 @@ package: source: # only pull sources after upstream PyPI release... url: https://github.com/jax-ml/jax/archive/jax-v{{ version }}.tar.gz - sha256: 04cc2eeb2e7ce1916674cea03a7d75a59d583ddb779d5104e103a2798a283ce9 + sha256: e74b1209517682075933f757d646b73040d09fe39ee3e9e4cd398407dd0902d2 patches: - patches/0001-Allow-for-custom-CUDA-build.patch - patches/0002-Consolidated-build-fixes-for-XLA.patch - # cannot absorb this into the overall xla patch, because patching three vendored projects - # deep breaks the application of the patch using `patch` (presumably due to lines starting - # legitimately with `+++` being misinterpreted as a hunk separator). - # - patches/0007-xla-cpu-Fix-build-errors-from-ACL.patch # backport https://github.com/jax-ml/jax/pull/26691 - patches/0003-fix-member-access-to-packed-CUDA-struct.patch - - patches/0004-fix-an-ambiguous-type.patch # backport https://github.com/jax-ml/jax/pull/26690 - - patches/0005-fix-getting-gcc-major-version.patch + - patches/0004-fix-getting-gcc-major-version.patch build: number: {{ build }} diff --git a/recipe/patches/0001-Allow-for-custom-CUDA-build.patch b/recipe/patches/0001-Allow-for-custom-CUDA-build.patch index cd43d56a..0ad61c11 100644 --- a/recipe/patches/0001-Allow-for-custom-CUDA-build.patch +++ b/recipe/patches/0001-Allow-for-custom-CUDA-build.patch @@ -1,17 +1,17 @@ -From f8e41b61ce27bf6bc28a10703ddfceea39d650ed Mon Sep 17 00:00:00 2001 +From c97dee8337618e173ed5d617eb67deabdda026c5 Mon Sep 17 00:00:00 2001 From: "Uwe L. Korn" Date: Sun, 8 Oct 2023 19:34:34 +0200 -Subject: [PATCH 1/5] Allow for custom CUDA build +Subject: [PATCH 1/4] Allow for custom CUDA build --- build/build.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/build/build.py b/build/build.py -index d2f68f80e..4fbb675fc 100755 +index 0df7d646f..6b9690ca1 100755 --- a/build/build.py +++ b/build/build.py -@@ -548,6 +548,13 @@ async def main(): +@@ -549,6 +549,13 @@ async def main(): if args.cuda_version: logging.debug("Hermetic CUDA version: %s", args.cuda_version) @@ -25,7 +25,7 @@ index d2f68f80e..4fbb675fc 100755 wheel_build_command_base.append( f"--repo_env=HERMETIC_CUDA_VERSION={args.cuda_version}" ) -@@ -653,10 +660,9 @@ async def main(): +@@ -656,10 +663,9 @@ async def main(): wheel_build_command.append(f"--jaxlib_git_hash={git_hash}") diff --git a/recipe/patches/0002-Consolidated-build-fixes-for-XLA.patch b/recipe/patches/0002-Consolidated-build-fixes-for-XLA.patch index 7a8f303f..e1271db9 100644 --- a/recipe/patches/0002-Consolidated-build-fixes-for-XLA.patch +++ b/recipe/patches/0002-Consolidated-build-fixes-for-XLA.patch @@ -1,7 +1,7 @@ -From b0cb237b915315590486cdb74c8500b42c22bddc Mon Sep 17 00:00:00 2001 +From 1dd800f6f64d0948feccef00840d74b5c61094e0 Mon Sep 17 00:00:00 2001 From: "Uwe L. Korn" Date: Thu, 14 Dec 2023 17:06:15 +0100 -Subject: [PATCH 2/5] Consolidated build fixes for XLA +Subject: [PATCH 2/4] Consolidated build fixes for XLA jax vendors xla, but only populates the sources through bazel, so we cannot patch as usual through conda, but rather need to teach the bazel build file @@ -15,30 +15,28 @@ which is also where we're patching in the list of patches to apply to xla. Co-Authored-By: H. Vetinari --- ...pport-third-party-build-of-boringssl.patch | 51 ++++ - third_party/xla/0002-Fix-abseil-headers.patch | 73 ++++++ + third_party/xla/0002-Fix-abseil-headers.patch | 73 +++++ .../xla/0003-Omit-usage-of-StrFormat.patch | 25 ++ - ...0004-Add-missing-bits-absl-systemlib.patch | 245 ++++++++++++++++++ - ...ther-absl-log-is-already-initialized.patch | 58 +++++ - .../xla/0006-Add-conda-cuda-path.patch | 31 +++ - third_party/xla/workspace.bzl | 10 + - 7 files changed, 493 insertions(+) + ...0004-Add-missing-bits-absl-systemlib.patch | 262 ++++++++++++++++++ + ...ther-absl-log-is-already-initialized.patch | 58 ++++ + third_party/xla/workspace.bzl | 7 + + 6 files changed, 476 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 000000000..74b1c5310 +index 000000000..8554f5126 --- /dev/null +++ b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch @@ -0,0 +1,51 @@ -+From 163c1236a63376eb43be7c3b055f49607d3ff07d Mon Sep 17 00:00:00 2001 ++From c6bf0f14e159dcbdc99c725823d23e1d26b73a04 Mon Sep 17 00:00:00 2001 +From: "Uwe L. Korn" +Date: Thu, 14 Dec 2023 15:04:51 +0100 -+Subject: [PATCH 1/7] Support third-party build of boringssl ++Subject: [PATCH 1/5] Support third-party build of boringssl + +--- + third_party/boringssl.BUILD | 21 +++++++++++++++++++++ @@ -74,7 +72,7 @@ index 000000000..74b1c5310 ++ ], ++) +diff --git a/workspace2.bzl b/workspace2.bzl -+index f8b1cf3764..a8a36cf364 100644 ++index cded32c8ca..5fb10ca45f 100644 +--- a/workspace2.bzl ++++ b/workspace2.bzl +@@ -108,7 +108,7 @@ def _tf_repositories(): @@ -88,14 +86,14 @@ index 000000000..74b1c5310 + 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 000000000..219d663e5 +index 000000000..7b2084d01 --- /dev/null +++ b/third_party/xla/0002-Fix-abseil-headers.patch @@ -0,0 +1,73 @@ -+From 8bdb1195668d1ecef8c9bee70dc20bce11529a5d Mon Sep 17 00:00:00 2001 ++From 1620a0eb9135e380b59583539080337853c130aa Mon Sep 17 00:00:00 2001 +From: "Uwe L. Korn" +Date: Thu, 23 May 2024 15:45:52 +0200 -+Subject: [PATCH 2/7] Fix abseil headers ++Subject: [PATCH 2/5] Fix abseil headers + +--- + xla/python/ifrt_proxy/common/BUILD | 3 +++ @@ -105,7 +103,7 @@ index 000000000..219d663e5 + 4 files changed, 10 insertions(+) + +diff --git a/xla/python/ifrt_proxy/common/BUILD b/xla/python/ifrt_proxy/common/BUILD -+index f994ff87ef..87033dec85 100644 ++index 5da3b67ecb..c684b54d52 100644 +--- a/xla/python/ifrt_proxy/common/BUILD ++++ b/xla/python/ifrt_proxy/common/BUILD +@@ -51,6 +51,9 @@ cc_library( @@ -119,7 +117,7 @@ index 000000000..219d663e5 + ], + alwayslink = True, +diff --git a/xla/tsl/distributed_runtime/rpc/BUILD b/xla/tsl/distributed_runtime/rpc/BUILD -+index 383846735f..45d6314c5a 100644 ++index 83aaccd1d9..dec829ddc7 100644 +--- a/xla/tsl/distributed_runtime/rpc/BUILD ++++ b/xla/tsl/distributed_runtime/rpc/BUILD +@@ -39,6 +39,7 @@ cc_library( @@ -144,7 +142,7 @@ index 000000000..219d663e5 + ) + +diff --git a/xla/tsl/profiler/rpc/client/BUILD b/xla/tsl/profiler/rpc/client/BUILD -+index 58144f683b..c571538247 100644 ++index c2554c43cd..5430b49357 100644 +--- a/xla/tsl/profiler/rpc/client/BUILD ++++ b/xla/tsl/profiler/rpc/client/BUILD +@@ -99,6 +99,8 @@ cc_library( @@ -167,14 +165,14 @@ index 000000000..219d663e5 + "@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 000000000..c82e6e469 +index 000000000..ac7be8f59 --- /dev/null +++ b/third_party/xla/0003-Omit-usage-of-StrFormat.patch @@ -0,0 +1,25 @@ -+From bd00fc129d585e972712923f2458d315261d5f32 Mon Sep 17 00:00:00 2001 ++From 7240379e230c6095e8e490aa8c00e9dc97c0c4f8 Mon Sep 17 00:00:00 2001 +From: "Uwe L. Korn" +Date: Thu, 4 Jul 2024 10:36:03 +0200 -+Subject: [PATCH 3/7] Omit usage of StrFormat ++Subject: [PATCH 3/5] Omit usage of StrFormat + +--- + xla/stream_executor/gpu/gpu_executor.h | 5 +++++ @@ -198,23 +196,23 @@ index 000000000..c82e6e469 + #include 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 000000000..b23805861 +index 000000000..d1ffe736c --- /dev/null +++ b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch -@@ -0,0 +1,245 @@ -+From 5856e6ee773620f29fe762b8c39d988d762e49e3 Mon Sep 17 00:00:00 2001 +@@ -0,0 +1,262 @@ ++From fa66d4f7f1597cd20962cae167c0a9f09080954b Mon Sep 17 00:00:00 2001 +From: "Uwe L. Korn" +Date: Thu, 4 Jul 2024 15:58:32 +0200 -+Subject: [PATCH 4/7] Add missing bits absl systemlib ++Subject: [PATCH 4/5] Add missing bits absl systemlib + +--- -+ .../third_party/absl/system.absl.base.BUILD | 16 +++++ -+ .../third_party/absl/system.absl.crc.BUILD | 70 +++++++++++++++++++ ++ .../third_party/absl/system.absl.base.BUILD | 16 ++++ ++ .../third_party/absl/system.absl.crc.BUILD | 70 +++++++++++++++++ + .../third_party/absl/system.absl.flags.BUILD | 2 +- -+ .../third_party/absl/system.absl.log.BUILD | 61 ++++++++++++++++ -+ .../absl/system.absl.strings.BUILD | 12 ++++ ++ .../third_party/absl/system.absl.log.BUILD | 78 +++++++++++++++++++ ++ .../absl/system.absl.strings.BUILD | 12 +++ + .../tsl/third_party/absl/workspace.bzl | 2 + -+ 6 files changed, 162 insertions(+), 1 deletion(-) ++ 6 files changed, 179 insertions(+), 1 deletion(-) + create mode 100644 third_party/tsl/third_party/absl/system.absl.crc.BUILD + create mode 100644 third_party/tsl/third_party/absl/system.absl.log.BUILD + @@ -340,10 +338,10 @@ index 000000000..b23805861 + ":flag_internal", +diff --git a/third_party/tsl/third_party/absl/system.absl.log.BUILD b/third_party/tsl/third_party/absl/system.absl.log.BUILD +new file mode 100644 -+index 0000000000..67350d613b ++index 0000000000..c23aa51acf +--- /dev/null ++++ b/third_party/tsl/third_party/absl/system.absl.log.BUILD -+@@ -0,0 +1,61 @@ ++@@ -0,0 +1,78 @@ ++load("@rules_cc//cc:defs.bzl", "cc_library") ++ ++package(default_visibility = ["//visibility:public"]) @@ -386,6 +384,23 @@ index 000000000..b23805861 ++) ++ ++cc_library( +++ name = "globals", +++ linkopts = [ +++ "-labsl_log_globals", +++ "-labsl_log_severity", +++ "-labsl_raw_logging_internal", +++ "-labsl_vlog_config_internal", +++ ], +++ deps = [ +++ ":log", +++ "//absl/base:config", +++ "//absl/base:core_headers", +++ "//absl/hash", +++ "//absl/strings", +++ ], +++) +++ +++cc_library( ++ name = "initialize", ++ linkopts = [ ++ "-labsl_log_initialize", @@ -449,14 +464,14 @@ index 000000000..b23805861 + "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 000000000..64b675274 +index 000000000..271a3ca1e --- /dev/null +++ b/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch @@ -0,0 +1,58 @@ -+From b79360dc6d626fbf368e7a61a846aad9af11863d Mon Sep 17 00:00:00 2001 ++From 9ebd25248eda049702cb781894ca0e097f6843b0 Mon Sep 17 00:00:00 2001 +From: "Uwe L. Korn" +Date: Fri, 22 Nov 2024 10:51:18 +0100 -+Subject: [PATCH 5/7] Check whether absl log is already initialized ++Subject: [PATCH 5/5] Check whether absl log is already initialized + +--- + xla/pjrt/c/pjrt_c_api_gpu.cc | 7 ++++++- @@ -511,48 +526,11 @@ index 000000000..64b675274 ++} + + } // 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 000000000..fdc5ee661 ---- /dev/null -+++ b/third_party/xla/0006-Add-conda-cuda-path.patch -@@ -0,0 +1,31 @@ -+From 005ea05408a4331da63dc9d75a0bd146aec3eef2 Mon Sep 17 00:00:00 2001 -+From: Silvio Traversaro -+Date: Tue, 10 Dec 2024 09:42:51 +0100 -+Subject: [PATCH 6/7] Add conda cuda path -+ -+--- -+ xla/tsl/platform/default/cuda_root_path.cc | 11 +++++++++++ -+ 1 file changed, 11 insertions(+) -+ -+diff --git a/xla/tsl/platform/default/cuda_root_path.cc b/xla/tsl/platform/default/cuda_root_path.cc -+index 578d8b05c7..e6cdcda7fe 100644 -+--- a/xla/tsl/platform/default/cuda_root_path.cc -++++ b/xla/tsl/platform/default/cuda_root_path.cc -+@@ -76,6 +76,17 @@ std::vector CandidateCudaRoots() { -+ // Also add the path to the copy of libdevice.10.bc that we include within -+ // the Python wheel. -+ roots.emplace_back(io::JoinPath(dir, "cuda")); -++ -++ // In case cuda was installed with nvidia's official conda packages, we also -++ // include the root prefix of the environment in the candidate roots dir, -++ // we assume that the lib binaries are either in the python package's root dir -++ // or in a 'python' subdirectory, as done by the previous for. -++ // python packages on non-Windows platforms are installed in -++ // $CONDA_PREFIX/lib/python3.12/site-packages/pkg_name, so if we want -++ // to add $CONDA_PREFIX to the candidate roots dirs we need to add -++ // ../../../.. -++ for (auto path : {"../../../..", "../../../../.."}) -++ roots.emplace_back(io::JoinPath(dir, path)); -+ } -+ #endif // defined(PLATFORM_POSIX) && !defined(__APPLE__) -+ diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl -index a015d9427..4848b6986 100644 +index 185ee6a39..a7d083e08 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl -@@ -30,6 +30,16 @@ def repo(): +@@ -30,6 +30,13 @@ def repo(): sha256 = XLA_SHA256, strip_prefix = "xla-{commit}".format(commit = XLA_COMMIT), urls = tf_mirror_urls("https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)), @@ -562,9 +540,6 @@ index a015d9427..4848b6986 100644 + "//third_party/xla:0003-Omit-usage-of-StrFormat.patch", + "//third_party/xla:0004-Add-missing-bits-absl-systemlib.patch", + "//third_party/xla:0005-Check-whether-absl-log-is-already-initialized.patch", -+ "//third_party/xla:0006-Add-conda-cuda-path.patch", -+ # backport https://github.com/openxla/xla/commit/14ba309d98db689ba5c185287483d555e3307b7f -+ "//third_party/xla:0007-xla-cpu-Fix-build-errors-from-ACL.patch", + ], ) diff --git a/recipe/patches/0003-fix-member-access-to-packed-CUDA-struct.patch b/recipe/patches/0003-fix-member-access-to-packed-CUDA-struct.patch index f14d2832..42a6f488 100644 --- a/recipe/patches/0003-fix-member-access-to-packed-CUDA-struct.patch +++ b/recipe/patches/0003-fix-member-access-to-packed-CUDA-struct.patch @@ -1,17 +1,17 @@ -From a7f66bd413f1c148652b6cf4a6ee00cf507e5e45 Mon Sep 17 00:00:00 2001 +From 094cf4d9cb2dc56c26ae984f7c31db0343c2a686 Mon Sep 17 00:00:00 2001 From: "H. Vetinari" Date: Sat, 22 Feb 2025 16:39:41 +1100 -Subject: [PATCH 3/5] fix member access to packed CUDA struct +Subject: [PATCH 3/4] fix member access to packed CUDA struct --- jaxlib/mosaic/gpu/mosaic_gpu_ext.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc b/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc -index e5c85ac58..0112a044d 100644 +index 312605df8..2f415912f 100644 --- a/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc +++ b/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc -@@ -183,8 +183,9 @@ void callback_complete(CUcontext context, uint32_t streamId, +@@ -182,8 +182,9 @@ void callback_complete(CUcontext context, uint32_t streamId, // Convert integer nanoseconds to floating point milliseconds to match // the interface of the events-based profiler. double duration_ms = (kernel->end - kernel->start) / 1e6; diff --git a/recipe/patches/0004-fix-an-ambiguous-type.patch b/recipe/patches/0004-fix-an-ambiguous-type.patch deleted file mode 100644 index 05e07c39..00000000 --- a/recipe/patches/0004-fix-an-ambiguous-type.patch +++ /dev/null @@ -1,22 +0,0 @@ -From 4dba4c17b8efc75832cf7cf17128541eb75cc510 Mon Sep 17 00:00:00 2001 -From: "H. Vetinari" -Date: Sun, 23 Feb 2025 02:27:18 +1100 -Subject: [PATCH 4/5] fix an ambiguous type - ---- - jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc | 2 +- - 1 file changed, 1 insertion(+), 1 deletion(-) - -diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc -index c4035fea2..93f4f8d74 100644 ---- a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc -+++ b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc -@@ -406,7 +406,7 @@ llvm::LogicalResult WGMMAOp::verify() { - - int groups_m = 0; - auto a_shape = a_shaped_type.getShape(); -- if (auto a_memref = dyn_cast(getA().getType())) { -+ if (auto a_memref = mlir::dyn_cast(getA().getType())) { - if (a_shape.size() != 4) { - return error("When `a` is a memref, it must have rank 4."); - } diff --git a/recipe/patches/0005-fix-getting-gcc-major-version.patch b/recipe/patches/0004-fix-getting-gcc-major-version.patch similarity index 77% rename from recipe/patches/0005-fix-getting-gcc-major-version.patch rename to recipe/patches/0004-fix-getting-gcc-major-version.patch index 658a9477..0022a9ff 100644 --- a/recipe/patches/0005-fix-getting-gcc-major-version.patch +++ b/recipe/patches/0004-fix-getting-gcc-major-version.patch @@ -1,14 +1,14 @@ -From 3d55ab98e2f62513bd1ac473c6905fd3059c90ba Mon Sep 17 00:00:00 2001 +From 295e968fca6184dd4b412bf3be3a0402ab82e8cc Mon Sep 17 00:00:00 2001 From: "H. Vetinari" Date: Mon, 24 Feb 2025 07:59:42 +1100 -Subject: [PATCH 5/5] fix getting gcc major version +Subject: [PATCH 4/4] fix getting gcc major version --- build/tools/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build/tools/utils.py b/build/tools/utils.py -index 03a762ac3..0a6d33f51 100644 +index f38dc5bda..255dd98f4 100644 --- a/build/tools/utils.py +++ b/build/tools/utils.py @@ -208,7 +208,7 @@ def get_gcc_major_version(gcc_path: str): diff --git a/recipe/patches/0007-xla-cpu-Fix-build-errors-from-ACL.patch b/recipe/patches/0007-xla-cpu-Fix-build-errors-from-ACL.patch deleted file mode 100644 index 811f2ca3..00000000 --- a/recipe/patches/0007-xla-cpu-Fix-build-errors-from-ACL.patch +++ /dev/null @@ -1,128 +0,0 @@ -From 65e170bb822f37e5846c03dfeb1ebe0bc7e97f96 Mon Sep 17 00:00:00 2001 -From: Penporn Koanantakool -Date: Fri, 24 Jan 2025 13:33:38 -0800 -Subject: [PATCH 7/7] [xla:cpu] Fix build errors from ACL - -Reported in: -https://github.com/jax-ml/jax/issues/26062 - -PiperOrigin-RevId: 719416696 ---- - .../exclude_omp_scheduler.patch | 23 +++++++++++++++++++ - .../compute_library/include_string.patch | 12 ++++++++++ - .../exclude_omp_scheduler.patch | 23 +++++++++++++++++++ - .../compute_library/include_string.patch | 12 ++++++++++ - third_party/tsl/workspace2.bzl | 2 ++ - 5 files changed, 72 insertions(+) - create mode 100644 third_party/compute_library/exclude_omp_scheduler.patch - create mode 100644 third_party/compute_library/include_string.patch - create mode 100644 third_party/tsl/third_party/compute_library/exclude_omp_scheduler.patch - create mode 100644 third_party/tsl/third_party/compute_library/include_string.patch - -diff --git a/third_party/compute_library/exclude_omp_scheduler.patch b/third_party/compute_library/exclude_omp_scheduler.patch -new file mode 100644 -index 0000000000..7ccfebbd36 ---- /dev/null -+++ b/third_party/compute_library/exclude_omp_scheduler.patch -@@ -0,0 +1,23 @@ -+diff --git a/src/BUILD.bazel b/src/BUILD.bazel -+index bf71e534e2..22377f1a32 100644 -+--- a/src/BUILD.bazel -++++ b/src/BUILD.bazel -+@@ -971,7 +971,6 @@ filegroup( -+ "runtime/NEON/functions/NETranspose.cpp", -+ "runtime/NEON/functions/NEUnstack.cpp", -+ "runtime/NEON/functions/NEWinogradConvolutionLayer.cpp", -+- "runtime/OMP/OMPScheduler.cpp", -+ "runtime/OffsetLifetimeManager.cpp", -+ "runtime/OffsetMemoryPool.cpp", -+ "runtime/OperatorTensor.cpp", -+@@ -984,6 +983,10 @@ filegroup( -+ "runtime/Tensor.cpp", -+ "runtime/TensorAllocator.cpp", -+ "runtime/Utils.cpp"] + -++ select({ -++ "//:openmp_flag": ["runtime/OMP/OMPScheduler.cpp"], -++ "//conditions:default": [], -++ }) + -+ glob(["**/*.h", -+ "**/*.hpp", -+ "**/*.inl"]), -diff --git a/third_party/compute_library/include_string.patch b/third_party/compute_library/include_string.patch -new file mode 100644 -index 0000000000..709aeea621 ---- /dev/null -+++ b/third_party/compute_library/include_string.patch -@@ -0,0 +1,12 @@ -+diff --git a/arm_compute/core/utils/logging/IPrinter.h b/arm_compute/core/utils/logging/IPrinter.h -+index 42dca58ea1..9c8b1598e8 100644 -+--- a/arm_compute/core/utils/logging/IPrinter.h -++++ b/arm_compute/core/utils/logging/IPrinter.h -+@@ -25,6 +25,7 @@ -+ #define ARM_COMPUTE_LOGGING_PRINTER_H -+ -+ #include "support/Mutex.h" -++#include -+ -+ namespace arm_compute -+ { -diff --git a/third_party/tsl/third_party/compute_library/exclude_omp_scheduler.patch b/third_party/tsl/third_party/compute_library/exclude_omp_scheduler.patch -new file mode 100644 -index 0000000000..7ccfebbd36 ---- /dev/null -+++ b/third_party/tsl/third_party/compute_library/exclude_omp_scheduler.patch -@@ -0,0 +1,23 @@ -+diff --git a/src/BUILD.bazel b/src/BUILD.bazel -+index bf71e534e2..22377f1a32 100644 -+--- a/src/BUILD.bazel -++++ b/src/BUILD.bazel -+@@ -971,7 +971,6 @@ filegroup( -+ "runtime/NEON/functions/NETranspose.cpp", -+ "runtime/NEON/functions/NEUnstack.cpp", -+ "runtime/NEON/functions/NEWinogradConvolutionLayer.cpp", -+- "runtime/OMP/OMPScheduler.cpp", -+ "runtime/OffsetLifetimeManager.cpp", -+ "runtime/OffsetMemoryPool.cpp", -+ "runtime/OperatorTensor.cpp", -+@@ -984,6 +983,10 @@ filegroup( -+ "runtime/Tensor.cpp", -+ "runtime/TensorAllocator.cpp", -+ "runtime/Utils.cpp"] + -++ select({ -++ "//:openmp_flag": ["runtime/OMP/OMPScheduler.cpp"], -++ "//conditions:default": [], -++ }) + -+ glob(["**/*.h", -+ "**/*.hpp", -+ "**/*.inl"]), -diff --git a/third_party/tsl/third_party/compute_library/include_string.patch b/third_party/tsl/third_party/compute_library/include_string.patch -new file mode 100644 -index 0000000000..709aeea621 ---- /dev/null -+++ b/third_party/tsl/third_party/compute_library/include_string.patch -@@ -0,0 +1,12 @@ -+diff --git a/arm_compute/core/utils/logging/IPrinter.h b/arm_compute/core/utils/logging/IPrinter.h -+index 42dca58ea1..9c8b1598e8 100644 -+--- a/arm_compute/core/utils/logging/IPrinter.h -++++ b/arm_compute/core/utils/logging/IPrinter.h -+@@ -25,6 +25,7 @@ -+ #define ARM_COMPUTE_LOGGING_PRINTER_H -+ -+ #include "support/Mutex.h" -++#include -+ -+ namespace arm_compute -+ { -diff --git a/third_party/tsl/workspace2.bzl b/third_party/tsl/workspace2.bzl -index f3a237dd7f..102b0624fc 100644 ---- a/third_party/tsl/workspace2.bzl -+++ b/third_party/tsl/workspace2.bzl -@@ -179,6 +179,8 @@ def _tf_repositories(): - patch_file = [ - "//third_party/compute_library:compute_library.patch", - "//third_party/compute_library:acl_thread_local_scheduler.patch", -+ "//third_party/compute_library:exclude_omp_scheduler.patch", -+ "//third_party/compute_library:include_string.patch", - ], - sha256 = "c4ca329a78da380163b2d86e91ba728349b6f0ee97d66e260a694ef37f0b0d93", - strip_prefix = "ComputeLibrary-23.05.1",