[xla:ffi] Make TypeInfo mandatory in XLA_FFI_REGISTER_TYPE

Add placeholders for future Type serialization/deserialization. It's not an ABI breaking change as it's unused today, and it allows to avoid ABI breaking change in the future when FFI will add proper ser/des support for user defined types.

PiperOrigin-RevId: 820676169

Author: ezhulenev

xla/ffi/api/c_api.h

@@ -272,6 +272,8 @@ typedef struct XLA_FFI_TypeId {
 // to destroy their state when executable is being destroyed.
 typedef struct XLA_FFI_TypeInfo {
   void (*deleter)(void* object);
+  void (*serialize)();    // placeholder for future use
+  void (*deserialize)();  // placeholder for future use
 } XLA_FFI_TypeInfo;
 
 // We use byte spans to pass strings to handlers because strings might not be

xla/ffi/api/ffi.h

@@ -1466,33 +1466,20 @@ struct CtxDecoding<FfiExecutionContext> {
 //===----------------------------------------------------------------------===//
 
 template <typename T>
-XLA_FFI_TypeInfo TypeInfo() {
+constexpr XLA_FFI_TypeInfo TypeInfo() {
   return XLA_FFI_TypeInfo{[](void* ptr) { delete static_cast<T*>(ptr); }};
 }
 
-#define XLA_FFI_REGISTER_TYPE_WITH_INFO(API, NAME, TYPE_ID, TYPE_INFO) \
-  XLA_FFI_REGISTER_TYPE_WITH_INFO_(API, NAME, TYPE_ID, TYPE_INFO, __COUNTER__)
-#define XLA_FFI_REGISTER_TYPE_WITH_INFO_(API, NAME, TYPE_ID, TYPE_INFO, N) \
-  XLA_FFI_REGISTER_TYPE_WITH_INFO__(API, NAME, TYPE_ID, TYPE_INFO, N)
-#define XLA_FFI_REGISTER_TYPE_WITH_INFO__(API, NAME, TYPE_ID, TYPE_INFO, N)    \
+#define XLA_FFI_REGISTER_TYPE(API, NAME, TYPE_ID, TYPE_INFO) \
+  XLA_FFI_REGISTER_TYPE_(API, NAME, TYPE_ID, TYPE_INFO, __COUNTER__)
+#define XLA_FFI_REGISTER_TYPE_(API, NAME, TYPE_ID, TYPE_INFO, N) \
+  XLA_FFI_REGISTER_TYPE__(API, NAME, TYPE_ID, TYPE_INFO, N)
+#define XLA_FFI_REGISTER_TYPE__(API, NAME, TYPE_ID, TYPE_INFO, N)              \
   XLA_FFI_ATTRIBUTE_UNUSED static const XLA_FFI_Error*                         \
       xla_ffi_type_##N##_registered_ = [] {                                    \
         return ::xla::ffi::Ffi::RegisterTypeId(API, NAME, TYPE_ID, TYPE_INFO); \
       }()
 
-#define XLA_FFI_REGISTER_TYPE_X(x, API, NAME, TYPE_ID, TYPE_INFO, FUNC, ...) \
-  FUNC
-
-// Registers external type with XLA runtime and assigns it a unique type id.
-//
-// This is a trick to define macro with optional parameters.
-// Source: https://stackoverflow.com/a/8814003
-#define XLA_FFI_REGISTER_TYPE(API, NAME, TYPE_ID, ...)                  \
-  XLA_FFI_REGISTER_TYPE_X(                                              \
-      , API, NAME, TYPE_ID, ##__VA_ARGS__,                              \
-      XLA_FFI_REGISTER_TYPE_WITH_INFO(API, NAME, TYPE_ID, __VA_ARGS__), \
-      XLA_FFI_REGISTER_TYPE_WITH_INFO(API, NAME, TYPE_ID, {nullptr}))
-
 //===----------------------------------------------------------------------===//
 // UserData
 //===----------------------------------------------------------------------===//

xla/ffi/api/ffi_test.cc

@@ -1202,13 +1202,14 @@ struct MyDataWithExplicitTypeId {
 
 // Rely on XLA to assign unique type id for the type.
 TypeId MyDataWithAutoTypeId::id = XLA_FFI_UNKNOWN_TYPE_ID;
-XLA_FFI_REGISTER_TYPE(GetXlaFfiApi(), "my_data_auto",
-                      &MyDataWithAutoTypeId::id);
+XLA_FFI_REGISTER_TYPE(GetXlaFfiApi(), "my_data_auto", &MyDataWithAutoTypeId::id,
+                      TypeInfo<MyDataWithAutoTypeId>());
 
 // Provide explicit type id and rely on XLA to check that it's unique.
 TypeId MyDataWithExplicitTypeId::id = {42};
 XLA_FFI_REGISTER_TYPE(GetXlaFfiApi(), "my_data_explicit",
-                      &MyDataWithExplicitTypeId::id);
+                      &MyDataWithExplicitTypeId::id,
+                      TypeInfo<MyDataWithExplicitTypeId>());
 
 TEST(FfiTest, UserData) {
   MyDataWithAutoTypeId data0{"foo"};
@@ -1253,7 +1254,8 @@ struct MyState {
 };
 
 TypeId MyState::id = {};  // zero-initialize type id
-XLA_FFI_REGISTER_TYPE(GetXlaFfiApi(), "state", &MyState::id);
+XLA_FFI_REGISTER_TYPE(GetXlaFfiApi(), "state", &MyState::id,
+                      TypeInfo<MyState>());
 
 TEST(FfiTest, StatefulHandler) {
   ExecutionState execution_state;

xla/pjrt/host_callback.cc

@@ -181,6 +181,7 @@ CreateHostCallbackStateAndAppendSendRecvCallbacks(
 // First 64 bits of SHA-512 of "xla::FfiLoadedHostCallbacks".
 ffi::TypeId FfiLoadedHostCallbacks::id = {7357244197867843242};
 XLA_FFI_REGISTER_TYPE(ffi::GetXlaFfiApi(), "FfiLoadedHostCallbacks",
-                      &FfiLoadedHostCallbacks::id);
+                      &FfiLoadedHostCallbacks::id,
+                      ffi::TypeInfo<FfiLoadedHostCallbacks>());
 
 }  // namespace xla