[pjrt] Add an API to add user data to FFI context via PJRT_ExecuteCon…

…text

PiperOrigin-RevId: 636260703

third_party/xla/xla/pjrt/c/BUILD

@@ -34,6 +34,30 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "pjrt_c_api_ffi_extension_hdrs",
+    hdrs = ["pjrt_c_api_ffi_extension.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":pjrt_c_api_hdrs",
+    ],
+)
+
+cc_library(
+    name = "pjrt_c_api_ffi_internal",
+    srcs = ["pjrt_c_api_ffi_internal.cc"],
+    hdrs = ["pjrt_c_api_ffi_internal.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":pjrt_c_api_ffi_extension_hdrs",
+        ":pjrt_c_api_hdrs",
+        ":pjrt_c_api_helpers",
+        ":pjrt_c_api_wrapper_impl",
+        "//xla/ffi:execution_context",
+        "@com_google_absl//absl/status",
+    ],
+)
+
 cc_library(
     name = "pjrt_c_api_gpu_extension_hdrs",
     hdrs = ["pjrt_c_api_gpu_extension.h"],
@@ -218,6 +242,8 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":pjrt_c_api_custom_partitioner_extension_hdrs",
+        ":pjrt_c_api_ffi_extension_hdrs",
+        ":pjrt_c_api_ffi_internal",
         ":pjrt_c_api_gpu_extension_hdrs",
         ":pjrt_c_api_hdrs",
         ":pjrt_c_api_helpers",
@@ -321,6 +347,7 @@ xla_cc_test(
     srcs = ["pjrt_c_api_gpu_test.cc"],
     use_gpu = True,
     deps = [
+        ":pjrt_c_api_ffi_extension_hdrs",
         ":pjrt_c_api_gpu",
         ":pjrt_c_api_gpu_extension_hdrs",
         ":pjrt_c_api_hdrs",
@@ -333,6 +360,7 @@ xla_cc_test(
         "//xla:shape_util",
         "//xla:status",
         "//xla:statusor",
+        "//xla/ffi:execution_context",
         "//xla/ffi:ffi_api",
         "//xla/ffi/api:ffi",
         "//xla/pjrt:pjrt_client",

third_party/xla/xla/pjrt/c/CHANGELOG.md

@@ -1,5 +1,10 @@
 # PJRT C API changelog
 
+
+## 0.53
+* Added ``PJRT_FFI_Extension` extension to support passing user data to FFI
+  handlers on compatible PJRT backends.
+
 ## 0.52
 * Added ``PJRT_ExecuteContext`` struct corresponding to ``xla::ExecuteContext``.
 

third_party/xla/xla/pjrt/c/pjrt_c_api.h

@@ -44,6 +44,7 @@ typedef enum {
   PJRT_Extension_Type_Custom_Partitioner,
   PJRT_Extension_Type_Stream,
   PJRT_Extension_Type_Layouts,
+  PJRT_Extension_Type_FFI,
 } PJRT_Extension_Type;
 
 // PJRT_Extension_Base contains a type and a pointer to next
@@ -78,7 +79,7 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Extension_Base, next);
 // Changes include:
 // * Adding a new field to the PJRT_Api or argument structs
 // * Renaming a method or argument (doesn't affect ABI)
-#define PJRT_API_MINOR 52
+#define PJRT_API_MINOR 53
 
 // The plugin should set the major_version and minor_version of
 // PJRT_Api.pjrt_api_version to be the `PJRT_API_MAJOR` and `PJRT_API_MINOR` in

third_party/xla/xla/pjrt/c/pjrt_c_api_ffi_extension.h

@@ -0,0 +1,68 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_PJRT_C_PJRT_C_API_FFI_EXTENSION_H_
+#define XLA_PJRT_C_PJRT_C_API_FFI_EXTENSION_H_
+
+#include <stddef.h>
+
+#include <cstdint>
+
+#include "xla/pjrt/c/pjrt_c_api.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// PJRT FFI extension provides capabilities for integrating with a
+// backend-specific FFI (foreign function interface) library, i.e. for XLA CPU
+// and GPU backends it gives access to the XLA FFI internals.
+//
+// See: https://en.wikipedia.org/wiki/Foreign_function_interface
+#define PJRT_API_FFI_EXTENSION_VERSION 1
+
+// User-data that will be forwarded to the FFI handlers. Deleter is optional,
+// and can be nullptr. Deleter will be called when the context is destroyed.
+struct PJRT_FFI_UserData {
+  int64_t type_id;
+  void* data;
+  void (*deleter)(void* data);
+};
+
+struct PJRT_FFI_UserData_Add_Args {
+  size_t struct_size;
+  PJRT_Extension_Base* extension_start;
+
+  PJRT_ExecuteContext* context;
+  PJRT_FFI_UserData user_data;
+};
+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);
+
+typedef struct PJRT_FFI_Extension {
+  size_t struct_size;
+  PJRT_Extension_Type type;
+  PJRT_Extension_Base* next;
+  PJRT_FFI_UserData_Add* user_data_add;
+} PJRT_FFI;
+PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_Extension, user_data_add);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // XLA_PJRT_C_PJRT_C_API_FFI_EXTENSION_H_

third_party/xla/xla/pjrt/c/pjrt_c_api_ffi_internal.cc

@@ -0,0 +1,52 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+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 "absl/status/status.h"
+#include "xla/ffi/execution_context.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"
+
+namespace pjrt {
+
+static PJRT_Error* PJRT_FFI_UserData_Add(PJRT_FFI_UserData_Add_Args* args) {
+  PJRT_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
+      "PJRT_FFI_UserData_Add_Args", PJRT_FFI_UserData_Add_Args_STRUCT_SIZE,
+      args->struct_size));
+
+  if (args->context == nullptr) {
+    return new PJRT_Error{absl::InvalidArgumentError(
+        "PJRT FFI extension requires execute context to be not nullptr")};
+  }
+
+  xla::ffi::ExecutionContext::TypeId type_id(args->user_data.type_id);
+  PJRT_RETURN_IF_ERROR(args->context->execute_context->ffi_context().Insert(
+      type_id, args->user_data.data, args->user_data.deleter));
+  return nullptr;
+}
+
+PJRT_FFI_Extension CreateFfiExtension(PJRT_Extension_Base* next) {
+  return {
+      /*struct_size=*/PJRT_FFI_Extension_STRUCT_SIZE,
+      /*type=*/PJRT_Extension_Type::PJRT_Extension_Type_FFI,
+      /*next=*/next,
+      /*custom_call=*/PJRT_FFI_UserData_Add,
+  };
+}
+
+}  // namespace pjrt

third_party/xla/xla/pjrt/c/pjrt_c_api_ffi_internal.h

@@ -0,0 +1,28 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_PJRT_C_PJRT_C_API_FFI_INTERNAL_H_
+#define XLA_PJRT_C_PJRT_C_API_FFI_INTERNAL_H_
+
+#include "xla/pjrt/c/pjrt_c_api.h"
+#include "xla/pjrt/c/pjrt_c_api_ffi_extension.h"
+
+namespace pjrt {
+
+PJRT_FFI_Extension CreateFfiExtension(PJRT_Extension_Base* next);
+
+}  // namespace pjrt
+
+#endif  // XLA_PJRT_C_PJRT_C_API_FFI_INTERNAL_H_

third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc

@@ -33,6 +33,8 @@ limitations under the License.
 #include "xla/ffi/ffi_api.h"
 #include "xla/pjrt/c/pjrt_c_api.h"
 #include "xla/pjrt/c/pjrt_c_api_custom_partitioner_extension.h"
+#include "xla/pjrt/c/pjrt_c_api_ffi_extension.h"
+#include "xla/pjrt/c/pjrt_c_api_ffi_internal.h"
 #include "xla/pjrt/c/pjrt_c_api_gpu_extension.h"
 #include "xla/pjrt/c/pjrt_c_api_helpers.h"
 #include "xla/pjrt/c/pjrt_c_api_layouts_extension.h"
@@ -299,15 +301,20 @@ const PJRT_Api* GetGpuPjrtApi() {
       /*next=*/reinterpret_cast<PJRT_Extension_Base*>(&stream),
       /*custom_call=*/PJRT_Gpu_Register_Custom_Call,
   };
+
   static PJRT_Layouts_Extension layouts_extension =
       pjrt::CreateLayoutsExtension(
           reinterpret_cast<PJRT_Extension_Base*>(&custom_call));
+
+  static PJRT_FFI_Extension ffi_extension = pjrt::CreateFfiExtension(
+      reinterpret_cast<PJRT_Extension_Base*>(&layouts_extension));
+
   static const PJRT_Api pjrt_api = pjrt::CreatePjrtApi(
       pjrt::gpu_plugin::PJRT_Client_Create,
       pjrt::gpu_plugin::PJRT_ExecuteContext_Create,
       pjrt::gpu_plugin::PJRT_GpuDeviceTopology_Create,
       pjrt::PJRT_Plugin_Initialize_NoOp,
-      reinterpret_cast<PJRT_Extension_Base*>(&layouts_extension),
+      reinterpret_cast<PJRT_Extension_Base*>(&ffi_extension),
       pjrt::PJRT_Plugin_Attributes_Xla);
 
   return &pjrt_api;

third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc

@@ -31,10 +31,12 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "xla/ffi/api/ffi.h"
+#include "xla/ffi/execution_context.h"
 #include "xla/ffi/ffi_api.h"
 #include "xla/literal.h"
 #include "xla/literal_util.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_gpu_extension.h"
 #include "xla/pjrt/c/pjrt_c_api_helpers.h"
 #include "xla/pjrt/c/pjrt_c_api_test.h"
@@ -158,10 +160,31 @@ TEST_F(PjrtCApiGpuTest, CreateAndDestroyExecuteContext) {
   create_arg.extension_start = nullptr;
   create_arg.context = nullptr;
 
-  PJRT_Error* error = api_->PJRT_ExecuteContext_Create(&create_arg);
-  EXPECT_EQ(error, nullptr);
+  EXPECT_EQ(api_->PJRT_ExecuteContext_Create(&create_arg), nullptr);
   EXPECT_NE(create_arg.context, nullptr);
 
+  const PJRT_FFI_Extension* ffi_extension =
+      pjrt::FindExtension<PJRT_FFI_Extension>(
+          api_, PJRT_Extension_Type::PJRT_Extension_Type_FFI);
+  ASSERT_NE(ffi_extension, nullptr);
+
+  std::string string_data = "string_data";
+
+  PJRT_FFI_UserData_Add_Args add_args;
+  add_args.struct_size = PJRT_FFI_UserData_Add_Args_STRUCT_SIZE;
+  add_args.extension_start = nullptr;
+  add_args.user_data.type_id = 42;
+  add_args.user_data.data = &string_data;
+  add_args.user_data.deleter = nullptr;
+  add_args.context = create_arg.context;
+  EXPECT_EQ(ffi_extension->user_data_add(&add_args), nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto lookup_user_data,
+      create_arg.context->execute_context->ffi_context().Lookup(
+          xla::ffi::ExecutionContext::TypeId(42)));
+  EXPECT_EQ(lookup_user_data, &string_data);
+
   PJRT_ExecuteContext_Destroy_Args destroy_args;
   destroy_args.struct_size = PJRT_Error_Destroy_Args_STRUCT_SIZE;
   destroy_args.extension_start = nullptr;