| 22 | namespace oneflow { |
| 23 | |
| 24 | StreamIndexGenerator* TaskStreamIndexManager::GetGenerator(const DeviceId& device_id) { |
| 25 | std::unique_lock<std::mutex> lck(mtx_); |
| 26 | auto iter = generators_.find(device_id); |
| 27 | if (iter == generators_.end()) { |
| 28 | uint32_t init_stream_index = 0; |
| 29 | const int64_t i64_device_id = EncodeDeviceIdToInt64(device_id); |
| 30 | if (stream_index_init_state_.count(i64_device_id) != 0) { |
| 31 | init_stream_index = stream_index_init_state_.at(i64_device_id); |
| 32 | } |
| 33 | iter = generators_.emplace(device_id, std::make_unique<StreamIndexGenerator>(init_stream_index)) |
| 34 | .first; |
| 35 | } |
| 36 | return iter->second.get(); |
| 37 | } |
| 38 | |
| 39 | TaskStreamIndexManager::stream_index_t TaskStreamIndexManager::GetTaskStreamIndex( |
| 40 | TaskType task_type, const DeviceId& device_id) { |
nothing calls this directly
no test coverage detected