[xla:ffi] Add TypeRegistry::TypeInfo to be able to register functions to manipulate user-defined types

PiperOrigin-RevId: 820811829

Author: ezhulenev

xla/ffi/BUILD

@@ -62,7 +62,7 @@ cc_library(
     srcs = ["execution_context.cc"],
     hdrs = ["execution_context.h"],
     deps = [
-        ":type_id_registry",
+        ":type_registry",
         "//xla:util",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:logging",
@@ -79,7 +79,7 @@ xla_cc_test(
     srcs = ["execution_context_test.cc"],
     deps = [
         ":execution_context",
-        ":type_id_registry",
+        ":type_registry",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",
@@ -94,13 +94,14 @@ cc_library(
     srcs = ["execution_state.cc"],
     hdrs = ["execution_state.h"],
     deps = [
-        ":type_id_registry",
+        ":type_registry",
         "//xla:util",
+        "//xla/tsl/platform:statusor",
+        "//xla/tsl/util:safe_reinterpret_cast",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
-        "@tsl//tsl/platform:logging",
-        "@tsl//tsl/platform:statusor",
     ],
 )
 
@@ -109,11 +110,12 @@ xla_cc_test(
     srcs = ["execution_state_test.cc"],
     deps = [
         ":execution_state",
+        ":type_registry",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:statusor",
+        "//xla/tsl/platform:test",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/platform:test",
     ],
 )
 
@@ -124,7 +126,7 @@ cc_library(
         ":api",
         ":execution_context",
         ":execution_state",
-        ":type_id_registry",
+        ":type_registry",
         "//xla:executable_run_options",
         "//xla:shape_util",
         "//xla:types",
@@ -141,6 +143,7 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
@@ -160,7 +163,7 @@ cc_library(
         ":call_frame",
         ":execution_context",
         ":execution_state",
-        ":type_id_registry",
+        ":type_registry",
         "//xla:executable_run_options",
         "//xla:util",
         "//xla/ffi/api:c_api",
@@ -218,7 +221,7 @@ xla_cc_test(
         ":execution_state",
         ":ffi",
         ":ffi_api",
-        ":type_id_registry",
+        ":type_registry",
         "//xla:executable_run_options",
         "//xla:xla_data_proto_cc",
         "//xla/ffi/api:c_api",
@@ -243,27 +246,31 @@ xla_cc_test(
 )
 
 cc_library(
-    name = "type_id_registry",
-    srcs = ["type_id_registry.cc"],
-    hdrs = ["type_id_registry.h"],
+    name = "type_registry",
+    srcs = ["type_registry.cc"],
+    hdrs = ["type_registry.h"],
     deps = [
         "//xla:util",
         "//xla/tsl/lib/gtl:int_type",
+        "//xla/tsl/util:safe_reinterpret_cast",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",
     ],
 )
 
 xla_cc_test(
-    name = "type_id_registry_test",
-    srcs = ["type_id_registry_test.cc"],
+    name = "type_registry_test",
+    srcs = ["type_registry_test.cc"],
     deps = [
-        ":type_id_registry",
+        ":type_registry",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",

xla/ffi/api/BUILD

@@ -89,7 +89,7 @@ xla_cc_test(
         "//xla/ffi:execution_context",
         "//xla/ffi:execution_state",
         "//xla/ffi:ffi_api",
-        "//xla/ffi:type_id_registry",
+        "//xla/ffi:type_registry",
         "//xla/stream_executor:device_memory",
         "//xla/stream_executor:device_memory_allocator",
         "//xla/tsl/concurrency:async_value",

xla/ffi/api/api.h

@@ -348,11 +348,13 @@ inline XLA_FFI_Error* Ffi::RegisterTypeId(const XLA_FFI_Api* api,
                                           std::string_view name,
                                           XLA_FFI_TypeId* type_id,
                                           XLA_FFI_TypeInfo type_info) {
+  assert(type_id && "type_id must not be null");
   XLA_FFI_TypeId_Register_Args args;
   args.struct_size = XLA_FFI_TypeId_Register_Args_STRUCT_SIZE;
   args.extension_start = nullptr;
   args.name = XLA_FFI_ByteSpan{name.data(), name.size()};
   args.type_id = type_id;
+  args.type_info = &type_info;
   return api->XLA_FFI_TypeId_Register(&args);
 }
 

xla/ffi/api/c_api.h

@@ -67,7 +67,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Extension_Base, next);
 // * Deleting a method or argument
 // * Changing the type of an argument
 // * Rearranging fields in the XLA_FFI_Api or argument structs
-#define XLA_FFI_API_MAJOR 0
+#define XLA_FFI_API_MAJOR 1
 
 // Incremented when the interface is updated in a way that is potentially
 // ABI-compatible with older versions, if supported by the caller and/or
@@ -82,7 +82,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Extension_Base, next);
 // Minor changes include:
 // * Adding a new field to the XLA_FFI_Api or argument structs
 // * Renaming a method or argument (doesn't affect ABI)
-#define XLA_FFI_API_MINOR 1
+#define XLA_FFI_API_MINOR 0
 
 struct XLA_FFI_Api_Version {
   size_t struct_size;
@@ -491,6 +491,7 @@ struct XLA_FFI_TypeId_Register_Args {
 
   XLA_FFI_ByteSpan name;
   XLA_FFI_TypeId* type_id;  // in-out
+  XLA_FFI_TypeInfo* type_info;
 };
 
 XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_TypeId_Register_Args, type_id);

xla/ffi/api/ffi.h

@@ -437,7 +437,7 @@ class CountDownPromise {
     assert(state_->count.load() >= count && "Invalid count down value");
 
     if (XLA_FFI_PREDICT_FALSE(!error.success())) {
-      const std::lock_guard<std::mutex> lock(state_->mutex);
+      std::lock_guard<std::mutex> lock(state_->mutex);  // NOLINT
       state_->is_error.store(true, std::memory_order_release);
       state_->error = error;
     }
@@ -448,7 +448,7 @@ class CountDownPromise {
       bool is_error = state_->is_error.load(std::memory_order_acquire);
       if (XLA_FFI_PREDICT_FALSE(is_error)) {
         auto take_error = [&] {
-          const std::lock_guard<std::mutex> lock(state_->mutex);
+          std::lock_guard<std::mutex> lock(state_->mutex);  // NOLINT
           return state_->error;
         };
         state_->promise.SetError(take_error());
@@ -476,7 +476,7 @@ class CountDownPromise {
     std::atomic<int64_t> count;
     std::atomic<bool> is_error;
 
-    std::mutex mutex;
+    std::mutex mutex;  // NOLINT
     Error error;
   };
 

xla/ffi/api/ffi_test.cc

@@ -39,7 +39,7 @@ limitations under the License.
 #include "xla/ffi/execution_context.h"
 #include "xla/ffi/execution_state.h"
 #include "xla/ffi/ffi_api.h"
-#include "xla/ffi/type_id_registry.h"
+#include "xla/ffi/type_registry.h"
 #include "xla/primitive_util.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/device_memory_allocator.h"
@@ -1220,9 +1220,9 @@ TEST(FfiTest, UserData) {
 
   ExecutionContext execution_context;
   TF_ASSERT_OK(execution_context.Insert(
-      TypeIdRegistry::TypeId(MyDataWithAutoTypeId::id.type_id), &data0));
+      TypeRegistry::TypeId(MyDataWithAutoTypeId::id.type_id), &data0));
   TF_ASSERT_OK(execution_context.Insert(
-      TypeIdRegistry::TypeId(MyDataWithExplicitTypeId::id.type_id), &data1));
+      TypeRegistry::TypeId(MyDataWithExplicitTypeId::id.type_id), &data1));
 
   CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0);
   auto call_frame = builder.Build();

xla/ffi/execution_context.h

@@ -25,7 +25,7 @@ limitations under the License.
 #include "absl/functional/function_ref.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "xla/ffi/type_id_registry.h"
+#include "xla/ffi/type_registry.h"
 #include "xla/tsl/platform/logging.h"
 #include "xla/tsl/platform/statusor.h"
 
@@ -45,7 +45,7 @@ namespace xla::ffi {
 // unique between separate calls to XLA execute.
 class ExecutionContext {
  public:
-  using TypeId = TypeIdRegistry::TypeId;
+  using TypeId = TypeRegistry::TypeId;
 
   template <typename T>
   using Deleter = std::function<void(T*)>;
@@ -67,7 +67,7 @@ class ExecutionContext {
   template <typename T>
   absl::StatusOr<T*> Lookup() const {
     TF_ASSIGN_OR_RETURN(auto user_data,
-                        LookupUserData(TypeIdRegistry::GetTypeId<T>()));
+                        LookupUserData(TypeRegistry::GetTypeId<T>()));
     return static_cast<T*>(user_data->data());
   }
 
@@ -110,7 +110,7 @@ class ExecutionContext {
 
 template <typename T>
 absl::Status ExecutionContext::Insert(T* data, Deleter<T> deleter) {
-  return InsertUserData(TypeIdRegistry::GetTypeId<T>(),
+  return InsertUserData(TypeRegistry::GetTypeId<T>(),
                         std::make_unique<UserData>(
                             data, [deleter = std::move(deleter)](void* data) {
                               if (deleter) deleter(static_cast<T*>(data));
@@ -119,7 +119,7 @@ absl::Status ExecutionContext::Insert(T* data, Deleter<T> deleter) {
 
 template <typename T, typename... Args>
 absl::Status ExecutionContext::Emplace(Args&&... args) {
-  return InsertUserData(TypeIdRegistry::GetTypeId<T>(),
+  return InsertUserData(TypeRegistry::GetTypeId<T>(),
                         std::make_unique<UserData>(
                             new T(std::forward<Args>(args)...),
                             [](void* data) { delete static_cast<T*>(data); }));

xla/ffi/execution_context_test.cc

@@ -20,7 +20,7 @@ limitations under the License.
 
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
-#include "xla/ffi/type_id_registry.h"
+#include "xla/ffi/type_registry.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/test.h"
@@ -62,8 +62,9 @@ TEST(ExecutionContextTest, InsertUserOwned) {
 }
 
 TEST(ExecutionContextTest, InsertUserOwnedWithTypeId) {
-  TF_ASSERT_OK_AND_ASSIGN(TypeIdRegistry::TypeId type_id,
-                          TypeIdRegistry::AssignExternalTypeId("I32UserData"));
+  TF_ASSERT_OK_AND_ASSIGN(TypeRegistry::TypeId type_id,
+                          TypeRegistry::AssignExternalTypeId(
+                              "I32UserData", TypeRegistry::TypeInfo{}));
 
   I32UserData user_data(42);
 

xla/ffi/execution_state.cc

@@ -15,46 +15,62 @@ limitations under the License.
 
 #include "xla/ffi/execution_state.h"
 
-#include <utility>
-
+#include "absl/base/attributes.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "xla/ffi/type_id_registry.h"
+#include "xla/ffi/type_registry.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
-#include "tsl/platform/logging.h"
 
 namespace xla::ffi {
 
 ExecutionState::ExecutionState()
-    : type_id_(TypeIdRegistry::kUnknownTypeId),
-      state_(nullptr),
-      deleter_(nullptr) {}
+    : type_id_(TypeRegistry::kUnknownTypeId), state_(nullptr) {}
 
 ExecutionState::~ExecutionState() {
-  if (deleter_) deleter_(state_);
+  if (type_info_.deleter) {
+    type_info_.deleter(state_);
+  }
+}
+
+absl::Status ExecutionState::Set(TypeId type_id, void* state) {
+  TF_ASSIGN_OR_RETURN(auto type_info,
+                      TypeRegistry::GetExternalTypeInfo(type_id));
+  if (type_info.deleter == nullptr) {
+    return InvalidArgument(
+        "Type id %d does not have a registered type info with a deleter",
+        type_id.value());
+  }
+  return Set(type_id, type_info, state);
 }
 
+ABSL_DEPRECATED("FFI users must rely in TypeInfo registration")
 absl::Status ExecutionState::Set(TypeId type_id, void* state,
-                                 Deleter<void> deleter) {
-  DCHECK(state && deleter) << "State and deleter must not be null";
+                                 void (*deleter)(void*)) {
+  return Set(type_id, TypeInfo{deleter}, state);
+}
+
+absl::Status ExecutionState::Set(TypeId type_id, TypeInfo type_info,
+                                 void* state) {
+  DCHECK(state && type_info.deleter) << "State and deleter must not be null";
 
-  if (type_id_ != TypeIdRegistry::kUnknownTypeId) {
+  if (type_id_ != TypeRegistry::kUnknownTypeId) {
     return FailedPrecondition("State is already set with a type id %d",
                               type_id_.value());
   }
 
   type_id_ = type_id;
+  type_info_ = type_info;
   state_ = state;
-  deleter_ = std::move(deleter);
 
   return absl::OkStatus();
 }
 
 // Returns opaque state of the given type id. If set state type id does not
 // match the requested one, returns an error.
 absl::StatusOr<void*> ExecutionState::Get(TypeId type_id) const {
-  if (type_id_ == TypeIdRegistry::kUnknownTypeId) {
+  if (type_id_ == TypeRegistry::kUnknownTypeId) {
     return NotFound("State is not set");
   }
 
@@ -68,7 +84,7 @@ absl::StatusOr<void*> ExecutionState::Get(TypeId type_id) const {
 }
 
 bool ExecutionState::IsSet() const {
-  return type_id_ != TypeIdRegistry::kUnknownTypeId;
+  return type_id_ != TypeRegistry::kUnknownTypeId;
 }
 
 }  // namespace xla::ffi