[xla:ffi] Add benchmarks for async FFI calls

-----------------------------------------------------------------
Benchmark                       Time             CPU   Iterations
-----------------------------------------------------------------
BM_AnyBufferArgX1            18.2 ns         18.2 ns     37983521
BM_AnyBufferArgX4            22.7 ns         22.7 ns     30682729
BM_AnyBufferArgX8            27.9 ns         27.9 ns     24817074
BM_AsyncAnyBufferArgX1       55.2 ns         55.2 ns     12514333

PiperOrigin-RevId: 674515693

xla/ffi/api/ffi.h

@@ -1019,9 +1019,6 @@ template <ExecutionStage stage>
 struct ResultEncoding<stage, Future> {
   static std::variant<XLA_FFI_Error*, XLA_FFI_Future*> Encode(
       const XLA_FFI_Api* api, XLA_FFI_ExecutionContext* ctx, Future future) {
-    // TODO(ezhulenev): Add benchmarks for asynchronous FFI handlers, and
-    // optimize the fast path for returning completed futures.
-
     // Create XLA_FFI_Future object that will signal completion to the runtime.
     XLA_FFI_Future_Create_Args args;
     args.struct_size = XLA_FFI_Future_Create_Args_STRUCT_SIZE;

xla/ffi/api/ffi_test.cc

@@ -1307,6 +1307,27 @@ void BM_AnyBufferArgX4(benchmark::State& state) {
 
 BENCHMARK(BM_AnyBufferArgX4);
 
+//===----------------------------------------------------------------------===//
+// BM_AsyncAnyBufferArgX1
+//===----------------------------------------------------------------------===//
+
+void BM_AsyncAnyBufferArgX1(benchmark::State& state) {
+  auto call_frame = WithBufferArgs(1).Build();
+
+  auto handler = Ffi::Bind().Arg<AnyBuffer>().To([](auto buffer) {
+    benchmark::DoNotOptimize(buffer);
+    Promise promise;
+    promise.SetAvailable();
+    return Future(promise);
+  });
+
+  for (auto _ : state) {
+    CHECK_OK(Call(*handler, call_frame));
+  }
+}
+
+BENCHMARK(BM_AsyncAnyBufferArgX1);
+
 //===----------------------------------------------------------------------===//
 // BM_BufferArgX1
 //===----------------------------------------------------------------------===//

xla/ffi/ffi_test.cc

@@ -1168,6 +1168,26 @@ void BM_AnyBufferArgX8(benchmark::State& state) {
 
 BENCHMARK(BM_AnyBufferArgX8);
 
+//===----------------------------------------------------------------------===//
+// BM_AsyncAnyBufferArgX1
+//===----------------------------------------------------------------------===//
+
+void BM_AsyncAnyBufferArgX1(benchmark::State& state) {
+  auto call_frame = WithBufferArgs(1).Build();
+
+  auto done = tsl::MakeAvailableAsyncValueRef<tsl::Chain>();
+  auto handler = Ffi::Bind().Arg<AnyBuffer>().To([&](auto buffer) {
+    benchmark::DoNotOptimize(buffer);
+    return done;
+  });
+
+  for (auto _ : state) {
+    CHECK_OK(Call(*handler, call_frame));
+  }
+}
+
+BENCHMARK(BM_AsyncAnyBufferArgX1);
+
 //===----------------------------------------------------------------------===//
 // BM_BufferArgX1
 //===----------------------------------------------------------------------===//