POC for zelUnloadDrivers

include/loader/ze_loader.h

@@ -170,6 +170,36 @@ zelLoaderTranslateHandle(
    void *handleIn,
    void **handleOut);
 
+/**
+ * @brief [PROOF OF CONCEPT] Unloads a single Level Zero driver identified by its handle.
+ *
+ * This function unloads a driver that was previously reported by zeDriverGet()/zeInitDrivers().
+ * The driver's shared library is freed, its DDI tables are cleared, and the driver is removed
+ * from the loader's internal driver lists so it is no longer reported by subsequent enumeration.
+ *
+ * Preconditions / limitations (proof of concept):
+ * - The driver handle must correspond to a currently loaded driver.
+ * - The driver must be unused: all child objects created through the driver (contexts, command
+ *   queues, command lists, events, event pools, modules, kernels, images, samplers, fences, and
+ *   physical memory) must have been destroyed first. If any remain live, the unload is rejected
+ *   as unsafe.
+ *
+ * After a successful unload, the supplied driver handle (and any handles derived from it) are
+ * invalid and must not be used.
+ *
+ * @param[in] hDriver
+ *   The driver handle to unload, as returned by zeDriverGet() or zeInitDrivers().
+ *
+ * @return
+ *   - ZE_RESULT_SUCCESS if the driver was successfully unloaded.
+ *   - ZE_RESULT_ERROR_INVALID_NULL_HANDLE if hDriver is NULL or does not match a loaded driver.
+ *   - ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE if the driver still owns live child objects.
+ *   - ZE_RESULT_ERROR_UNINITIALIZED if the loader has not been initialized.
+ */
+ZE_APIEXPORT ze_result_t ZE_APICALL
+zelUnloadDriver(
+   ze_driver_handle_t hDriver);
+
 /**
  * @brief Notifies the loader that a driver has been removed and forces prevention of subsequent API calls.
  *

source/inc/ze_singleton.h

@@ -70,6 +70,22 @@ class singleton_factory_t
         return map.find( getKey( _key ) ) != map.end();
     }
 
+    //////////////////////////////////////////////////////////////////////////
+    /// counts the live instances whose dditable pointer matches the argument.
+    /// used to detect whether a driver still owns outstanding child objects.
+    template<typename _dditable_t>
+    size_t countByDditable( const _dditable_t* dditable )
+    {
+        std::lock_guard<std::mutex> lk( mut );
+        size_t count = 0;
+        for( const auto& entry : map )
+        {
+            if( entry.second && entry.second->dditable == dditable )
+                ++count;
+        }
+        return count;
+    }
+
     //////////////////////////////////////////////////////////////////////////
     /// once the key is no longer valid, release the singleton
     void release( _key_t _key )

source/lib/ze_lib.cpp

@@ -462,6 +462,24 @@ zelLoaderTranslateHandle(
 #endif
 }
 
+ze_result_t ZE_APICALL
+zelUnloadDriver(
+   ze_driver_handle_t hDriver)
+{
+#ifdef L0_STATIC_LOADER_BUILD
+    if(nullptr == ze_lib::context->loader)
+        return ZE_RESULT_ERROR_UNINITIALIZED;
+    typedef ze_result_t (ZE_APICALL *zelUnloadDriverInternal_t)(ze_driver_handle_t hDriver);
+    auto unloadDriver = reinterpret_cast<zelUnloadDriverInternal_t>(
+            GET_FUNCTION_PTR(ze_lib::context->loader, "zelUnloadDriverInternal") );
+    if (nullptr == unloadDriver)
+        return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    return unloadDriver(hDriver);
+#else
+    return zelUnloadDriverInternal(hDriver);
+#endif
+}
+
 ze_result_t ZE_APICALL
 zelSetDriverTeardown()
 {

source/loader/ze_ldrddi.cpp

@@ -362,6 +362,9 @@ namespace loader
 
         uint32_t total_driver_handle_count = 0;
         for( auto& drv : loader::context->zeDrivers ) {
+            if (drv.unloaded) {
+                continue; // Never reload an explicitly unloaded driver.
+            }
             if (!drv.handle || !drv.ddiInitialized) {
                 auto res = loader::context->init_driver( drv, 0, desc);
                 if (res != ZE_RESULT_SUCCESS || drv.zeddiInitResult != ZE_RESULT_SUCCESS) {
@@ -384,6 +387,9 @@ namespace loader
 
         for( auto& drv : loader::context->zeDrivers )
         {
+            if (drv.unloaded) {
+                continue; // Unloaded drivers are a hole in the list; do not enumerate them.
+            }
             if (!drv.ddiInitialized || !drv.dditable.ze.Global.pfnInitDrivers) {
                 drv.initDriversStatus = ZE_RESULT_ERROR_UNINITIALIZED;
                 result = ZE_RESULT_ERROR_UNINITIALIZED;

source/loader/ze_loader.cpp

@@ -106,6 +106,9 @@ namespace loader
         // Group drivers by type and track their original indices
         for (uint32_t i = 0; i < originalDrivers.size(); ++i) {
             const auto& driver = originalDrivers[i];
+            if (driver.unloaded) {
+                continue; // Tombstones are not eligible for type/index based ordering.
+            }
             switch (driver.driverType) {
                 case ZEL_DRIVER_TYPE_DISCRETE_GPU:
                     discreteGPUDrivers.push_back(driver);
@@ -147,6 +150,7 @@ namespace loader
             switch (spec.type) {
                 case DriverOrderSpecType::BY_GLOBAL_INDEX:
                     if (spec.globalIndex < originalDrivers.size() &&
+                        !originalDrivers[spec.globalIndex].unloaded &&
                         usedGlobalIndices.find(spec.globalIndex) == usedGlobalIndices.end()) {
                         orderedDrivers.push_back(originalDrivers[spec.globalIndex]);
                         usedGlobalIndices.insert(spec.globalIndex);
@@ -253,6 +257,9 @@ namespace loader
         for (auto &driver : *drivers) {
             uint32_t pCount = 0;
             std::vector<ze_driver_handle_t> driverHandles;
+            if (driver.unloaded) {
+                continue; // Skip tombstoned drivers; they must not be probed or re-typed.
+            }
             driver.pciOrderingRequested = loader::context->pciOrderingRequested;
             ze_result_t res = ZE_RESULT_SUCCESS;
             if (desc && driver.dditable.ze.Global.pfnInitDrivers) {
@@ -480,6 +487,11 @@ namespace loader
 
     ze_result_t context_t::init_driver(driver_t &driver, ze_init_flags_t flags, ze_init_driver_type_desc_t* desc) {
         bool loadDriver = false;
+        // An unloaded driver must never be reloaded; doing so would re-dlopen the library and
+        // leave it mapped in the process with no way to reach it.
+        if (driver.unloaded) {
+            return ZE_RESULT_ERROR_UNINITIALIZED;
+        }
         if (debugTraceEnabled) {
             std::string message = "Initializing driver " + driver.name + " with type " + std::to_string(driver.driverType);\
             debug_trace_message(message, "");
@@ -949,6 +961,150 @@ namespace loader
         }
     };
 
+    bool context_t::isDriverInUse(const dditable_t *dditable)
+    {
+        // Proof-of-concept "state machine" detection: a driver is considered in use if any
+        // child object created through it is still live in the loader's object factories.
+        // These are the primary stateful resources an application creates and must destroy
+        // before a driver can be safely unloaded.
+        return ze_context_factory.countByDditable(dditable) > 0
+            || ze_command_queue_factory.countByDditable(dditable) > 0
+            || ze_command_list_factory.countByDditable(dditable) > 0
+            || ze_event_pool_factory.countByDditable(dditable) > 0
+            || ze_event_factory.countByDditable(dditable) > 0
+            || ze_fence_factory.countByDditable(dditable) > 0
+            || ze_image_factory.countByDditable(dditable) > 0
+            || ze_sampler_factory.countByDditable(dditable) > 0
+            || ze_module_factory.countByDditable(dditable) > 0
+            || ze_kernel_factory.countByDditable(dditable) > 0
+            || ze_physical_mem_factory.countByDditable(dditable) > 0;
+    }
+
+    ze_result_t context_t::unloadDriver(ze_driver_handle_t hDriver)
+    {
+        if (nullptr == hDriver) {
+            return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
+        }
+
+        // Locate the driver_t backing this user-facing handle. When the loader intercepts
+        // handles, hDriver is a ze_driver_object_t whose dditable points into a zeDrivers entry.
+        // Otherwise (single-driver / DDI-handle path) the raw driver handle was stored in
+        // driver_t::zerDriverHandle.
+        dditable_t *targetDdiTable = nullptr;
+        HMODULE targetModule = nullptr;
+        std::string targetName;
+        ze_driver_handle_t rawHandle = nullptr;
+        bool found = false;
+
+        std::lock_guard<std::mutex> lock(sortMutex);
+
+        if (intercept_enabled) {
+            auto obj = reinterpret_cast<ze_driver_object_t *>(hDriver);
+            for (auto &drv : zeDrivers) {
+                if (&drv.dditable == obj->dditable) {
+                    targetDdiTable = &drv.dditable;
+                    targetModule = drv.handle;
+                    targetName = drv.name;
+                    rawHandle = obj->handle;
+                    found = true;
+                    break;
+                }
+            }
+        } else {
+            for (auto &drv : zeDrivers) {
+                if (drv.zerDriverHandle == hDriver) {
+                    targetDdiTable = &drv.dditable;
+                    targetModule = drv.handle;
+                    targetName = drv.name;
+                    rawHandle = hDriver;
+                    found = true;
+                    break;
+                }
+            }
+        }
+
+        if (!found) {
+            if (debugTraceEnabled) {
+                debug_trace_message("zelUnloadDriver: driver handle not found", "");
+            }
+            return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
+        }
+
+        // Safety gate: refuse to unload a driver that still owns live child objects.
+        if (isDriverInUse(targetDdiTable)) {
+            if (debugTraceEnabled) {
+                debug_trace_message("zelUnloadDriver: driver still in use: ", targetName);
+            }
+            return ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE;
+        }
+
+        // Release the wrapper object for the driver handle so the factory no longer tracks it.
+        if (rawHandle) {
+            ze_driver_factory.release(rawHandle);
+        }
+
+        // Clear every copy of this driver across the loader's driver lists in place, zeroing its
+        // DDI tables and marking it uninitialized so it is no longer reported by enumeration.
+        // Entries are tombstoned rather than erased so that pointers held by other drivers'
+        // handles and their child objects (which reference driver_t::dditable by address) remain
+        // valid -- unloading one driver must not disturb another.
+        auto clearMatching = [&](driver_vector_t &vec) {
+            for (auto &drv : vec) {
+                if (drv.handle == targetModule && drv.name == targetName) {
+                    drv.unloaded = true;
+                    drv.dditable = {};
+                    drv.properties = {};
+                    drv.handle = nullptr;
+                    drv.zerDriverHandle = nullptr;
+                    // Neutralize the sort key so the tombstone is never bucketed by ordering.
+                    drv.driverType = ZEL_DRIVER_TYPE_FORCE_UINT32;
+                    drv.driverInuse = false;
+                    drv.ddiInitialized = false;
+                    drv.legacyInitAttempted = false;
+                    drv.driverDDIHandleSupportQueried = false;
+                    drv.initStatus = ZE_RESULT_ERROR_UNINITIALIZED;
+                    drv.initSysManStatus = ZE_RESULT_ERROR_UNINITIALIZED;
+                    drv.initDriversStatus = ZE_RESULT_ERROR_UNINITIALIZED;
+                    drv.zeddiInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
+                    drv.zetddiInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
+                    drv.zesddiInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
+                    drv.zerddiInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
+                }
+            }
+        };
+        clearMatching(zeDrivers);
+        clearMatching(zesDrivers);
+        clearMatching(allDrivers);
+
+        // Free the driver library. All copies were just cleared, so nothing else references it.
+        if (targetModule) {
+            auto free_result = FREE_DRIVER_LIBRARY(targetModule);
+            auto failure = FREE_DRIVER_LIBRARY_FAILURE_CHECK(free_result);
+            if (debugTraceEnabled && failure) {
+                std::string freeLibraryErrorValue;
+                GET_LIBRARY_ERROR(freeLibraryErrorValue);
+                if (!freeLibraryErrorValue.empty()) {
+                    debug_trace_message("zelUnloadDriver: Free Library Failed for " + targetName + " with ", freeLibraryErrorValue);
+                }
+            }
+        }
+
+        // Keep the default ZER DDI table pointing at a driver that is still loaded, if any.
+        loader::defaultZerDdiTable = nullptr;
+        for (auto &drv : zeDrivers) {
+            if (drv.handle) {
+                loader::defaultZerDdiTable = &drv.dditable.zer;
+                break;
+            }
+        }
+
+        if (debugTraceEnabled) {
+            debug_trace_message("zelUnloadDriver: unloaded driver ", targetName);
+        }
+
+        return ZE_RESULT_SUCCESS;
+    }
+
     void context_t::add_loader_version(){
         zel_component_version_t compVersion = {};
         string_copy_s(compVersion.component_name, LOADER_COMP_NAME, ZEL_COMPONENT_STRING_SIZE - 1);

source/loader/ze_loader_api.cpp

@@ -292,6 +292,16 @@ zelLoaderTranslateHandleInternal(
     return ZE_RESULT_SUCCESS;
 }
 
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelUnloadDriverInternal(
+   ze_driver_handle_t hDriver)
+{
+    if (!loader::context) {
+        return ZE_RESULT_ERROR_UNINITIALIZED;
+    }
+    return loader::context->unloadDriver(hDriver);
+}
+
 
 #if defined(__cplusplus)
 }

source/loader/ze_loader_api.h

@@ -88,6 +88,18 @@ zelLoaderTranslateHandleInternal(
    void **handleOut);                      //Output: Pointer to handleOut is set to driver handle if successful
 
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Proof-of-concept: unload a single driver identified by its handle.
+///
+/// @returns
+///     - ::ZE_RESULT_SUCCESS
+///     - ::ZE_RESULT_ERROR_INVALID_NULL_HANDLE
+///     - ::ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelUnloadDriverInternal(
+   ze_driver_handle_t hDriver);           //Input: driver handle to unload
+
+
 #if defined(__cplusplus)
 }
 #endif

source/loader/ze_loader_internal.h

@@ -67,6 +67,10 @@ namespace loader
         ze_result_t zetddiInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
         ze_result_t zesddiInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
         ze_result_t zerddiInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
+        // Set once the driver has been explicitly unloaded via zelUnloadDriver. Distinguishes a
+        // tombstoned slot from a not-yet-loaded one (both have handle==nullptr) so the loader
+        // never reloads it and skips it during ordering/enumeration.
+        bool unloaded = false;
     };
 
     using driver_vector_t = std::vector< driver_t >;
@@ -164,6 +168,11 @@ namespace loader
         void add_loader_version();
         bool driverSorting(driver_vector_t *drivers, ze_init_driver_type_desc_t* desc, bool sysmanOnly);
         void driverOrdering(driver_vector_t *drivers);
+
+        // Proof-of-concept: unload a single driver identified by its user-facing handle.
+        ze_result_t unloadDriver(ze_driver_handle_t hDriver);
+        // Returns true if the driver (identified by its dditable) still owns live child objects.
+        bool isDriverInUse(const dditable_t *dditable);
         ~context_t();
         bool intercept_enabled = false;
         bool debugTraceEnabled = false;

test/CMakeLists.txt

@@ -225,6 +225,23 @@ else()
   set_property(TEST tests_multi_driver_zeandzesdriverget_sort APPEND PROPERTY ENVIRONMENT "ZE_ENABLE_LOADER_DEBUG_TRACE=1;ZE_ENABLE_ALT_DRIVERS=$<TARGET_FILE:ze_null_test1>,$<TARGET_FILE:ze_null_test2>")
 endif()
 
+# Proof of concept: unload one of two drivers and confirm the survivor still works.
+# Intercept + DDI-ext disabled so the loader wraps handles in its object factories.
+add_test(NAME tests_unload_driver_multi COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenTwoDriversWhenUnloadingSecondDriverThenFirstDriverStillExecutes)
+if (MSVC)
+  set_property(TEST tests_unload_driver_multi PROPERTY ENVIRONMENT "ZE_ENABLE_LOADER_INTERCEPT=1;ZEL_TEST_NULL_DRIVER_DISABLE_DDI_EXT=3;ZE_ENABLE_LOADER_DEBUG_TRACE=1;ZE_ENABLE_ALT_DRIVERS=$<TARGET_FILE_DIR:ze_null_test1>/ze_null_test1.dll,$<TARGET_FILE_DIR:ze_null_test1>/ze_null_test2.dll")
+else()
+  set_property(TEST tests_unload_driver_multi PROPERTY ENVIRONMENT "ZE_ENABLE_LOADER_INTERCEPT=1;ZEL_TEST_NULL_DRIVER_DISABLE_DDI_EXT=3;ZE_ENABLE_LOADER_DEBUG_TRACE=1;ZE_ENABLE_ALT_DRIVERS=$<TARGET_FILE:ze_null_test1>,$<TARGET_FILE:ze_null_test2>")
+endif()
+
+# Design probe: unload a driver then attempt to reload it via re-initialization.
+add_test(NAME tests_unload_driver_reload COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenUnloadedDriverWhenReinitializedThenDriverRemainsUnloaded)
+if (MSVC)
+  set_property(TEST tests_unload_driver_reload PROPERTY ENVIRONMENT "ZE_ENABLE_LOADER_INTERCEPT=1;ZEL_TEST_NULL_DRIVER_DISABLE_DDI_EXT=3;ZE_ENABLE_LOADER_DEBUG_TRACE=1;ZE_ENABLE_ALT_DRIVERS=$<TARGET_FILE_DIR:ze_null_test1>/ze_null_test1.dll,$<TARGET_FILE_DIR:ze_null_test1>/ze_null_test2.dll")
+else()
+  set_property(TEST tests_unload_driver_reload PROPERTY ENVIRONMENT "ZE_ENABLE_LOADER_INTERCEPT=1;ZEL_TEST_NULL_DRIVER_DISABLE_DDI_EXT=3;ZE_ENABLE_LOADER_DEBUG_TRACE=1;ZE_ENABLE_ALT_DRIVERS=$<TARGET_FILE:ze_null_test1>,$<TARGET_FILE:ze_null_test2>")
+endif()
+
 add_test(NAME tests_loader_teardown_check COMMAND tests --gtest_filter=*GivenLoaderNotInDestructionStateWhenCallingzelCheckIsLoaderInTearDownThenFalseIsReturned)
 set_property(TEST tests_loader_teardown_check PROPERTY ENVIRONMENT "ZE_ENABLE_LOADER_DEBUG_TRACE=1;ZE_ENABLE_NULL_DRIVER=1")