OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 00:19:58 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

[xla:ffi] Make TypeInfo mandatory in XLA_FFI_REGISTER_TYPE

Browse Source 
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
This commit is contained in:
Eugene Zhulenev 2025-10-17 07:09:12 -07:00 committed by TensorFlower Gardener
parent 30d25d6d18
commit 4752801386
4 changed files with 16 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
third_party/xla/xla/ffi/api/c_api.h vendored
View File

@ -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

25
third_party/xla/xla/ffi/api/ffi.h vendored
View File

@ -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
//===----------------------------------------------------------------------===//

10
third_party/xla/xla/ffi/api/ffi_test.cc vendored
View File

@ -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;

3
third_party/xla/xla/pjrt/host_callback.cc vendored
View File

@ -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