diff options
author | 2018-07-09 01:49:04 -0700 | |
---|---|---|
committer | 2018-07-09 01:54:33 -0700 | |
commit | 955e356e4c69d3fce4ac2bac5966671e964f9627 (patch) | |
tree | 95183ec6d5122e5dae704b8a80796f01dfa1b85a /tensorflow/stream_executor | |
parent | caf711b6be448d46354e0cbef23989fa837efcf7 (diff) |
[SE,XLA] Switch to using multiple streams in xla_device_context
Instead of having one stream for compute, host-to-device and device-to-host transfers, switch to having separate streams, just like the GPU does.
Add a se::Event field to XlaTensor to allow accurate inter-stream dependencies to be created.
As part of this:
- Fix TransferManager::TransferLiteralFrom/ToDevice to correctly make generated substreams wait on their master stream.
- Fix Stream::BlockHostUntilDone() to not block on or return substreams. This behavior is completely broken and not only nondeterministically returns substreams to the pool but causes indefinite hangs with the HostStream.
PiperOrigin-RevId: 203726543
Diffstat (limited to 'tensorflow/stream_executor')
-rw-r--r-- | tensorflow/stream_executor/event.cc | 11 | ||||
-rw-r--r-- | tensorflow/stream_executor/event.h | 3 | ||||
-rw-r--r-- | tensorflow/stream_executor/stream.cc | 19 |
3 files changed, 13 insertions, 20 deletions
diff --git a/tensorflow/stream_executor/event.cc b/tensorflow/stream_executor/event.cc index 50a6edd80b..52efe771bc 100644 --- a/tensorflow/stream_executor/event.cc +++ b/tensorflow/stream_executor/event.cc @@ -15,9 +15,9 @@ limitations under the License. #include "tensorflow/stream_executor/event.h" +#include "tensorflow/stream_executor/stream.h" #include "tensorflow/stream_executor/stream_executor_internal.h" #include "tensorflow/stream_executor/stream_executor_pimpl.h" -#include "tensorflow/stream_executor/stream.h" namespace stream_executor { @@ -27,9 +27,12 @@ Event::Event(StreamExecutor* stream_exec) stream_exec_->implementation()->CreateEventImplementation()) {} Event::~Event() { - auto status = stream_exec_->DeallocateEvent(this); - if (!status.ok()) { - LOG(ERROR) << status.error_message(); + // Deal with nullptr implementation_, as this event may have been std::moved. + if (stream_exec_ && implementation_) { + auto status = stream_exec_->DeallocateEvent(this); + if (!status.ok()) { + LOG(ERROR) << status.error_message(); + } } } diff --git a/tensorflow/stream_executor/event.h b/tensorflow/stream_executor/event.h index 1f37262c78..9cc87a7c12 100644 --- a/tensorflow/stream_executor/event.h +++ b/tensorflow/stream_executor/event.h @@ -61,6 +61,9 @@ class Event { // Returns a pointer to the underlying platform-specific implementation. internal::EventInterface* implementation() { return implementation_.get(); } + Event(Event&&) = default; + Event& operator=(Event&&) = default; + private: friend class Stream; diff --git a/tensorflow/stream_executor/stream.cc b/tensorflow/stream_executor/stream.cc index 0cd0790a72..9369183133 100644 --- a/tensorflow/stream_executor/stream.cc +++ b/tensorflow/stream_executor/stream.cc @@ -5228,24 +5228,11 @@ port::Status Stream::BlockHostUntilDone() { return status; } - port::Status first_error; - { - // Wait until all active sub-streams have done their tasks. - mutex_lock lock(mu_); - for (auto &stream : sub_streams_) { - if (!stream.second) { - first_error.Update(stream.first->BlockHostUntilDone()); - // Set this sub-stream as available. - stream.second = true; - } - } - } - temporary_memory_manager_.DeallocateFinalizedTemporaries(); - first_error.Update(parent_->BlockHostUntilDone(this)); - CheckError(first_error.ok()); - return first_error; + port::Status error = parent_->BlockHostUntilDone(this); + CheckError(error.ok()); + return error; } } // namespace stream_executor |