ensure UR will clear context on unloading

Author: omarahmed1111

source/adapters/level_zero/CMakeLists.txt

@@ -140,6 +140,11 @@ if(UR_BUILD_ADAPTER_L0)
             PRIVATE
                 ${CMAKE_CURRENT_SOURCE_DIR}/adapter_lib_init_linux.cpp
         )
+    else()
+        target_sources(ur_adapter_level_zero
+            PRIVATE
+                ${CMAKE_CURRENT_SOURCE_DIR}/adapter_lib_init_windows.cpp
+        )
     endif()
 
     # TODO: fix level_zero adapter conversion warnings

source/adapters/level_zero/adapter_lib_init_windows.cpp

@@ -0,0 +1,43 @@
+//===--------- adapter_lib_init_linux.cpp - Level Zero Adapter ------------===//
+//
+// Copyright (C) 2023 Intel Corporation
+//
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+// Exceptions. See LICENSE.TXT
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "adapter.hpp"
+#include "ur_level_zero.hpp"
+
+#include <windows.h>
+
+BOOL WINAPI DllMain(HINSTANCE hinstDLL, // handle to DLL module
+                    DWORD fdwReason,    // reason for calling function
+                    LPVOID lpReserved)  // reserved
+{
+  switch (fdwReason) {
+  case DLL_PROCESS_ATTACH:
+    break;
+  case DLL_PROCESS_DETACH: {
+    std::lock_guard<std::mutex> Lock{GlobalAdapter->Mutex};
+    const auto *platforms = GlobalAdapter->PlatformCache->get_value();
+    for (const auto &p : *platforms) {
+      std::scoped_lock<ur_shared_mutex> ContextsLock(p->ContextsMutex);
+      while (!p->Contexts.empty()) {
+        ur_context_handle_t &ctx = p->Contexts.front();
+        ctx->deleteCachedObjectsOnDestruction();
+        UR_CALL(urContextRelease(ctx));
+        deleteFromCachedList<ur_context_handle_t>(ctx, p->Contexts);
+      }
+    }
+    break;
+  }
+  case DLL_THREAD_ATTACH:
+    break;
+  case DLL_THREAD_DETACH:
+    break;
+  }
+  return TRUE;
+}

source/adapters/level_zero/common.hpp

@@ -450,6 +450,9 @@ extern const bool UseUSMAllocator;
 
 // Controls support of the indirect access kernels and deferred memory release.
 const bool IndirectAccessTrackingEnabled = [] {
+#ifdef _WIN32
+  return false;
+#endif
   char *UrRet = std::getenv("UR_L0_TRACK_INDIRECT_ACCESS_MEMORY");
   char *PiRet = std::getenv("SYCL_PI_LEVEL_ZERO_TRACK_INDIRECT_ACCESS_MEMORY");
   const bool RetVal = UrRet ? std::stoi(UrRet) : (PiRet ? std::stoi(PiRet) : 0);
@@ -530,4 +533,19 @@ extern thread_local int32_t ErrorAdapterNativeCode;
                                       ur_result_t ErrorCode,
                                       int32_t AdapterErrorCode);
 
+template <class T>
+void addToCachedList(T &CachedObject, std::list<T> &CachedList) {
+  auto It = std::find(CachedList.begin(), CachedList.end(), CachedObject);
+  if (It == CachedList.end()) {
+    CachedList.push_back(CachedObject);
+  }
+}
+
+template <class T>
+void deleteFromCachedList(T &CachedObject, std::list<T> &CachedList) {
+  auto It = std::find(CachedList.begin(), CachedList.end(), CachedObject);
+  if (It != CachedList.end())
+    CachedList.erase(It);
+}
+
 #define L0_DRIVER_INORDER_MIN_VERSION 29534

source/adapters/level_zero/context.cpp

@@ -40,10 +40,15 @@ UR_APIEXPORT ur_result_t UR_APICALL urContextCreate(
 
     Context->initialize();
     *RetContext = reinterpret_cast<ur_context_handle_t>(Context);
+#ifdef _WIN32
+    std::scoped_lock<ur_shared_mutex> Lock(Platform->ContextsMutex);
+    addToCachedList<ur_context_handle_t>(*RetContext, Platform->Contexts);
+#else
     if (IndirectAccessTrackingEnabled) {
       std::scoped_lock<ur_shared_mutex> Lock(Platform->ContextsMutex);
       Platform->Contexts.push_back(*RetContext);
     }
+#endif
   } catch (const std::bad_alloc &) {
     return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
   } catch (...) {
@@ -355,13 +360,21 @@ ur_result_t ContextReleaseHelper(ur_context_handle_t Context) {
   if (!Context->RefCount.decrementAndTest())
     return UR_RESULT_SUCCESS;
 
-  if (IndirectAccessTrackingEnabled) {
+  auto DeleteFromContextsCache = [&]() {
     ur_platform_handle_t Plt = Context->getPlatform();
     auto &Contexts = Plt->Contexts;
     auto It = std::find(Contexts.begin(), Contexts.end(), Context);
     if (It != Contexts.end())
       Contexts.erase(It);
+  };
+
+#ifdef _WIN32
+  DeleteFromContextsCache();
+#else
+  if (IndirectAccessTrackingEnabled) {
+    DeleteFromContextsCache();
   }
+#endif
   ze_context_handle_t DestroyZeContext =
       Context->OwnNativeHandle ? Context->ZeContext : nullptr;
 
@@ -451,6 +464,12 @@ ur_result_t ur_context_handle_t_::finalize() {
       }
     }
   }
+
+  for (auto &kernel : KernelsCache) {
+    UR_CALL(urKernelRelease(kernel));
+  }
+  KernelsCache.clear();
+
   return UR_RESULT_SUCCESS;
 }
 

source/adapters/level_zero/context.hpp

@@ -175,6 +175,8 @@ struct ur_context_handle_t_ : _ur_object {
   std::vector<std::unordered_map<ur_device_handle_t, size_t>>
       EventCachesDeviceMap{4};
 
+  std::list<ur_kernel_handle_t> KernelsCache;
+
   // Initialize the PI context.
   ur_result_t initialize();
 
@@ -309,6 +311,14 @@ struct ur_context_handle_t_ : _ur_object {
   // Get handle to the L0 context
   ze_context_handle_t getZeHandle() const;
 
+  void deleteCachedObjectsOnDestruction() {
+    while (!KernelsCache.empty()) {
+      ur_kernel_handle_t &kernel = KernelsCache.front();
+      UR_CALL_THROWS(urKernelRelease(kernel));
+      deleteFromCachedList(kernel, KernelsCache);
+    }
+  }
+
 private:
   // Get the cache of events for a provided scope and profiling mode.
   auto getEventCache(bool HostVisible, bool WithProfiling,

source/adapters/level_zero/kernel.cpp

@@ -591,6 +591,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelCreate(
   try {
     ur_kernel_handle_t_ *UrKernel = new ur_kernel_handle_t_(true, Program);
     *RetKernel = reinterpret_cast<ur_kernel_handle_t>(UrKernel);
+
+#ifdef _WIN32
+    auto &Context = Program->Context;
+    auto It = std::find(Context->KernelsCache.begin(),
+                        Context->KernelsCache.end(), *RetKernel);
+    if (It == Context->KernelsCache.end()) {
+      Context->KernelsCache.push_back(*RetKernel);
+    }
+#endif
+
   } catch (const std::bad_alloc &) {
     return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
   } catch (...) {
@@ -902,10 +912,13 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelRelease(
         return ze2urResult(ZeResult);
     }
   }
+
   Kernel->ZeKernelMap.clear();
+
   if (IndirectAccessTrackingEnabled) {
     UR_CALL(urContextRelease(KernelProgram->Context));
   }
+
   // do a release on the program this kernel was part of without delete of the
   // program handle
   KernelProgram->ur_release_program_resources(false);