[XLA] Guard stream pool with mutex.

A data race can occur while populating the map.

PiperOrigin-RevId: 157647183
This commit is contained in:
Jacques Pienaar 2017-05-31 15:53:17 -07:00 committed by TensorFlower Gardener
parent ccdb30763a
commit 414470329b
2 changed files with 4 additions and 1 deletions

View File

@ -118,6 +118,7 @@ StatusOr<Backend::StreamPtr> Backend::BorrowStream(int device_ordinal) {
StatusOr<Backend::StreamPtr> Backend::BorrowStream(
se::StreamExecutor* executor) {
tensorflow::mutex_lock l(mu_);
if (0 == stream_pools_.count(executor)) {
stream_pools_.emplace(std::piecewise_construct,
std::forward_as_tuple(executor),

View File

@ -190,10 +190,12 @@ class Backend {
// Vector of stream executors. stream_executors_[0] is the default executor.
std::vector<perftools::gputools::StreamExecutor*> stream_executors_;
tensorflow::mutex mu_;
// Mapping from stream executor to stream pools, used by `BorrowStream` above.
std::map<perftools::gputools::StreamExecutor*,
Pool<perftools::gputools::Stream>>
stream_pools_;
stream_pools_ GUARDED_BY(mu_);
// The default memory allocator to use.
std::unique_ptr<StreamExecutorMemoryAllocator> memory_allocator_;