Add BlockHostUntilDoneWithStatus, which returns Status rather than bool.

Also fixed a deadlock in Stream::BlockHostUntilDone.  The problem with the
original code was that it grabbed mu_ before looping over substreams, and would
call CheckError with mu_ still held.  But CheckError will attempt to lock mu_ in
the failure case, which would deadlock.

PiperOrigin-RevId: 178191634

1 files changed, 9 insertions, 6 deletions

diff --git a/tensorflow/stream_executor/stream_executor_pimpl.cc b/tensorflow/stream_executor/stream_executor_pimpl.cc
index 76afb85068..5630255b5d 100644
--- a/tensorflow/stream_executor/stream_executor_pimpl.cc
+++ b/tensorflow/stream_executor/stream_executor_pimpl.cc
@@ -432,12 +432,15 @@ bool StreamExecutor::Launch(Stream *stream, const ThreadDim &thread_dims,
   return implementation_->Launch(stream, thread_dims, block_dims, kernel, args);
 }
 
-bool StreamExecutor::BlockHostUntilDone(Stream *stream) {
-  bool result;
-  SCOPED_TRACE(TraceListener::BlockHostUntilDone, &result, stream);
-
-  result = implementation_->BlockHostUntilDone(stream);
-  return result;
+port::Status StreamExecutor::BlockHostUntilDoneWithStatus(Stream *stream) {
+  // TODO(toddw): Change TraceListener::BlockHostUntilDone to record Status
+  // rather than bool.
+  bool trace_result;
+  SCOPED_TRACE(TraceListener::BlockHostUntilDone, &trace_result, stream);
+
+  port::Status status = implementation_->BlockHostUntilDoneWithStatus(stream);
+  trace_result = status.ok();
+  return status;
 }
 
 void *StreamExecutor::Allocate(uint64 size) {