| 157 | } |
| 158 | |
| 159 | XRTCompilationCache::CompiledSubgraph* XRTCompilationCache::InitializeEntry( |
| 160 | const string& key, |
| 161 | const std::function<Status(std::unique_ptr<xla::LocalExecutable>*)>& |
| 162 | initialize_program) { |
| 163 | CompiledSubgraph* entry = new CompiledSubgraph(); |
| 164 | entry->parent = this; |
| 165 | entry->key = key; |
| 166 | entry->uid = get_uid(); |
| 167 | // Add the entry to the cache. Once the computation has been compiled, |
| 168 | // UpdateEntryAfterCompilation will be called to potentially mark old entries |
| 169 | // that don't fit any more for eviction. |
| 170 | // |
| 171 | // At this point there is one reference to entry, which is owned by the caller |
| 172 | // who created the entry. A second reference, owned by the cache, will be |
| 173 | // added below since we leave the entry in the 'marked for eviction' state |
| 174 | // here. |
| 175 | auto cache_inserted = |
| 176 | cache_.insert(std::pair<string, CompiledSubgraph*>(key, entry)); |
| 177 | CHECK(cache_inserted.second); |
| 178 | |
| 179 | // Initialize the program outside the lock so that other cache operations |
| 180 | // can proceed during the (potentially lengthy) initialization. |
| 181 | Status s; |
| 182 | std::unique_ptr<xla::LocalExecutable> program; |
| 183 | { |
| 184 | mu_.Unlock(); |
| 185 | { s = initialize_program(&program); } |
| 186 | mu_.Lock(); |
| 187 | } |
| 188 | |
| 189 | // Add the entry to the uid index. |
| 190 | auto uid_inserted = entries_by_uid_.insert( |
| 191 | std::pair<int64, CompiledSubgraph*>(entry->uid, entry)); |
| 192 | CHECK(uid_inserted.second); |
| 193 | |
| 194 | entry->initialized = true; |
| 195 | entry->initialization_status = s; |
| 196 | if (s.ok()) { |
| 197 | entry->program = std::move(program); |
| 198 | } |
| 199 | // Add the entry to marked_for_eviction_entries_ since it will be adjusted |
| 200 | // down again when the newly-created entry gets unmarked. |
| 201 | ++marked_for_eviction_entries_; |
| 202 | return entry; |
| 203 | } |
| 204 | |
| 205 | Status XRTCompilationCache::CompileIfKeyAbsent( |
| 206 | const string& key, int64* uid, |