Skip to content

Commit 84116c1

Browse files
ezhulenevcopybara-github
authored andcommitted
[xla:ffi] NFC: Remove function name from XLA_FFI binding
FFI module itself defines under what name FFI binding should be exported to the Xla runtime and ffi binding name was not really used anywhere. PiperOrigin-RevId: 501955340
1 parent 13a941a commit 84116c1

File tree

5 files changed

+14
-23
lines changed

5 files changed

+14
-23
lines changed

xla/runtime/ffi/ffi_api.h

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -146,12 +146,11 @@ class Ffi {
146146
public:
147147
virtual ~Ffi() = default;
148148

149-
virtual std::string_view name() const = 0;
150149
virtual XLA_FFI_Error* operator()(const XLA_FFI_Api* api,
151150
XLA_FFI_ExecutionContext* ctx, void** args,
152151
void** attrs, void** rets) const = 0;
153152

154-
static FfiBinding<> Bind(std::string name);
153+
static FfiBinding<> Binding();
155154

156155
template <typename T>
157156
static bool Isa(const XLA_FFI_Api* api, XLA_FFI_TypeId type_id);
@@ -510,32 +509,29 @@ class FfiBinding {
510509

511510
template <typename Fn>
512511
std::unique_ptr<FfiHandler<Fn, Ts...>> To(Fn fn) {
513-
return std::unique_ptr<FfiHandler<Fn, Ts...>>(new FfiHandler<Fn, Ts...>(
514-
std::forward<Fn>(fn), std::move(name_), std::move(attrs_)));
512+
return std::unique_ptr<FfiHandler<Fn, Ts...>>(
513+
new FfiHandler<Fn, Ts...>(std::forward<Fn>(fn), std::move(attrs_)));
515514
}
516515

517516
private:
518517
template <typename...>
519518
friend class FfiBinding;
520519
friend class Ffi;
521520

522-
explicit FfiBinding(std::string name) : name_(std::move(name)) {
521+
explicit FfiBinding() {
523522
static_assert(sizeof...(Ts) == 0, "ffi arguments must be empty");
524523
}
525524

526525
template <typename... TTs>
527526
FfiBinding(FfiBinding<TTs...>&& other) // NOLINT
528-
: name_(std::move(other.name_)), attrs_(std::move(other.attrs_)) {}
527+
: attrs_(std::move(other.attrs_)) {}
529528

530529
FfiBinding(FfiBinding&) = delete;
531530

532-
std::string name_; // ffi name
533531
std::vector<std::string> attrs_; // names of bound attributes
534532
};
535533

536-
inline FfiBinding<> Ffi::Bind(std::string name) {
537-
return FfiBinding<>(std::move(name));
538-
}
534+
inline FfiBinding<> Ffi::Binding() { return FfiBinding<>(); }
539535

540536
//===----------------------------------------------------------------------===//
541537
// Helpers for decoding opaque arguments and attributes' memory.
@@ -798,8 +794,6 @@ class FfiHandler : public Ffi {
798794
}
799795

800796
public:
801-
std::string_view name() const final { return name_; }
802-
803797
XLA_FFI_Error* operator()(const XLA_FFI_Api* api,
804798
XLA_FFI_ExecutionContext* ctx, void** args,
805799
void** attrs, void** rets) const final {
@@ -867,9 +861,8 @@ class FfiHandler : public Ffi {
867861
return ToError(api, FfiStatus::Ok());
868862
}
869863

870-
FfiHandler(Fn fn, std::string name, std::vector<std::string> attrs)
864+
FfiHandler(Fn fn, std::vector<std::string> attrs)
871865
: fn_(std::move(fn)),
872-
name_(std::move(name)),
873866
attrs_(std::move(attrs)),
874867
attrs_idx_(attrs_.size()) {
875868
// Sort attributes names.
@@ -886,7 +879,6 @@ class FfiHandler : public Ffi {
886879

887880
Fn fn_;
888881

889-
std::string name_;
890882
std::vector<std::string> attrs_;
891883

892884
// A mapping from the attribute index to its index in the lexicographically

xla/runtime/ffi_test.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ struct TestModule : public ffi::StatefulModule<TestModuleState> {
156156
// Function that tests that we can successfully decode various kinds of
157157
// attributes attached to custom calls.
158158
XLA_FFI_DEFINE_FUNCTION(FFI_AttrsDecoding, AttrsDecoding,
159-
ffi::Ffi::Bind("ffi.attrs")
159+
ffi::Ffi::Binding()
160160
.State<TestModuleState>() // state
161161
.Attr<std::string_view>("str")
162162
.Attr<float>("f32")
@@ -173,7 +173,7 @@ struct TestModule : public ffi::StatefulModule<TestModuleState> {
173173
// Function that tests that we can successfully decode various kinds of
174174
// arguments passed to custom calls.
175175
XLA_FFI_DEFINE_FUNCTION(FFI_Fill, Fill,
176-
ffi::Ffi::Bind("ffi.fill")
176+
ffi::Ffi::Binding()
177177
.State<TestModuleState>() // state
178178
.Arg<int32_t>() // arg0
179179
.Arg<ffi::BufferArg>() // arg1

xla/service/gpu/custom_call_test.cc

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -417,12 +417,11 @@ struct TestFfiModule : ffi::StatelessModule {
417417
{{"ffi.always_fail", FFI_AlwaysFail}, {"ffi.memcpy", FFI_Memcpy}}) {
418418
}
419419

420-
XLA_FFI_DEFINE_FUNCTION(
421-
FFI_AlwaysFail, AlwaysFail,
422-
ffi::Ffi::Bind("ffi.always_fail").Arg<ffi::StridedBufferArg>());
420+
XLA_FFI_DEFINE_FUNCTION(FFI_AlwaysFail, AlwaysFail,
421+
ffi::Ffi::Binding().Arg<ffi::StridedBufferArg>());
423422

424423
XLA_FFI_DEFINE_FUNCTION(FFI_Memcpy, Memcpy,
425-
ffi::Ffi::Bind("ffi.memcpy")
424+
ffi::Ffi::Binding()
426425
.Stream<se::gpu::GpuStreamHandle>()
427426
.Arg<ffi::StridedBufferArg>()
428427
.Arg<ffi::StridedBufferArg>());

xla/tests/custom_call_test.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,7 @@ struct TestFfiModule : ffi::StatelessModule {
371371
{{"ffi.add_const", FFI_AddConst}}) {}
372372

373373
XLA_FFI_DEFINE_FUNCTION(FFI_AddConst, AddConst,
374-
ffi::Ffi::Bind("ffi.add_const")
374+
ffi::Ffi::Binding()
375375
.Arg<ffi::StridedBufferArg>()
376376
.Arg<ffi::StridedBufferArg>()
377377
.Attr<float>("cst"));

xla/tests/xla_ffi_test.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ struct TestModule : public runtime::ffi::StatefulModule<TestModuleState> {
5757

5858
// XLA runtime binding for the C++ function.
5959
XLA_FFI_DEFINE_FUNCTION(FFI_Impl, Impl,
60-
Ffi::Bind("test.ffi")
60+
Ffi::Binding()
6161
.State<TestModuleState>()
6262
.Arg<StridedBufferArg>()
6363
.Arg<StridedBufferArg>()

0 commit comments

Comments
 (0)