From 47af6b2b407da3a3655d792e2731fdc6df01902e Mon Sep 17 00:00:00 2001
From: "Neil R. Spruit" <neil.r.spruit@intel.com>
Date: Wed, 14 May 2025 14:02:43 -0700
Subject: [PATCH] Add ability to Register a TeardownCallback to notify release
 of L0 resources

Signed-off-by: Neil R. Spruit <neil.r.spruit@intel.com>
---
 include/loader/ze_loader.h | 11 +++++++++++
 source/lib/ze_lib.cpp      | 35 +++++++++++++++++++++++++++++++++++
 source/lib/ze_lib.h        |  4 ++++
 3 files changed, 50 insertions(+)

diff --git a/include/loader/ze_loader.h b/include/loader/ze_loader.h
index dec810e1..8ae2ccbc 100644
--- a/include/loader/ze_loader.h
+++ b/include/loader/ze_loader.h
@@ -94,6 +94,17 @@ 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_teardown_callback_t)(void);
+
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelRegisterTeardownCallback(
+   zel_teardown_callback_t callback // [in] Pointer to the callback function
+);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for Disabling the Tracing Layer During Runtime.
 ///
diff --git a/source/lib/ze_lib.cpp b/source/lib/ze_lib.cpp
index 2046f5fa..7c327e97 100644
--- a/source/lib/ze_lib.cpp
+++ b/source/lib/ze_lib.cpp
@@ -27,6 +27,10 @@ namespace ze_lib
         }
     }
     bool delayContextDestruction = false;
+    bool loaderTeardownCallbackReceived = false;
+    void teardownCallback() {
+        loaderTeardownCallbackReceived = true;
+    }
     #endif
     bool destruction = false;
 
@@ -39,6 +43,14 @@ namespace ze_lib
     ///////////////////////////////////////////////////////////////////////////////
     __zedlllocal context_t::~context_t()
     {
+        if (debugTraceEnabled) {
+            debug_trace_message("ze_lib Context Destructor", "");
+        }
+        // Call the teardown callbacks
+        for (auto &callback : teardownCallbacks) {
+            callback();
+        }
+        teardownCallbacks.clear();
 #ifdef DYNAMIC_LOAD_LOADER
         if (loader) {
             FREE_DRIVER_LIBRARY( loader );
@@ -346,6 +358,7 @@ namespace ze_lib
         if (!delayContextDestruction) {
             std::atexit(context_at_exit_destructor);
         }
+        zelRegisterTeardownCallback(teardownCallback);
         #endif
         return result;
     }
@@ -397,6 +410,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();
+            }
+        }
+
         ze_lib::destruction = true;
     }
     return result;
@@ -476,6 +496,18 @@ void stabilityCheck(std::promise<int> stabilityPromise) {
 }
 #endif
 
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelRegisterTeardownCallback(
+   zel_teardown_callback_t callback // [in] Pointer to the callback function
+) {
+    ze_result_t result = ZE_RESULT_SUCCESS;
+    if (nullptr == callback) {
+        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    ze_lib::context->teardownCallbacks.push_back(callback);
+    return result;
+}
+
 /**
  * @brief Checks if the loader is in the process of tearing down.
  *
@@ -497,6 +529,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;
diff --git a/source/lib/ze_lib.h b/source/lib/ze_lib.h
index f4a2bad8..589ce2b5 100644
--- a/source/lib/ze_lib.h
+++ b/source/lib/ze_lib.h
@@ -24,6 +24,8 @@
 #include <typeinfo>
 #include <iostream>
 
+typedef void (*zel_teardown_callback_t)(void);
+
 namespace ze_lib
 {
     ///////////////////////////////////////////////////////////////////////////////
@@ -175,12 +177,14 @@ namespace ze_lib
         bool debugTraceEnabled = false;
         bool dynamicTracingSupported = true;
         ze_pfnDriverGet_t loaderDriverGet = nullptr;
+        std::vector<zel_teardown_callback_t> teardownCallbacks;
     };
 
     extern bool destruction;
     extern context_t *context;
     #ifdef DYNAMIC_LOAD_LOADER
     extern bool delayContextDestruction;
+    extern bool loaderTeardownCallbackReceived;
     #endif
 
 } // namespace ze_lib