xla: bump to 22ea002 (#37)

Author: hugomano

.github/workflows/build.yaml

@@ -19,7 +19,7 @@ permissions:
   contents: write
 
 env:
-  XLA_COMMIT: ${{ inputs.xla_commit || '6e48cbb8d33d771c964697e39bfaf678bcc6de31' }}
+  XLA_COMMIT: ${{ inputs.xla_commit || '22ea002baf40702b90a0b431485d07b63e3fe2fa' }}
 
 jobs:
   setup_openxla:

openxla/bazelrc/cuda.bazelrc

@@ -10,6 +10,7 @@ build      --action_env CLANG_CUDA_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
 build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="5.0,5.2,5.3,6.0,6.1,6.2,7.0,7.2,7.5,8.0,8.6,8.7,8.9,9.0,10.0,10.1,12.0"
 build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.8.0"
 build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.8.0"
+build:cuda --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5"
 build:cuda --@//xla/stream_executor/cuda:enable_libnvjitlink_support=True
 build:cuda --@//xla/stream_executor/cuda:enable_libnvptxcompiler_support=True
 

openxla/patches/0001-Added-FFI-handler-registration-API-to-the-FFI-PjRt.patch

@@ -1,17 +1,17 @@
-From a33cf614d0e351a9123c0bc4435f53231675f5a0 Mon Sep 17 00:00:00 2001
+From 50621f9cde6abe56bd3a63d1d3acf6513657ff4a Mon Sep 17 00:00:00 2001
 From: Hugo Mano <hugo@zml.ai>
 Date: Tue, 27 May 2025 11:48:17 +0200
-Subject: [PATCH 1/6] Added FFI handler registration API to the FFI PjRt
+Subject: [PATCH 1/7] Added FFI handler registration API to the FFI PjRt
 
 PR: https://github.com/openxla/xla/pull/13420
 ---
  xla/pjrt/c/BUILD                      |  5 +++++
- xla/pjrt/c/pjrt_c_api_ffi_extension.h | 21 ++++++++++++++++++
+ xla/pjrt/c/pjrt_c_api_ffi_extension.h | 23 ++++++++++++++++++-
  xla/pjrt/c/pjrt_c_api_ffi_internal.cc | 32 ++++++++++++++++++++++++++-
- 3 files changed, 57 insertions(+), 1 deletion(-)
+ 3 files changed, 58 insertions(+), 2 deletions(-)
 
 diff --git a/xla/pjrt/c/BUILD b/xla/pjrt/c/BUILD
-index 054d966f47..7b78ca8241 100644
+index f60f939edb..3943d172aa 100644
 --- a/xla/pjrt/c/BUILD
 +++ b/xla/pjrt/c/BUILD
 @@ -69,8 +69,13 @@ cc_library(
@@ -29,10 +29,10 @@ index 054d966f47..7b78ca8241 100644
  )
 
 diff --git a/xla/pjrt/c/pjrt_c_api_ffi_extension.h b/xla/pjrt/c/pjrt_c_api_ffi_extension.h
-index 995a2c7e50..b8f10bc2f7 100644
+index 995a2c7e50..7c83ca9fbc 100644
 --- a/xla/pjrt/c/pjrt_c_api_ffi_extension.h
 +++ b/xla/pjrt/c/pjrt_c_api_ffi_extension.h
-@@ -69,10 +69,31 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_UserData_Add_Args, user_data);
+@@ -69,12 +69,33 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_UserData_Add_Args, user_data);
  // Adds a user data to the execute context.
  typedef PJRT_Error* PJRT_FFI_UserData_Add(PJRT_FFI_UserData_Add_Args* args);
 
@@ -62,8 +62,11 @@ index 995a2c7e50..b8f10bc2f7 100644
    PJRT_FFI_UserData_Add* user_data_add;
 +  PJRT_FFI_Register_Handler* register_handler;
  } PJRT_FFI;
- PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_Extension, user_data_add);
+-PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_Extension, user_data_add);
++PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_Extension, register_handler);
 
+ #ifdef __cplusplus
+ }
 diff --git a/xla/pjrt/c/pjrt_c_api_ffi_internal.cc b/xla/pjrt/c/pjrt_c_api_ffi_internal.cc
 index 5fa88eab33..63ad65bc69 100644
 --- a/xla/pjrt/c/pjrt_c_api_ffi_internal.cc

openxla/patches/0002-Drop-the-.so-suffix-for-PjRt-plugin-targets-.dylib-o.patch

@@ -1,19 +1,19 @@
-From 6479728ad9b2534776d4224c6d58fcc6b2a667ee Mon Sep 17 00:00:00 2001
+From 9494c46866ff2d00e6208aabf9316b3ea4544ceb Mon Sep 17 00:00:00 2001
 From: Hugo Mano <hugo@zml.ai>
 Date: Tue, 27 May 2025 11:49:51 +0200
-Subject: [PATCH 2/6] Drop the .so suffix for PjRt plugin targets (.dylib on
+Subject: [PATCH 2/7] Drop the .so suffix for PjRt plugin targets (.dylib on
  macOS)
 
 PR: https://github.com/openxla/xla/pull/16696
 ---
- xla/pjrt/c/BUILD | 5 +++--
- 1 file changed, 3 insertions(+), 2 deletions(-)
+ xla/pjrt/c/BUILD | 4 ++--
+ 1 file changed, 2 insertions(+), 2 deletions(-)
 
 diff --git a/xla/pjrt/c/BUILD b/xla/pjrt/c/BUILD
-index 7b78ca8241..207e991972 100644
+index 3943d172aa..3047c382d4 100644
 --- a/xla/pjrt/c/BUILD
 +++ b/xla/pjrt/c/BUILD
-@@ -332,13 +332,14 @@ cc_library(
+@@ -359,7 +359,7 @@ cc_library(
 
  # PJRT CPU plugin.
  xla_cc_binary(
@@ -22,14 +22,7 @@ index 7b78ca8241..207e991972 100644
      additional_linker_inputs = [
          ":pjrt_c_api_cpu_version_script.lds",
      ],
-     linkopts = if_macos(
-         [
-             "-Wl,-exported_symbol,_GetPjrtApi",
-+            "-Wl,-install_name,@rpath/pjrt_c_api_cpu_plugin.dylib",
-             "-Wl,-undefined,error",
-         ],
-         [
-@@ -425,7 +426,7 @@ cc_library(
+@@ -453,7 +453,7 @@ cc_library(
 
  # PJRT GPU plugin. Can be configured to be built for CUDA or ROCM.
  xla_cc_binary(

openxla/patches/0003-Expose-nvptxcompiler-to-link-against-in-XLA-if-enabl.patch

@@ -1,58 +1,58 @@
-From 08f8bb1431e1536b07118e0fe85826d7958b56a7 Mon Sep 17 00:00:00 2001
+From ff930c6b192a4c72aec42fe9fe1c2d51661db95b Mon Sep 17 00:00:00 2001
 From: Hugo Mano <hugo@zml.ai>
 Date: Tue, 27 May 2025 11:50:40 +0200
-Subject: [PATCH 3/6] Expose nvptxcompiler to link against in XLA if
+Subject: [PATCH 3/7] Expose nvptxcompiler to link against in XLA if
  enable_libnvptxcompiler_support is set
 
 Only for ZML, no PR on XLA side.
 ---
- third_party/gpus/cuda/hermetic/cuda_nvcc.BUILD.tpl |  8 ++++++++
- xla/stream_executor/cuda/BUILD                     | 12 +++++++++++-
- 2 files changed, 19 insertions(+), 1 deletion(-)
+ third_party/gpus/cuda/hermetic/BUILD.tpl | 1 +
+ workspace0.bzl                           | 6 +++---
+ xla/stream_executor/cuda/BUILD           | 7 ++++++-
+ 3 files changed, 10 insertions(+), 4 deletions(-)
 
-diff --git a/third_party/gpus/cuda/hermetic/cuda_nvcc.BUILD.tpl b/third_party/gpus/cuda/hermetic/cuda_nvcc.BUILD.tpl
-index 733d95d371..f6c1c5ab53 100644
---- a/third_party/gpus/cuda/hermetic/cuda_nvcc.BUILD.tpl
-+++ b/third_party/gpus/cuda/hermetic/cuda_nvcc.BUILD.tpl
-@@ -52,6 +52,14 @@ filegroup(
-     visibility = ["//visibility:public"],
+diff --git a/third_party/gpus/cuda/hermetic/BUILD.tpl b/third_party/gpus/cuda/hermetic/BUILD.tpl
+index 41a273ebbe..d789e71883 100644
+--- a/third_party/gpus/cuda/hermetic/BUILD.tpl
++++ b/third_party/gpus/cuda/hermetic/BUILD.tpl
+@@ -307,6 +307,7 @@ cc_library(
+     # This is not yet fully supported, but we need the rule
+     # to make bazel query happy.
+     name = "nvptxcompiler",
++    actual = "@cuda_nvcc//:libnvptxcompiler",
  )
 
-+filegroup(
-+    name = "libnvptxcompiler",
-+    srcs = [
-+        "lib/libnvptxcompiler_static.a",
-+    ],
-+    visibility = ["//visibility:public"],
-+)
-+
- filegroup(
-     name = "bin",
-     srcs = glob([
+ alias(
+diff --git a/workspace0.bzl b/workspace0.bzl
+index 25db007e51..e6ccc318f0 100644
+--- a/workspace0.bzl
++++ b/workspace0.bzl
+@@ -138,10 +138,10 @@ def workspace():
+     # Details: https://github.com/google-ml-infra/rules_ml_toolchain
+     http_archive(
+         name = "rules_ml_toolchain",
+-        sha256 = "de3b14418657eeacd8afc2aa89608be6ec8d66cd6a5de81c4f693e77bc41bee1",
+-        strip_prefix = "rules_ml_toolchain-5653e5a0ca87c1272069b4b24864e55ce7f129a1",
++        sha256 = "fb78d09234528aef2be856820b69b76486829f65e4eb3c7ffaa5803b667fa441",
++        strip_prefix = "rules_ml_toolchain-f4ad89fa906be2c1374785a79335c8a7dcd49df7",
+         urls = [
+-            "https://github.com/google-ml-infra/rules_ml_toolchain/archive/5653e5a0ca87c1272069b4b24864e55ce7f129a1.tar.gz",
++            "https://github.com/zml/rules_ml_toolchain/archive/f4ad89fa906be2c1374785a79335c8a7dcd49df7.tar.gz",
+         ],
+     )
+ 
 diff --git a/xla/stream_executor/cuda/BUILD b/xla/stream_executor/cuda/BUILD
-index 4a62a44357..d5f723fd62 100644
+index 08abf6197f..ac06e394bc 100644
 --- a/xla/stream_executor/cuda/BUILD
 +++ b/xla/stream_executor/cuda/BUILD
-@@ -84,6 +84,11 @@ config_setting(
-     },
- )
- 
-+cc_import(
-+    name = "nvptxcompiler",
-+    static_library = "@cuda_nvcc//:libnvptxcompiler",
-+)
-+
- cc_library(
-     name = "cuda_platform_id",
-     srcs = ["cuda_platform_id.cc"],
-@@ -128,7 +133,12 @@ cc_library(
+@@ -124,7 +124,12 @@ cc_library(
              "@tsl//tsl/platform:errors",
              "@tsl//tsl/platform:status",
              "@tsl//tsl/platform:statusor",
 -        ] + tf_additional_cuda_platform_deps(),
 +        ] + tf_additional_cuda_platform_deps() + select({
 +            ":libnvptxcompiler_support_enabled": [
-+                ":nvptxcompiler",
++                "@local_config_cuda//cuda:nvptxcompiler",
 +            ],
 +            "//conditions:default": [],
 +        }),

openxla/patches/0004-build-use-hermetic-cc-toolchain-for-Linux-CPU-use-gl.patch

@@ -1,23 +1,22 @@
-From fed224a7f56daabfe127fce715090c44966cdcc5 Mon Sep 17 00:00:00 2001
+From cc37221ecd006b6685ab788820b07c8c1180a03f Mon Sep 17 00:00:00 2001
 From: Hugo Mano <hugo@zml.ai>
 Date: Tue, 27 May 2025 11:51:17 +0200
-Subject: [PATCH 4/6] build: use hermetic cc toolchain for Linux CPU (use glibc
+Subject: [PATCH 4/7] build: use hermetic cc toolchain for Linux CPU (use glibc
  2.31)
 
 Only for ZML, no PR on XLA side.
 ---
- WORKSPACE | 24 ++++++++++++++++++++++++
- 1 file changed, 24 insertions(+)
+ WORKSPACE | 23 +++++++++++++++++++++++
+ 1 file changed, 23 insertions(+)
 
 diff --git a/WORKSPACE b/WORKSPACE
-index d8f4ae6e60..50701a23d6 100644
+index 67d505addc..e67c73e3f5 100644
 --- a/WORKSPACE
 +++ b/WORKSPACE
-@@ -126,3 +126,27 @@ load(
- )
+@@ -137,3 +137,26 @@ cc_toolchain_deps()
+ register_toolchains("@rules_ml_toolchain//cc_toolchain:lx64_lx64")
 
- nvshmem_configure(name = "local_config_nvshmem")
-+
+ register_toolchains("@rules_ml_toolchain//cc_toolchain:lx64_lx64_cuda")
 +load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 +
 +HERMETIC_CC_TOOLCHAIN_VERSION = "v3.1.1"
@@ -41,7 +40,6 @@ index d8f4ae6e60..50701a23d6 100644
 +register_toolchains(
 +    "@zig_sdk//toolchain:linux_amd64_gnu.2.31",
 +)
-\ No newline at end of file
 -- 
 2.39.5 (Apple Git-154)