openxla: patch PJRT client should_stage_host_to_device_transfers (#14)

Author: hugomano

third_party/openxla/patches/20250128-001-PJRT-Expose-should_stage_host_to_device_transfers.patch

@@ -0,0 +1,46 @@
+From 8b54dc4825628b088ea47e1dda7189477a0989fb Mon Sep 17 00:00:00 2001
+From: Hugo Mano <hugo@zml.ai>
+Date: Tue, 28 Jan 2025 16:15:15 +0100
+Subject: [PATCH] [PJRT]  Expose should_stage_host_to_device_transfers  as
+ create option
+
+Expose GPU option `should_stage_host_to_device_transfers `
+as configuration in PJRT client create func.
+---
+ xla/pjrt/c/pjrt_c_api_gpu_internal.cc | 7 +++++++
+ 1 file changed, 7 insertions(+)
+
+diff --git a/xla/pjrt/c/pjrt_c_api_gpu_internal.cc b/xla/pjrt/c/pjrt_c_api_gpu_internal.cc
+index cf23c997fb..b3157b59b1 100644
+--- a/xla/pjrt/c/pjrt_c_api_gpu_internal.cc
++++ b/xla/pjrt/c/pjrt_c_api_gpu_internal.cc
+@@ -84,6 +84,7 @@ PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) {
+           {"visible_devices", PJRT_NamedValue_Type::PJRT_NamedValue_kInt64List},
+           {"node_id", PJRT_NamedValue_Type::PJRT_NamedValue_kInt64},
+           {"num_nodes", PJRT_NamedValue_Type::PJRT_NamedValue_kInt64},
++          {"should_stage_host_to_device_transfers", PJRT_NamedValue_Type::PJRT_NamedValue_kBool},
+           {"enable_mock_nccl", PJRT_NamedValue_Type::PJRT_NamedValue_kBool},
+           {"mock_gpu_topology", PJRT_NamedValue_Type::PJRT_NamedValue_kString},
+       });
+@@ -139,6 +140,11 @@ PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) {
+   if (auto it = create_options.find("num_nodes"); it != create_options.end()) {
+     num_nodes = std::get<int64_t>(it->second);
+   }
++  bool should_stage_host_to_device_transfers = true;
++  if (auto it = create_options.find("should_stage_host_to_device_transfers");
++      it != create_options.end()) {
++    should_stage_host_to_device_transfers = std::get<bool>(it->second);
++  }
+   bool enable_mock_nccl = false;
+   if (auto it = create_options.find("enable_mock_nccl");
+       it != create_options.end()) {
+@@ -159,6 +165,7 @@ PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) {
+   options.kv_store = pjrt::ToCppKeyValueStore(
+       args->kv_get_callback, args->kv_get_user_arg, args->kv_try_get_callback,
+       args->kv_try_get_user_arg, args->kv_put_callback, args->kv_put_user_arg);
++  options.should_stage_host_to_device_transfers = should_stage_host_to_device_transfers;
+   options.enable_mock_nccl = enable_mock_nccl;
+   options.mock_gpu_topology = mock_gpu_topology;
+   PJRT_ASSIGN_OR_RETURN(std::unique_ptr<xla::PjRtClient> client,
+--
+2.39.5 (Apple Git-154)

third_party/openxla/xla.bzl

@@ -30,6 +30,7 @@ filegroup(
             "patches/20240901-003-Only-export-GetPjrtApi-symbol-on-macOS.patch", # PR: https://github.com/openxla/xla/pull/16696
             "patches/20250120-001-Enable-nvptxcompiler-with-nvjitlink.patch", # Allow us to levarage technologies flagged for Google only
             "patches/20250122-001-Fix-LoadedNvJitLinkHasKnownIssues-check.patch", # PR: https://github.com/openxla/xla/pull/2172
+            "patches/20250128-001-PJRT-Expose-should_stage_host_to_device_transfers.patch", # PR: https://github.com/openxla/xla/pull/21965
         ],
     )