[xla:ffi] Allow registering type ids assigned by the user

ezhulenev · Google-ML-Automation · commit 4dd28604913f · 2025-03-20T14:34:13.000-07:00
PiperOrigin-RevId: 738942202
diff --git a/xla/ffi/BUILD b/xla/ffi/BUILD
@@ -81,11 +81,11 @@ xla_cc_test(
         ":execution_context",
         ":type_id_registry",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:statusor",
+        "//xla/tsl/platform:test",
         "@com_google_absl//absl/status",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/platform:test",
     ],
 )
 
@@ -247,9 +247,12 @@ cc_library(
     deps = [
         "//xla:util",
         "//xla/tsl/lib/gtl:int_type",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",
     ],
 )
@@ -260,10 +263,10 @@ xla_cc_test(
     deps = [
         ":type_id_registry",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:statusor",
+        "//xla/tsl/platform:test",
         "@com_google_absl//absl/status",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/platform:test",
     ],
 )
diff --git a/xla/ffi/api/c_api.h b/xla/ffi/api/c_api.h
@@ -468,17 +468,22 @@ typedef XLA_FFI_Error* XLA_FFI_Handler_Register(
 // TypeId
 //===----------------------------------------------------------------------===//
 
+#define XLA_FFI_UNKNOWN_TYPE_ID XLA_FFI_TypeId{0}
+
 struct XLA_FFI_TypeId_Register_Args {
   size_t struct_size;
   XLA_FFI_Extension_Base* extension_start;
 
   XLA_FFI_ByteSpan name;
-  XLA_FFI_TypeId* type_id;  // out
+  XLA_FFI_TypeId* type_id;  // in-out
 };
 
 XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_TypeId_Register_Args, type_id);
 
-// Registers user type `name` and returns a unique `type_id`.
+// Registers user type `name` with XLA. If type id is `XLA_FFI_UNKNOWN_TYPE_ID`,
+// XLA will assign a unique type id and return it in `type_id` out argument,
+// otherwise XLA will verify that type id is unique and matches the type id of
+// the type registered with the same `name` earlier.
 typedef XLA_FFI_Error* XLA_FFI_TypeId_Register(
     XLA_FFI_TypeId_Register_Args* args);
 
diff --git a/xla/ffi/execution_context_test.cc b/xla/ffi/execution_context_test.cc
@@ -22,8 +22,8 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "xla/ffi/type_id_registry.h"
 #include "xla/tsl/lib/core/status_test_util.h"
-#include "tsl/platform/statusor.h"
-#include "tsl/platform/test.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/test.h"
 
 namespace xla::ffi {
 
@@ -62,9 +62,8 @@ TEST(ExecutionContextTest, InsertUserOwned) {
 }
 
 TEST(ExecutionContextTest, InsertUserOwnedWithTypeId) {
-  TF_ASSERT_OK_AND_ASSIGN(
-      TypeIdRegistry::TypeId type_id,
-      TypeIdRegistry::RegisterExternalTypeId("I32UserData"));
+  TF_ASSERT_OK_AND_ASSIGN(TypeIdRegistry::TypeId type_id,
+                          TypeIdRegistry::AssignExternalTypeId("I32UserData"));
 
   I32UserData user_data(42);
 
diff --git a/xla/ffi/ffi_api.cc b/xla/ffi/ffi_api.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
 #include "xla/executable_run_options.h"
 #include "xla/ffi/api/api.h"
 #include "xla/ffi/api/c_api.h"
@@ -635,13 +636,27 @@ static XLA_FFI_Error* XLA_FFI_TypeId_Register(
       "XLA_FFI_ExecutionContext_Get_Args",
       XLA_FFI_ExecutionContext_Get_Args_STRUCT_SIZE, args->struct_size));
 
-  auto type_id = TypeIdRegistry::RegisterExternalTypeId(
-      std::string_view(args->name.ptr, args->name.len));
-  if (!type_id.ok()) {
-    return new XLA_FFI_Error{std::move(type_id).status()};
+  absl::string_view type_name(args->name.ptr, args->name.len);
+  TypeIdRegistry::TypeId type_id(args->type_id->type_id);
+
+  // If type_id is unknown, we are registering a new type and XLA will assign a
+  // unique type id to it.
+  if (type_id == TypeIdRegistry::kUnknownTypeId) {
+    auto assigned_type_id = TypeIdRegistry::AssignExternalTypeId(type_name);
+    if (!assigned_type_id.ok()) {
+      return new XLA_FFI_Error{std::move(assigned_type_id).status()};
+    }
+
+    args->type_id->type_id = assigned_type_id->value();
+    return nullptr;
+  }
+
+  // If type_id is set, we are relying on the caller-provided unique type id.
+  if (auto status = TypeIdRegistry::RegisterExternalTypeId(type_name, type_id);
+      !status.ok()) {
+    return new XLA_FFI_Error{std::move(status)};
   }
 
-  args->type_id->type_id = type_id->value();
   return nullptr;
 }
 
diff --git a/xla/ffi/type_id_registry.cc b/xla/ffi/type_id_registry.cc
@@ -20,10 +20,13 @@ limitations under the License.
 #include <string>
 #include <string_view>
 
+#include "absl/algorithm/container.h"
 #include "absl/base/attributes.h"
 #include "absl/base/const_init.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "xla/util.h"
 
@@ -39,24 +42,56 @@ static ExternalTypeIdRegistry& StaticExternalTypeIdRegistry() {
   return *registry;
 }
 
-TypeIdRegistry::TypeId TypeIdRegistry::GetNextTypeId() {
+TypeIdRegistry::TypeId TypeIdRegistry::GetNextInternalTypeId() {
   static auto* counter = new std::atomic<int64_t>(1);
   return TypeId(counter->fetch_add(1));
 }
 
-absl::StatusOr<TypeIdRegistry::TypeId> TypeIdRegistry::RegisterExternalTypeId(
+TypeIdRegistry::TypeId TypeIdRegistry::GetNextExternalTypeId() {
+  static auto* counter = new std::atomic<int64_t>(1);
+  return TypeId(counter->fetch_add(1));
+}
+
+absl::StatusOr<TypeIdRegistry::TypeId> TypeIdRegistry::AssignExternalTypeId(
     std::string_view name) {
   absl::MutexLock lock(&type_registry_mutex);
   auto& registry = StaticExternalTypeIdRegistry();
 
-  // Try to emplace with type id zero and fill it with real type id only if we
+  // Try to emplace with unknow type id and fill it with real type id only if we
   // successfully acquired an entry for a given name.
-  auto emplaced = registry.emplace(name, TypeId(0));
+  auto emplaced = registry.emplace(name, kUnknownTypeId);
   if (!emplaced.second) {
-    return Internal("Type id %d already registered for type name %s",
-                    emplaced.first->second.value(), name);
+    return Internal("Type name %s already registered with type id %d", name,
+                    emplaced.first->second.value());
   }
-  return emplaced.first->second = GetNextTypeId();
+
+  // Returns true if the registry contains an entry with a given type id.
+  auto type_id_is_in_use = [&registry](TypeId type_id) {
+    return absl::c_any_of(registry,
+                          [&](const auto& e) { return e.second == type_id; });
+  };
+
+  // Create a new type id that is not already in use.
+  TypeId type_id = GetNextExternalTypeId();
+  while (type_id_is_in_use(type_id)) {
+    type_id = GetNextExternalTypeId();
+  }
+
+  return emplaced.first->second = type_id;
+}
+
+absl::Status TypeIdRegistry::RegisterExternalTypeId(absl::string_view name,
+                                                    TypeId type_id) {
+  absl::MutexLock lock(&type_registry_mutex);
+  auto& registry = StaticExternalTypeIdRegistry();
+
+  auto emplaced = registry.emplace(name, type_id);
+  if (!emplaced.second && emplaced.first->second != type_id) {
+    return Internal("Type name %s already registered with type id %d vs %d)",
+                    name, emplaced.first->second.value(), type_id.value());
+  }
+
+  return absl::OkStatus();
 }
 
 }  // namespace xla::ffi
diff --git a/xla/ffi/type_id_registry.h b/xla/ffi/type_id_registry.h
@@ -17,9 +17,10 @@ limitations under the License.
 #define XLA_FFI_TYPE_ID_REGISTRY_H_
 
 #include <cstdint>
-#include <string_view>
 
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "xla/tsl/lib/gtl/int_type.h"
 
 namespace xla::ffi {
@@ -48,20 +49,30 @@ class TypeIdRegistry {
 
   static constexpr TypeId kUnknownTypeId = TypeId(0);
 
-  // Registers external type with a given name in a static type registry.
-  static absl::StatusOr<TypeId> RegisterExternalTypeId(std::string_view name);
+  // Assigns a unique type id to an external type with a given name. Returns an
+  // error if a type with a given name is already registered in the process.
+  static absl::StatusOr<TypeId> AssignExternalTypeId(absl::string_view name);
+
+  // Registers external type with a given name and type id. Type id is provided
+  // by the caller, and must be unique. Returns an error if a type with a given
+  // name is already registered with a different type id.
+  static absl::Status RegisterExternalTypeId(absl::string_view name,
+                                             TypeId type_id);
 
   // Returns a type id for a given type. For internal type ids only.
   template <typename T>
   static TypeId GetTypeId();
 
  private:
-  static TypeId GetNextTypeId();
+  // We never mix external and internal type ids, so we can use different type
+  // id spaces to assign unique ids to each type.
+  static TypeId GetNextInternalTypeId();
+  static TypeId GetNextExternalTypeId();
 };
 
 template <typename T>
 TypeIdRegistry::TypeId TypeIdRegistry::GetTypeId() {
-  static const TypeId id = GetNextTypeId();
+  static const TypeId id = GetNextInternalTypeId();
   return id;
 }
 
diff --git a/xla/ffi/type_id_registry_test.cc b/xla/ffi/type_id_registry_test.cc
@@ -16,12 +16,14 @@ limitations under the License.
 #include "xla/ffi/type_id_registry.h"
 
 #include <cstdint>
+#include <limits>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
-#include "tsl/platform/statusor.h"
-#include "tsl/platform/test.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/test.h"
 
 namespace xla::ffi {
 namespace {
@@ -30,12 +32,25 @@ using ::testing::HasSubstr;
 
 TEST(TypeIdRegistryTest, RegisterExternalTypeId) {
   TF_ASSERT_OK_AND_ASSIGN(auto type_id,
-                          TypeIdRegistry::RegisterExternalTypeId("foo"));
+                          TypeIdRegistry::AssignExternalTypeId("foo"));
   EXPECT_GE(type_id.value(), 0);
 
-  auto duplicate_type_id = TypeIdRegistry::RegisterExternalTypeId("foo");
+  auto duplicate_type_id = TypeIdRegistry::AssignExternalTypeId("foo");
   EXPECT_THAT(duplicate_type_id.status().message(),
-              HasSubstr("already registered for type name foo"));
+              HasSubstr("Type name foo already registered with type id"));
+
+  // It's ok to register the same type with same type id.
+  TF_ASSERT_OK(TypeIdRegistry::RegisterExternalTypeId("foo", type_id));
+
+  // It's an error to register the same type with a different type id.
+  auto wrong_type_id = TypeIdRegistry::RegisterExternalTypeId(
+      "foo", TypeIdRegistry::TypeId(std::numeric_limits<int64_t>::max()));
+  EXPECT_THAT(wrong_type_id.message(),
+              HasSubstr("Type name foo already registered with type id"));
+
+  // It's ok to register a new type with a user-provided type id.
+  TF_ASSERT_OK(TypeIdRegistry::RegisterExternalTypeId(
+      "bar", TypeIdRegistry::TypeId(std::numeric_limits<int64_t>::max())));
 }
 
 TEST(TypeIdRegistryTest, RegisterInternalTypeId) {
diff --git a/xla/pjrt/c/BUILD b/xla/pjrt/c/BUILD
@@ -70,6 +70,7 @@ cc_library(
         "//xla/ffi:execution_context",
         "//xla/ffi:type_id_registry",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
     ],
 )
 
diff --git a/xla/pjrt/c/pjrt_c_api_ffi_extension.h b/xla/pjrt/c/pjrt_c_api_ffi_extension.h
@@ -38,11 +38,14 @@ struct PJRT_FFI_TypeID_Register_Args {
 
   const char* type_name;
   size_t type_name_size;
-  int64_t type_id;  // out
+  int64_t type_id;  // in-out
 };
 PJRT_DEFINE_STRUCT_TRAITS(PJRT_FFI_TypeID_Register_Args, type_id);
 
-// Registers external type in a static type registry.
+// Registers external type in a static type registry. If `type_id` is set to `0`
+// XLA will assign a unique type id to it and return via out argument, otherwise
+// it will verify that user-provided type id matches previously registered type
+// id for the given type name.
 typedef PJRT_Error* PJRT_FFI_TypeID_Register(
     PJRT_FFI_TypeID_Register_Args* args);
 
diff --git a/xla/pjrt/c/pjrt_c_api_ffi_internal.cc b/xla/pjrt/c/pjrt_c_api_ffi_internal.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/pjrt/c/pjrt_c_api_ffi_internal.h"
 
 #include "absl/status/status.h"
+#include "absl/strings/string_view.h"
 #include "xla/ffi/execution_context.h"
 #include "xla/ffi/type_id_registry.h"
 #include "xla/pjrt/c/pjrt_c_api.h"
@@ -31,11 +32,23 @@ static PJRT_Error* PJRT_FFI_TypeID_Register(
       "PJRT_FFI_TypeID_Register_Args",
       PJRT_FFI_TypeID_Register_Args_STRUCT_SIZE, args->struct_size));
 
-  PJRT_ASSIGN_OR_RETURN(
-      auto type_id,
-      xla::ffi::TypeIdRegistry::RegisterExternalTypeId(
-          absl::string_view(args->type_name, args->type_name_size)));
-  args->type_id = type_id.value();
+  absl::string_view type_name(args->type_name, args->type_name_size);
+  xla::ffi::TypeIdRegistry::TypeId type_id(args->type_id);
+
+  if (type_id == xla::ffi::TypeIdRegistry::kUnknownTypeId) {
+    // If type_id is unknown, we are registering a new type and XLA will assign
+    // a unique type id to it.
+    PJRT_ASSIGN_OR_RETURN(
+        auto assigned_type_id,
+        xla::ffi::TypeIdRegistry::AssignExternalTypeId(type_name));
+    args->type_id = assigned_type_id.value();
+
+  } else {
+    // If type_id is set, we are relying on the caller-provided unique type id.
+    PJRT_RETURN_IF_ERROR(
+        xla::ffi::TypeIdRegistry::RegisterExternalTypeId(type_name, type_id));
+  }
+
   return nullptr;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -70,6 +70,7 @@ cc_library(`
`70`	`70`	`"//xla/ffi:execution_context",`
`71`	`71`	`"//xla/ffi:type_id_registry",`
`72`	`72`	`"@com_google_absl//absl/status",`
	`73`	`+ "@com_google_absl//absl/strings:string_view",`
`73`	`74`	`],`
`74`	`75`	`)`
`75`	`76`