Add ability to Register a TeardownCallback to notify release of L0 resources

Signed-off-by: Neil R. Spruit <neil.r.spruit@intel.com>

Author: nrspruit

include/loader/ze_loader.h

@@ -94,6 +94,20 @@ zelEnableTracingLayer();
 ZE_DLLEXPORT bool ZE_APICALL
 zelCheckIsLoaderInTearDown();
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Exported function for registering a callback to indicate teardown.
+/// @details The callback function will be invoked when the loader is in teardown.
+///
+typedef void (*zel_loader_teardown_callback_t)();
+typedef void (*zel_application_teardown_callback_t)(uint32_t index);
+
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelRegisterTeardownCallback(
+   zel_loader_teardown_callback_t application_callback, // [in] Pointer to the user's application callback function
+   zel_application_teardown_callback_t *loader_callback,      // [out] Pointer to the L0 Loader's callback function 
+   uint32_t *index // [out] Index of the callback function
+);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for Disabling the Tracing Layer During Runtime.
 ///

source/lib/windows/lib_init.cpp

@@ -16,6 +16,8 @@ namespace ze_lib
 #ifndef DYNAMIC_LOAD_LOADER
     extern "C" BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {
         if (fdwReason == DLL_PROCESS_DETACH) {
+            printf("DLL_PROCESS_DETACH\n");
+            ze_lib::context->teardownCallbacks.clear();
             delete context;
             delete loader::context;
         } else if (fdwReason == DLL_PROCESS_ATTACH) {

source/lib/ze_lib.cpp

@@ -27,7 +27,18 @@ namespace ze_lib
         }
     }
     bool delayContextDestruction = false;
+    bool loaderTeardownCallbackReceived = false;
+    void staticLoaderTeardownCallback() {
+        printf("ze_lib Static Teardown Callback\n");
+        loaderTeardownCallbackReceived = true;
+    }
     #endif
+    void applicationTeardownCallback(uint32_t index) {
+        printf("ze_lib Application Teardown Callback %d\n", index);
+        if (ze_lib::context->teardownCallbacks.find(index) != ze_lib::context->teardownCallbacks.end()) {
+            ze_lib::context->teardownCallbacks.erase(index);
+        }
+    }
     bool destruction = false;
 
     ///////////////////////////////////////////////////////////////////////////////
@@ -39,10 +50,22 @@ namespace ze_lib
     ///////////////////////////////////////////////////////////////////////////////
     __zedlllocal context_t::~context_t()
     {
+        if (debugTraceEnabled) {
+            debug_trace_message("ze_lib Context Destructor", "");
+        }
 #ifdef DYNAMIC_LOAD_LOADER
+        if (!loaderTeardownCallbackReceived) {
+            loaderTeardownCallback(loaderTeardownCallbackIndex);
+        }
         if (loader) {
             FREE_DRIVER_LIBRARY( loader );
         }
+#else
+        // Call the teardown callbacks
+        for (auto &callback : teardownCallbacks) {
+            callback.second();
+        }
+        teardownCallbacks.clear();
 #endif
         ze_lib::destruction = true;
     };
@@ -343,9 +366,19 @@ namespace ze_lib
             isInitialized = true;
         }
         #ifdef DYNAMIC_LOAD_LOADER
-        if (!delayContextDestruction) {
-            std::atexit(context_at_exit_destructor);
-        }
+        std::call_once(ze_lib::context->initTeardownCallbacksOnce, [this]() {
+            if (!delayContextDestruction) {
+                std::atexit(context_at_exit_destructor);
+            }
+            // Get the function pointer for zelRegisterTeardownCallback from the loader
+            typedef ze_result_t (ZE_APICALL *zelRegisterTeardownCallback_t)(
+                zel_loader_teardown_callback_t, 
+                zel_application_teardown_callback_t*, 
+                uint32_t*);
+            auto pfnZelRegisterTeardownCallback = reinterpret_cast<zelRegisterTeardownCallback_t>(
+                GET_FUNCTION_PTR(loader, "zelRegisterTeardownCallback"));
+            pfnZelRegisterTeardownCallback(staticLoaderTeardownCallback, &loaderTeardownCallback, &loaderTeardownCallbackIndex);
+        });
         #endif
         return result;
     }
@@ -397,6 +430,13 @@ zelSetDriverTeardown()
 {
     ze_result_t result = ZE_RESULT_SUCCESS;
     if (!ze_lib::destruction) {
+        if (ze_lib::context) {
+            // Call the teardown callbacks
+            for (auto &callback : ze_lib::context->teardownCallbacks) {
+                callback.second();
+            }
+        }
+
         ze_lib::destruction = true;
     }
     return result;
@@ -476,6 +516,30 @@ void stabilityCheck(std::promise<int> stabilityPromise) {
 }
 #endif
 
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelRegisterTeardownCallback(
+   zel_loader_teardown_callback_t application_callback, // [in] Pointer to the user's application callback function
+   zel_application_teardown_callback_t *loader_callback,      // [out] Pointer to the L0 Loader's callback function 
+   uint32_t *index // [out] Index of the callback function
+) {
+    ze_result_t result = ZE_RESULT_SUCCESS;
+    if (nullptr == application_callback) {
+        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (!ze_lib::context) {
+        return ZE_RESULT_ERROR_UNINITIALIZED;
+    }
+    *loader_callback = ze_lib::applicationTeardownCallback;
+    ze_lib::context->teardownCallbacksCount.fetch_add(1);
+    *index = ze_lib::context->teardownCallbacksCount.load();
+    ze_lib::context->teardownCallbacks.insert(std::pair<uint32_t, zel_loader_teardown_callback_t>(*index, application_callback));
+    if (ze_lib::context->debugTraceEnabled) {
+        std::string message = "Registered teardown callback with index: " + std::to_string(*index);
+        ze_lib::context->debug_trace_message(message, "");
+    }
+    return result;
+}
+
 /**
  * @brief Checks if the loader is in the process of tearing down.
  *
@@ -497,6 +561,9 @@ zelCheckIsLoaderInTearDown() {
         return true;
     }
     #if defined(DYNAMIC_LOAD_LOADER) && defined(_WIN32)
+    if (ze_lib::loaderTeardownCallbackReceived) {
+        return true;
+    }
     std::promise<int> stabilityPromise;
     std::future<int> resultFuture = stabilityPromise.get_future();
     int result = -1;

source/lib/ze_lib.h

@@ -24,6 +24,9 @@
 #include <typeinfo>
 #include <iostream>
 
+typedef void (*zel_loader_teardown_callback_t)();
+typedef void (*zel_application_teardown_callback_t)(uint32_t index);
+
 namespace ze_lib
 {
     ///////////////////////////////////////////////////////////////////////////////
@@ -175,12 +178,20 @@ namespace ze_lib
         bool debugTraceEnabled = false;
         bool dynamicTracingSupported = true;
         ze_pfnDriverGet_t loaderDriverGet = nullptr;
+        std::atomic<uint32_t> teardownCallbacksCount{0};
+        std::map<uint32_t, zel_loader_teardown_callback_t> teardownCallbacks;
+        #ifdef DYNAMIC_LOAD_LOADER
+        std::once_flag initTeardownCallbacksOnce;
+        zel_application_teardown_callback_t loaderTeardownCallback = nullptr;
+        uint32_t loaderTeardownCallbackIndex = 0;
+        #endif
     };
 
     extern bool destruction;
     extern context_t *context;
     #ifdef DYNAMIC_LOAD_LOADER
     extern bool delayContextDestruction;
+    extern bool loaderTeardownCallbackReceived;
     #endif
 
 } // namespace ze_lib