| 54 | } // end namespace |
| 55 | |
| 56 | port::StatusOr<StreamExecutor*> ExecutorCache::GetOrCreate( |
| 57 | const StreamExecutorConfig& config, |
| 58 | const std::function<ExecutorFactory>& factory) { |
| 59 | static int64 cuda_contexts_count = 1; |
| 60 | |
| 61 | std::string key(std::to_string(config.ordinal)); |
| 62 | // Use MPS |
| 63 | // config.virtual_ordinal equal tf_gpu_id, |
| 64 | // Executor key is config.ordinal:config.virtual_ordinal |
| 65 | // For example: 0 -> 0,1 |
| 66 | // 1 -> 2,3,4 |
| 67 | // The executor keys are: 0:0, 0:1, 1:2, 1:3, 1:4 |
| 68 | if (config.virtual_ordinal >= 0) { |
| 69 | absl::call_once(flag_init, &GetCudaContextsCount, |
| 70 | config.ordinal, &cuda_contexts_count); |
| 71 | |
| 72 | int mod_virtual_ordinal = |
| 73 | config.virtual_ordinal % cuda_contexts_count; |
| 74 | key = std::to_string(config.ordinal) + ":" + |
| 75 | std::to_string(mod_virtual_ordinal); |
| 76 | } |
| 77 | // In the fast path case, the cache already has an entry and we can just |
| 78 | // return after Get() which only takes a shared lock and not a unique lock. |
| 79 | // If we need to create, we take a unique lock on cache_. |
| 80 | auto fast_result = Get(config, key); |
| 81 | if (fast_result.ok()) { |
| 82 | return fast_result; |
| 83 | } |
| 84 | |
| 85 | Entry* entry = nullptr; |
| 86 | { |
| 87 | absl::MutexLock lock{&mutex_}; |
| 88 | entry = &cache_[key]; |
| 89 | // Release the map lock; the address of 'entry' is stable because |
| 90 | // std::map guarantees reference stability. |
| 91 | } |
| 92 | |
| 93 | // Acquire the per-Entry mutex without holding the map mutex. Initializing |
| 94 | // an Executor may be expensive, so we want to allow concurrent |
| 95 | // initialization of different entries. |
| 96 | absl::MutexLock lock{&entry->configurations_mutex}; |
| 97 | for (const auto& iter : entry->configurations) { |
| 98 | if (iter.first.plugin_config == config.plugin_config && |
| 99 | iter.first.device_options == config.device_options) { |
| 100 | VLOG(2) << "hit in cache"; |
| 101 | return iter.second.get(); |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | VLOG(2) << "building executor"; |
| 106 | port::StatusOr<std::unique_ptr<StreamExecutor>> result = factory(); |
| 107 | if (!result.ok()) { |
| 108 | VLOG(2) << "failed to get build executor: " << result.status(); |
| 109 | // If construction failed, leave the cache Entry around, but with a null |
| 110 | // executor. |
| 111 | return result.status(); |
| 112 | } |
| 113 | entry->configurations.emplace_back(config, std::move(result.ValueOrDie())); |