zml · hugomano · Feb 6, 2025 · Feb 4, 2025 · Feb 4, 2025 · Feb 4, 2025
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -19,7 +19,7 @@ permissions:
   contents: write
 
 env:
-  XLA_COMMIT: ${{ inputs.xla_commit || 'cc075beb6148c2777da2b6749c63830856ee6c2a' }}
+  XLA_COMMIT: ${{ inputs.xla_commit || '6789523d0a42c2cf83309a263f417122c4c63811' }}
 
 jobs:
   setup_openxla:
@@ -56,6 +56,7 @@ jobs:
   pjrt-artifacts:
     runs-on: ${{ matrix.pjrt.runs_on }}
     strategy:
+      fail-fast: false
       matrix:
         pjrt:
           - target: cuda
@@ -83,7 +84,7 @@ jobs:
             config: "--config=release_macos_base"
             artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
             renamed_artifact: libpjrt_cpu.dylib
-            runs_on: ["macos-14"]
+            runs_on: ["self-hosted", "macOS"]
             platform: darwin-amd64
             bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
           - target: cpu
@@ -159,15 +160,9 @@ jobs:
         working-directory: ./xla
         run: |
           ~/bin/bazel \
+              --output_base=/tmp/bazel_output_base_${{ matrix.pjrt.platform }} \
               build \
-              --bes_results_url=https://app.buildbuddy.io/invocation/ \
-              --bes_backend=grpcs://remote.buildbuddy.io \
-              --remote_cache=grpcs://remote.buildbuddy.io \
-              --experimental_remote_build_event_upload=minimal \
-              --nolegacy_important_outputs \
-              --remote_download_toplevel \
-              --remote_timeout=3600 \
-              --remote_header=x-buildbuddy-api-key=${{ secrets.BUILDBUDDY_ORG_API_KEY }} \
+              --spawn_strategy=processwrapper-sandbox \
               ${{ matrix.pjrt.config }} \
               ${{ matrix.pjrt.bazel_target }}
 

diff --git a/openxla/bazelrc/common.bazelrc b/openxla/bazelrc/common.bazelrc
@@ -7,3 +7,4 @@ build --repo_env=HERMETIC_PYTHON_VERSION=3.11
 
 test  --test_env LD_LIBRARY_PATH
 test  --test_size_filters small,medium
+
diff --git a/openxla/bazelrc/cpu.bazelrc b/openxla/bazelrc/cpu.bazelrc
@@ -0,0 +1,6 @@
+common          --enable_platform_specific_config
+build           --incompatible_strict_action_env
+build           --incompatible_enable_cc_toolchain_resolution
+build:linux     --strip=always
+build:linux     --repo_env BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1
+build:linux     --sandbox_add_mount_pair=/tmp
diff --git a/openxla/patches/20240318-001-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extension.patch b/openxla/patches/20240318-001-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extension.patch
@@ -3,6 +3,7 @@ From: Steeve Morin <steeve.morin@gmail.com>
 Date: Mon, 18 Mar 2024 17:17:34 +0100
 Subject: [PATCH] [PJRT C API] Ensure C compliance for Profiler Extension
 
+PR: https://github.com/openxla/xla/pull/22082
 ---
  xla/pjrt/c/pjrt_c_api_profiler_extension.h | 2 ++
  1 file changed, 2 insertions(+)

diff --git a/...a/patches/20240605-001-Added-FFI-handler-registration-API-to-the-FFI-PjRt-extension.patch b/...a/patches/20240605-001-Added-FFI-handler-registration-API-to-the-FFI-PjRt-extension.patch
@@ -0,0 +1,131 @@
+From c79f202be6fde802b4e5d697a5925d7eccea3d25 Mon Sep 17 00:00:00 2001
+From: Hugo Mano <hugo@zml.ai>
+Date: Wed, 5 Feb 2025 19:25:03 +0100
+Subject: [PATCH] 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 | 16 ++++++++++++
+ xla/pjrt/c/pjrt_c_api_ffi_internal.cc | 35 +++++++++++++++++++++++++--
+ 3 files changed, 54 insertions(+), 2 deletions(-)
+
+diff --git a/xla/pjrt/c/BUILD b/xla/pjrt/c/BUILD
+index ad1b3987fe..0598281ad1 100644
+--- a/xla/pjrt/c/BUILD
++++ b/xla/pjrt/c/BUILD
+@@ -69,7 +69,12 @@ cc_library(
+         ":pjrt_c_api_wrapper_impl",
+         "//xla/ffi:execution_context",
+         "//xla/ffi:type_id_registry",
++        "//xla/ffi:ffi_api",
++        "//xla/ffi/api:c_api",
++        "//xla/ffi/api:ffi",
++        "//xla/service:custom_call_target_registry",
+         "@com_google_absl//absl/status",
++        "@com_google_absl//absl/strings:str_format",
+     ],
+ )
+
+diff --git a/xla/pjrt/c/pjrt_c_api_ffi_extension.h b/xla/pjrt/c/pjrt_c_api_ffi_extension.h
+index c5766f2a19..3d74e7cbf3 100644
+--- a/xla/pjrt/c/pjrt_c_api_ffi_extension.h
++++ b/xla/pjrt/c/pjrt_c_api_ffi_extension.h
+@@ -67,12 +67,28 @@ 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);
+
++struct PJRT_FFI_Register_Handler_Args {
++  size_t struct_size;
++  const char* target_name;
++  size_t target_name_size;
++  int api_version;  // 0 for an untyped call, 1 -- for typed
++  void* handler;
++  const char* platform_name;
++  size_t platform_name_size;
++};
++PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_Register_Handler_Args, handler);
++
++// Registers an FFI call handler for a specific platform.
++typedef PJRT_Error* PJRT_FFI_Register_Handler(
++    PJRT_FFI_Register_Handler_Args* args);
++
+ typedef struct PJRT_FFI_Extension {
+   size_t struct_size;
+   PJRT_Extension_Type type;
+   PJRT_Extension_Base* next;
+   PJRT_FFI_TypeID_Register* type_id_register;
+   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);
+
+diff --git a/xla/pjrt/c/pjrt_c_api_ffi_internal.cc b/xla/pjrt/c/pjrt_c_api_ffi_internal.cc
+index 0375b39d0b..3527a0756e 100644
+--- a/xla/pjrt/c/pjrt_c_api_ffi_internal.cc
++++ b/xla/pjrt/c/pjrt_c_api_ffi_internal.cc
+@@ -13,15 +13,20 @@ See the License for the specific language governing permissions and
+ limitations under the License.
+ ==============================================================================*/
+
+-#include "xla/pjrt/c/pjrt_c_api_ffi_internal.h"
++#include <string>
+
+ #include "absl/status/status.h"
++#include "absl/strings/str_format.h"
++#include "xla/ffi/api/c_api.h"
++#include "xla/ffi/api/ffi.h"
+ #include "xla/ffi/execution_context.h"
+-#include "xla/ffi/type_id_registry.h"
++ #include "xla/ffi/type_id_registry.h"
++#include "xla/ffi/ffi_api.h"
+ #include "xla/pjrt/c/pjrt_c_api.h"
+ #include "xla/pjrt/c/pjrt_c_api_ffi_extension.h"
+ #include "xla/pjrt/c/pjrt_c_api_helpers.h"
+ #include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h"
++#include "xla/service/custom_call_target_registry.h"
+
+ namespace pjrt {
+
+@@ -55,6 +60,31 @@ static PJRT_Error* PJRT_FFI_UserData_Add(PJRT_FFI_UserData_Add_Args* args) {
+   return nullptr;
+ }
+
++static PJRT_Error* PJRT_FFI_Register_Handler(
++    PJRT_FFI_Register_Handler_Args* args) {
++  PJRT_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
++      "PJRT_FFI_Register_Handler_Args",
++      PJRT_FFI_Register_Handler_Args_STRUCT_SIZE, args->struct_size));
++  std::string target_name(args->target_name, args->target_name_size);
++  std::string platform_name(args->platform_name, args->platform_name_size);
++  switch (args->api_version) {
++    case 0:
++      xla::CustomCallTargetRegistry::Global()->Register(
++          target_name, args->handler, platform_name);
++      return nullptr;
++    case 1:
++      xla::ffi::Ffi::RegisterStaticHandler(
++          xla::ffi::GetXlaFfiApi(), target_name, platform_name,
++          reinterpret_cast<XLA_FFI_Handler*>(args->handler));
++      return nullptr;
++    default:
++      return new PJRT_Error{absl::UnimplementedError(
++          absl::StrFormat("API version %d not supported for PJRT GPU plugin. "
++                          "Supported versions are 0 and 1.",
++                          args->api_version))};
++  }
++}
++
+ PJRT_FFI_Extension CreateFfiExtension(PJRT_Extension_Base* next) {
+   return {
+       /*struct_size=*/PJRT_FFI_Extension_STRUCT_SIZE,
+@@ -62,6 +92,7 @@ PJRT_FFI_Extension CreateFfiExtension(PJRT_Extension_Base* next) {
+       /*next=*/next,
+       /*type_id_register=*/PJRT_FFI_TypeID_Register,
+       /*user_data_add=*/PJRT_FFI_UserData_Add,
++      /*register_handler=*/PJRT_FFI_Register_Handler,
+   };
+ }
+
+-- 
+2.39.5 (Apple Git-154)
diff --git a/openxla/patches/20240901-001-Do-not-force-DEVELOPER_DIR-on-macOS.patch b/openxla/patches/20240901-001-Do-not-force-DEVELOPER_DIR-on-macOS.patch
@@ -5,6 +5,8 @@ Subject: [PATCH] Do not force DEVELOPER_DIR on macOS
 
 It's is autodetected by Bazel and supports building using
 only the Apple Command Line Tools.
+
+PR: https://github.com/openxla/xla/pull/16696
 ---
  tensorflow.bazelrc       | 1 -
  third_party/tsl/.bazelrc | 1 -

diff --git a/...la/patches/20240901-002-Set-the-macosx-deployment-target-via-the-bazel-command-line.patch b/...la/patches/20240901-002-Set-the-macosx-deployment-target-via-the-bazel-command-line.patch
@@ -4,6 +4,8 @@ Date: Sun, 1 Sep 2024 14:53:15 +0200
 Subject: [PATCH] Set the macosx deployment target via the bazel command line
 
 Instead of via an action env.
+
+PR: https://github.com/openxla/xla/pull/16696
 ---
  tensorflow.bazelrc       | 4 ++--
  third_party/tsl/.bazelrc | 4 ++--

diff --git a/openxla/patches/20240901-003-Only-export-GetPjrtApi-symbol-on-macOS.patch b/openxla/patches/20240901-003-Only-export-GetPjrtApi-symbol-on-macOS.patch
@@ -5,6 +5,8 @@ Subject: [PATCH] [PjRt] Only export GetPjrtApi symbol on macOS
 
 Also add missing macOS linkops, remove the .so suffix to the plugin targets
 and add additional_linker_inputs for the linker script instead of deps.
+
+PR: https://github.com/openxla/xla/pull/16696
 ---
  xla/pjrt/c/BUILD | 18 ++++++++++--------
  1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/openxla/patches/20250120-001-Enable-nvptxcompiler-with-nvjitlink.patch b/openxla/patches/20250120-001-Enable-nvptxcompiler-with-nvjitlink.patch
@@ -4,6 +4,7 @@ Date: Tue, 21 Jan 2025 14:41:42 +0100
 Subject: [PATCH] Expose nvptxcompiler to link against in XLA if
  enable_libnvptxcompiler_support is set
 
+Only for ZML, no PR on XLA side.
 ---
  .../gpus/cuda/hermetic/cuda_nvcc.BUILD.tpl           |  8 ++++++++
  xla/stream_executor/cuda/BUILD                       | 12 +++++++++++-

diff --git a/openxla/patches/20250128-001-PJRT-Expose-should_stage_host_to_device_transfers.patch b/openxla/patches/20250128-001-PJRT-Expose-should_stage_host_to_device_transfers.patch
@@ -6,6 +6,8 @@ Subject: [PATCH] [PJRT]  Expose should_stage_host_to_device_transfers  as
 
 Expose GPU option `should_stage_host_to_device_transfers `
 as configuration in PJRT client create func.
+
+PR: https://github.com/openxla/xla/pull/21965
 ---
  xla/pjrt/c/pjrt_c_api_gpu_internal.cc | 7 +++++++
  1 file changed, 7 insertions(+)

diff --git a/...s/20250129-0001-perfs-fixup-20-perf.patch → ...es/20250129-001-perfs-fixup-20-perf.patch b/...s/20250129-0001-perfs-fixup-20-perf.patch → ...es/20250129-001-perfs-fixup-20-perf.patch
@@ -3,6 +3,7 @@ From: Jean-Baptiste Dalido <jb@zml.ai>
 Date: Wed, 29 Jan 2025 15:42:34 +0000
 Subject: [PATCH] perfs: fixup -20% perf
 
+PR: https://github.com/openxla/xla/pull/22081
 ---
  xla/service/gpu/gpu_compiler.cc | 1 +
  1 file changed, 1 insertion(+)

diff --git a/openxla/patches/20250205-001-Use-hermetic-CC-toolchain-glibc-2.31.patch b/openxla/patches/20250205-001-Use-hermetic-CC-toolchain-glibc-2.31.patch
@@ -0,0 +1,44 @@
+From 86de695ba9579ace447bc6c4f5c54bc03d467f85 Mon Sep 17 00:00:00 2001
+From: Hugo Mano <hugo@zml.ai>
+Date: Wed, 5 Feb 2025 16:28:27 +0100
+Subject: [PATCH] build: use hermetic cc toolchain for Linux CPU (use glibc 2.31)
+
+Only for ZML, no PR on XLA side.
+---
+ WORKSPACE | 24 ++++++++++++++++++++++++
+ 1 files changed, 24 insertions(+)
+
+diff --git a/WORKSPACE b/WORKSPACE
+index 028dcdc7ef..55b6ed691f 100644
+--- a/WORKSPACE
++++ b/WORKSPACE
+@@ -99,3 +99,27 @@ load(
+ )
+
+ nccl_configure(name = "local_config_nccl")
++
++load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
++
++HERMETIC_CC_TOOLCHAIN_VERSION = "v3.1.1"
++
++http_archive(
++    name = "hermetic_cc_toolchain",
++    sha256 = "907745bf91555f77e8234c0b953371e6cac5ba715d1cf12ff641496dd1bce9d1",
++    urls = [
++        "https://mirror.bazel.build/github.com/uber/hermetic_cc_toolchain/releases/download/{0}/hermetic_cc_toolchain-{0}.tar.gz".format(HERMETIC_CC_TOOLCHAIN_VERSION),
++        "https://github.com/uber/hermetic_cc_toolchain/releases/download/{0}/hermetic_cc_toolchain-{0}.tar.gz".format(HERMETIC_CC_TOOLCHAIN_VERSION),
++    ],
++)
++
++load("@hermetic_cc_toolchain//toolchain:defs.bzl", zig_toolchains = "toolchains")
++
++# Plain zig_toolchains() will pick reasonable defaults. See
++# toolchain/defs.bzl:toolchains on how to change the Zig SDK version and
++# download URL.
++zig_toolchains()
++
++register_toolchains(
++    "@zig_sdk//toolchain:linux_amd64_gnu.2.31",
++)
+--
+2.39.5 (Apple Git-154)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -7,3 +7,4 @@ build --repo_env=HERMETIC_PYTHON_VERSION=3.11

		test --test_env LD_LIBRARY_PATH
		test --test_size_filters small,medium