| 203 | } |
| 204 | |
| 205 | Status XRTCompilationCache::CompileIfKeyAbsent( |
| 206 | const string& key, int64* uid, |
| 207 | const std::function<Status(std::unique_ptr<xla::LocalExecutable>*)>& |
| 208 | compile_function) { |
| 209 | CompiledSubgraph* entry = nullptr; |
| 210 | |
| 211 | absl::MutexLock lock(&mu_); |
| 212 | auto iter = cache_.find(key); |
| 213 | |
| 214 | if (iter == cache_.end()) { |
| 215 | // The single ref on the newly-created entry is owned by the caller. |
| 216 | VLOG(1) << "Before adding new entry for key " << key << " cache is " |
| 217 | << cache_.size() << " entries (" |
| 218 | << cache_entries_ + marked_for_eviction_entries_ << "), " |
| 219 | << " marked for eviction " |
| 220 | << (cache_.size() - entries_by_last_use_.size()) << " entries (" |
| 221 | << marked_for_eviction_entries_ << ")."; |
| 222 | entry = InitializeEntry(key, compile_function); |
| 223 | } else { |
| 224 | VLOG(1) << "Before refreshing entry for key " << key << " cache is " |
| 225 | << cache_.size() << " entries (" |
| 226 | << cache_entries_ + marked_for_eviction_entries_ << "), " |
| 227 | << " marked for eviction " |
| 228 | << (cache_.size() - entries_by_last_use_.size()) << " entries (" |
| 229 | << marked_for_eviction_entries_ << ")."; |
| 230 | entry = iter->second; |
| 231 | // Make a new reference that is owned by the caller. |
| 232 | entry->Ref(); |
| 233 | // Block if necessary until the subgraph has been initialized. |
| 234 | mu_.Await(absl::Condition( |
| 235 | +[](CompiledSubgraph* e) { return e->initialized; }, entry)); |
| 236 | } |
| 237 | |
| 238 | // Let the caller know the uid of the entry. |
| 239 | *uid = entry->uid; |
| 240 | |
| 241 | // Remove the old LRU-table entry if it wasn't already marked for eviction. |
| 242 | auto erased = entries_by_last_use_.erase(entry->last_use); |
| 243 | // Update the LRU table indicating this entry is the most recently used. |
| 244 | entry->last_use = use_counter_++; |
| 245 | entries_by_last_use_[entry->last_use] = entry; |
| 246 | if (erased == 0) { |
| 247 | // The entry had been marked for eviction, or is newly created. |
| 248 | LookupEntryMarkedForEviction(entry); |
| 249 | } |
| 250 | |
| 251 | VLOG(1) << "After refreshing entry for key " << key << " cache is " |
| 252 | << cache_.size() << " entries (" |
| 253 | << cache_entries_ + marked_for_eviction_entries_ << "), " |
| 254 | << " marked for eviction " |
| 255 | << (cache_.size() - entries_by_last_use_.size()) << " entries (" |
| 256 | << marked_for_eviction_entries_ << ")."; |
| 257 | |
| 258 | return entry->initialization_status; |
| 259 | } |
| 260 | |
| 261 | Status XRTCompilationCache::Lookup( |
| 262 | int64 uid, std::unique_ptr<XRTCompilationCacheEntryRef>* entry) { |