Skip to content

Commit

Permalink
[SYCL] Extend image cleanup for __sycl_unregister_lib (#17091)
Browse files Browse the repository at this point in the history 
fixes #16031

---------

Signed-off-by: Tikhomirova, Kseniya <[email protected]>
  • Loading branch information
@KseniyaTikhomirova
KseniyaTikhomirova authored Feb 27, 2025
1 parent 42ff9c2 commit 1e0c021
Show file tree
Hide file tree
Showing 9 changed files with 610 additions and 138 deletions.
174 changes: 100 additions & 74 deletions sycl/source/detail/kernel_program_cache.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,11 @@ class KernelProgramCache {
MProgramEvictionList.pop_front();
}
}

void erase(const ProgramCacheKeyT &CacheKey) {
MProgramEvictionList.remove(CacheKey);
MProgramToEvictionListMap.erase(CacheKey);
}
};

~KernelProgramCache() = default;
Expand Down Expand Up @@ -427,31 +432,100 @@ class KernelProgramCache {

template <typename KeyT, typename ValT>
void saveKernel(KeyT &&CacheKey, ValT &&CacheVal) {

ur_program_handle_t Program = std::get<3>(CacheVal);
if (SYCLConfig<SYCL_IN_MEM_CACHE_EVICTION_THRESHOLD>::
isProgramCacheEvictionEnabled()) {

ur_program_handle_t Program = std::get<3>(CacheVal);
// Save kernel in fast cache only if the corresponding program is also
// in the cache.
auto LockedCache = acquireCachedPrograms();
auto &ProgCache = LockedCache.get();
if (ProgCache.ProgramSizeMap.find(Program) ==
ProgCache.ProgramSizeMap.end())
return;

// Save reference between the program and the fast cache key.
std::unique_lock<std::mutex> Lock(MKernelFastCacheMutex);
MProgramToKernelFastCacheKeyMap[Program].emplace_back(CacheKey);
}

// Save reference between the program and the fast cache key.
std::unique_lock<std::mutex> Lock(MKernelFastCacheMutex);
MProgramToKernelFastCacheKeyMap[Program].emplace_back(CacheKey);

// if no insertion took place, thus some other thread has already inserted
// smth in the cache
traceKernel("Kernel inserted.", CacheKey.second, true);
MKernelFastCache.emplace(CacheKey, CacheVal);
}

// Expects locked program cache
size_t removeProgramByKey(const ProgramCacheKeyT &CacheKey,
ProgramCache &ProgCache) {
auto It = ProgCache.Cache.find(CacheKey);

if (It != ProgCache.Cache.end()) {
// We are about to remove this program now.
// (1) Remove it from KernelPerProgram cache.
// (2) Remove corresponding entries from KernelFastCache.
// (3) Remove it from ProgramCache KeyMap.
// (4) Remove it from the ProgramCache.
// (5) Remove it from ProgramSizeMap.
// (6) Update the cache size.

// Remove entry from the KernelsPerProgram cache.
ur_program_handle_t NativePrg = It->second->Val;
{
auto LockedCacheKP = acquireKernelsPerProgramCache();
// List kernels that are to be removed from the cache, if tracing is
// enabled.
if (SYCLConfig<SYCL_CACHE_TRACE>::isTraceInMemCache()) {
for (const auto &Kernel : LockedCacheKP.get()[NativePrg])
traceKernel("Kernel evicted.", Kernel.first);
}
LockedCacheKP.get().erase(NativePrg);
}

{
// Remove corresponding entries from KernelFastCache.
std::unique_lock<std::mutex> Lock(MKernelFastCacheMutex);
if (auto FastCacheKeyItr =
MProgramToKernelFastCacheKeyMap.find(NativePrg);
FastCacheKeyItr != MProgramToKernelFastCacheKeyMap.end()) {
for (const auto &FastCacheKey : FastCacheKeyItr->second) {
MKernelFastCache.erase(FastCacheKey);
traceKernel("Kernel evicted.", FastCacheKey.second, true);
}
MProgramToKernelFastCacheKeyMap.erase(FastCacheKeyItr);
}
}

// Remove entry from ProgramCache KeyMap.
CommonProgramKeyT CommonKey =
std::make_pair(CacheKey.first.second, CacheKey.second);
// Since KeyMap is a multi-map, we need to iterate over all entries
// with this CommonKey and remove those that match the CacheKey.
auto KeyMapItrRange = ProgCache.KeyMap.equal_range(CommonKey);
for (auto KeyMapItr = KeyMapItrRange.first;
KeyMapItr != KeyMapItrRange.second; ++KeyMapItr) {
if (KeyMapItr->second == CacheKey) {
ProgCache.KeyMap.erase(KeyMapItr);
break;
}
}

// Get size of the program.
size_t ProgramSize = MCachedPrograms.ProgramSizeMap[It->second->Val];
// Evict program from the cache.
ProgCache.Cache.erase(It);
// Remove program size from the cache size.
MCachedPrograms.ProgramCacheSizeInBytes -= ProgramSize;
MCachedPrograms.ProgramSizeMap.erase(NativePrg);

traceProgram("Program evicted.", CacheKey);
} else
// This should never happen.
throw sycl::exception(sycl::make_error_code(sycl::errc::runtime),
"Program not found in the cache.");

return MCachedPrograms.ProgramCacheSizeInBytes;
}

// Evict programs from cache to free up space.
void evictPrograms(size_t DesiredCacheSize, size_t CurrentCacheSize) {

Expand All @@ -471,73 +545,7 @@ class KernelProgramCache {
ProgramCacheKeyT CacheKey = ProgramEvictionList.front();
auto LockedCache = acquireCachedPrograms();
auto &ProgCache = LockedCache.get();
auto It = ProgCache.Cache.find(CacheKey);

if (It != ProgCache.Cache.end()) {
// We are about to remove this program now.
// (1) Remove it from KernelPerProgram cache.
// (2) Remove corresponding entries from KernelFastCache.
// (3) Remove it from ProgramCache KeyMap.
// (4) Remove it from the ProgramCache.
// (5) Remove it from ProgramSizeMap.
// (6) Update the cache size.

// Remove entry from the KernelsPerProgram cache.
ur_program_handle_t NativePrg = It->second->Val;
{
auto LockedCacheKP = acquireKernelsPerProgramCache();
// List kernels that are to be removed from the cache, if tracing is
// enabled.
if (SYCLConfig<SYCL_CACHE_TRACE>::isTraceInMemCache()) {
for (const auto &Kernel : LockedCacheKP.get()[NativePrg])
traceKernel("Kernel evicted.", Kernel.first);
}
LockedCacheKP.get().erase(NativePrg);
}

{
// Remove corresponding entries from KernelFastCache.
std::unique_lock<std::mutex> Lock(MKernelFastCacheMutex);
if (auto FastCacheKeyItr =
MProgramToKernelFastCacheKeyMap.find(NativePrg);
FastCacheKeyItr != MProgramToKernelFastCacheKeyMap.end()) {
for (const auto &FastCacheKey : FastCacheKeyItr->second) {
MKernelFastCache.erase(FastCacheKey);
traceKernel("Kernel evicted.", FastCacheKey.second, true);
}
MProgramToKernelFastCacheKeyMap.erase(FastCacheKeyItr);
}
}

// Remove entry from ProgramCache KeyMap.
CommonProgramKeyT CommonKey =
std::make_pair(CacheKey.first.second, CacheKey.second);
// Since KeyMap is a multi-map, we need to iterate over all entries
// with this CommonKey and remove those that match the CacheKey.
auto KeyMapItrRange = LockedCache.get().KeyMap.equal_range(CommonKey);
for (auto KeyMapItr = KeyMapItrRange.first;
KeyMapItr != KeyMapItrRange.second; ++KeyMapItr) {
if (KeyMapItr->second == CacheKey) {
LockedCache.get().KeyMap.erase(KeyMapItr);
break;
}
}

// Get size of the program.
size_t ProgramSize = MCachedPrograms.ProgramSizeMap[It->second->Val];
// Evict program from the cache.
ProgCache.Cache.erase(It);
// Remove program size from the cache size.
MCachedPrograms.ProgramCacheSizeInBytes -= ProgramSize;
MCachedPrograms.ProgramSizeMap.erase(NativePrg);

traceProgram("Program evicted.", CacheKey);
} else
// This should never happen.
throw sycl::exception(sycl::make_error_code(sycl::errc::runtime),
"Program not found in the cache.");

CurrCacheSize = MCachedPrograms.ProgramCacheSizeInBytes;
CurrCacheSize = removeProgramByKey(CacheKey, ProgCache);
// Remove the program from the eviction list.
MEvictionList.popFront();
}
Expand Down Expand Up @@ -724,6 +732,24 @@ class KernelProgramCache {
}
}

void removeAllRelatedEntries(uint32_t ImageId) {
auto LockedCache = acquireCachedPrograms();
auto &ProgCache = LockedCache.get();

auto It = std::find_if(
ProgCache.KeyMap.begin(), ProgCache.KeyMap.end(),
[&ImageId](const auto &Entry) { return ImageId == Entry.first.first; });
if (It == ProgCache.KeyMap.end())
return;

auto Key = It->second;
removeProgramByKey(Key, ProgCache);
{
auto LockedEvictionList = acquireEvictionList();
LockedEvictionList.get().erase(Key);
}
}

private:
std::mutex MProgramCacheMutex;
std::mutex MKernelsPerProgramCacheMutex;
Expand Down
Loading

0 comments on commit 1e0c021

Please sign in to comment.