mirror of
https://github.com/zebrajr/tensorflow.git
synced 2025-12-07 12:20:24 +01:00
[XLA] Guard stream pool with mutex.
A data race can occur while populating the map. PiperOrigin-RevId: 157647183
This commit is contained in:
parent
ccdb30763a
commit
414470329b
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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_;
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user