@@ -263,15 +263,15 @@ class KernelProgramCache {
263263 }
264264
265265 // Single subcache might be used by different contexts.
266- FastKernelSubcacheMapT &CacheMap = MSubcachePtr->Map ;
266+ // Remove all entries from the subcache that are associated with the
267+ // current context.
268+ FastKernelSubcacheEntriesT &Entries = MSubcachePtr->Entries ;
267269 FastKernelSubcacheWriteLockT Lock{MSubcachePtr->Mutex };
268- for (auto it = CacheMap.begin (); it != CacheMap.end ();) {
269- if (it->first .second == MUrContext) {
270- it = CacheMap.erase (it);
271- } else {
272- ++it;
273- }
274- }
270+ Entries.erase (std::remove_if (Entries.begin (), Entries.end (),
271+ [this ](const FastKernelEntryT &Entry) {
272+ return Entry.Key .second == MUrContext;
273+ }),
274+ Entries.end ());
275275 }
276276
277277 FastKernelSubcacheT &get () { return *MSubcachePtr; }
@@ -476,14 +476,21 @@ class KernelProgramCache {
476476 KernelSubcacheHint = &It.first ->second .get ();
477477 }
478478
479- const FastKernelSubcacheMapT &SubcacheMap = KernelSubcacheHint->Map ;
479+ const FastKernelSubcacheEntriesT &SubcacheEntries =
480+ KernelSubcacheHint->Entries ;
480481 FastKernelSubcacheReadLockT SubcacheLock{KernelSubcacheHint->Mutex };
481482 ur_context_handle_t Context = getURContext ();
482- auto It = SubcacheMap.find (FastKernelCacheKeyT (Device, Context));
483- if (It != SubcacheMap.end ()) {
483+ const FastKernelCacheKeyT RequiredKey (Device, Context);
484+ // Search for the kernel in the subcache.
485+ auto It = std::find_if (SubcacheEntries.begin (), SubcacheEntries.end (),
486+ [&](const FastKernelEntryT &Entry) {
487+ return Entry.Key == RequiredKey;
488+ });
489+ if (It != SubcacheEntries.end ()) {
484490 traceKernel (" Kernel fetched." , KernelName, true );
485- return It->second ;
491+ return It->Value ;
486492 }
493+
487494 return FastKernelCacheValPtr ();
488495 }
489496
@@ -516,8 +523,8 @@ class KernelProgramCache {
516523
517524 FastKernelSubcacheWriteLockT SubcacheLock{KernelSubcacheHint->Mutex };
518525 ur_context_handle_t Context = getURContext ();
519- KernelSubcacheHint->Map . emplace ( FastKernelCacheKeyT (Device, Context),
520- CacheVal);
526+ KernelSubcacheHint->Entries . emplace_back (
527+ FastKernelCacheKeyT (Device, Context), CacheVal);
521528 }
522529
523530 // Expects locked program cache
@@ -562,15 +569,21 @@ class KernelProgramCache {
562569 bool RemoveSubcache = false ;
563570 {
564571 FastKernelSubcacheWriteLockT SubcacheLock{Subcache.Mutex };
565- Subcache.Map .erase (
566- FastKernelCacheKeyT (FastCacheKey.second , Context));
572+ Subcache.Entries .erase (
573+ std::remove_if (
574+ Subcache.Entries .begin (), Subcache.Entries .end (),
575+ [&](const FastKernelEntryT &Entry) {
576+ return Entry.Key == FastKernelCacheKeyT (
577+ FastCacheKey.second , Context);
578+ }),
579+ Subcache.Entries .end ());
567580 traceKernel (" Kernel evicted." , FastCacheKey.first , true );
568581
569582 // Remove the subcache wrapper from this kernel program cache if
570583 // the subcache no longer contains entries for this context.
571584 RemoveSubcache = std::none_of (
572- Subcache.Map .begin (), Subcache.Map .end (),
573- [&](const auto &It) { return It.first .second == Context; });
585+ Subcache.Entries .begin (), Subcache.Entries .end (),
586+ [&](const auto &It) { return It.Key .second == Context; });
574587 }
575588 if (RemoveSubcache)
576589 MFastKernelCache.erase (FastKernelCacheItr);
0 commit comments