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.h b/tensorflow/stream_executor/stream.h
index 4d41409fef..e1629b5b30 100644
--- a/tensorflow/stream_executor/stream.h
+++ b/tensorflow/stream_executor/stream.h
@@ -2045,6 +2045,11 @@ class Stream {
   // negative effects on performance.
   Stream &ThenDoHostCallback(std::function<void()> callback);
 
+  // Entrains onto the stream a callback to the host (from the device).
+  // Behaves as ThenDoHostCallback above, but returns a Status instead of void.
+  // This overload should be preferred if the callback could fail.
+  Stream &ThenDoHostCallbackWithStatus(std::function<port::Status()> callback);
+
   // Returns the StreamExecutor (parent object) associated with this stream.
   StreamExecutor *parent() const {
     CHECK(parent_ != nullptr);