Implement DoHostCallbackWithStatus to allow callbacks to return a status

PiperOrigin-RevId: 207714420

1 files changed, 5 insertions, 0 deletions

diff --git a/tensorflow/stream_executor/stream_executor_pimpl.h b/tensorflow/stream_executor/stream_executor_pimpl.h
index 47b3a2b030..437f298616 100644
--- a/tensorflow/stream_executor/stream_executor_pimpl.h
+++ b/tensorflow/stream_executor/stream_executor_pimpl.h
@@ -549,6 +549,11 @@ class StreamExecutor {
   // See Stream::ThenDoHostCallback for full details.
   bool HostCallback(Stream *stream, std::function<void()> callback);
 
+  // Entrains on a stream a user-specified function to be run on the host.
+  // See Stream::ThenDoHostCallback for full details.
+  // This is the preferred form for a callback that may return an error.
+  bool HostCallback(Stream *stream, std::function<port::Status()> callback);
+
   // Performs platform-specific allocation and initialization of an event.
   port::Status AllocateEvent(Event *event);