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_loader.cpp

@@ -949,6 +949,137 @@ 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.dditable = {};
+                    drv.handle = nullptr;
+                    drv.ddiInitialized = false;
+                    drv.initStatus = ZE_RESULT_ERROR_UNINITIALIZED;
+                    drv.zerDriverHandle = nullptr;
+                }
+            }
+        };
+        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

@@ -164,6 +164,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,15 @@ 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()
+
 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")
 

test/loader_api.cpp

@@ -3997,4 +3997,72 @@ TEST_F(DriverOrderingTest,
     EXPECT_EQ(0, strcmp(errorDesc, "ERROR UNSUPPORTED FEATURE"));
   }
 
+// Proof of concept: unload one driver out of two and confirm the surviving driver
+// keeps working. Requires two drivers (ZE_ENABLE_ALT_DRIVERS) with loader intercept
+// enabled and the driver DDI-handle extension disabled so the loader wraps handles in
+// its object factories (which the unload safety check and reverse handle mapping rely on).
+TEST(
+    LoaderUnloadDriver,
+    GivenTwoDriversWhenUnloadingSecondDriverThenFirstDriverStillExecutes) {
+
+  ze_init_driver_type_desc_t desc = {ZE_STRUCTURE_TYPE_INIT_DRIVER_TYPE_DESC};
+  desc.flags = UINT32_MAX;
+  desc.pNext = nullptr;
+
+  uint32_t driverCount = 0;
+  ASSERT_EQ(ZE_RESULT_SUCCESS, zeInitDrivers(&driverCount, nullptr, &desc));
+  ASSERT_GE(driverCount, 2u)
+      << "This test requires at least two drivers via ZE_ENABLE_ALT_DRIVERS";
+  std::vector<ze_driver_handle_t> drivers(driverCount);
+  ASSERT_EQ(ZE_RESULT_SUCCESS, zeInitDrivers(&driverCount, drivers.data(), &desc));
+
+  ze_driver_handle_t firstDriver = drivers[0];
+  ze_driver_handle_t secondDriver = drivers[1];
+
+  // Exercise a driver end-to-end: create a context and a module, then tear them down.
+  auto executeOnDriver = [](ze_driver_handle_t driver) {
+    ze_context_desc_t contextDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC};
+    ze_context_handle_t context = nullptr;
+    ASSERT_EQ(ZE_RESULT_SUCCESS, zeContextCreate(driver, &contextDesc, &context));
+
+    uint32_t deviceCount = 0;
+    ASSERT_EQ(ZE_RESULT_SUCCESS, zeDeviceGet(driver, &deviceCount, nullptr));
+    ASSERT_GT(deviceCount, 0u);
+    std::vector<ze_device_handle_t> devices(deviceCount);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, zeDeviceGet(driver, &deviceCount, devices.data()));
+
+    ze_module_desc_t moduleDesc = {ZE_STRUCTURE_TYPE_MODULE_DESC};
+    ze_module_handle_t module = nullptr;
+    EXPECT_EQ(ZE_RESULT_SUCCESS,
+              zeModuleCreate(context, devices[0], &moduleDesc, &module, nullptr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeModuleDestroy(module));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeContextDestroy(context));
+  };
+
+  // 1. Execute on the first driver.
+  executeOnDriver(firstDriver);
+
+  // 2. A driver that still owns a live child object cannot be unloaded.
+  ze_context_desc_t contextDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC};
+  ze_context_handle_t liveContext = nullptr;
+  ASSERT_EQ(ZE_RESULT_SUCCESS, zeContextCreate(secondDriver, &contextDesc, &liveContext));
+  EXPECT_EQ(ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE, zelUnloadDriver(secondDriver));
+  EXPECT_EQ(ZE_RESULT_SUCCESS, zeContextDestroy(liveContext));
+
+  // 3. Now idle, the second driver unloads successfully.
+  EXPECT_EQ(ZE_RESULT_SUCCESS, zelUnloadDriver(secondDriver));
+  // Note: secondDriver is invalid past this point and must not be reused.
+
+  // 4. A null handle is rejected.
+  EXPECT_EQ(ZE_RESULT_ERROR_INVALID_NULL_HANDLE, zelUnloadDriver(nullptr));
+
+  // 5. The loader now reports one fewer driver.
+  uint32_t driverGetCount = 0;
+  EXPECT_EQ(ZE_RESULT_SUCCESS, zeDriverGet(&driverGetCount, nullptr));
+  EXPECT_EQ(driverGetCount, driverCount - 1);
+
+  // 6. The first driver is completely unaffected and still executes.
+  executeOnDriver(firstDriver);
+}
+
 } // namespace