diff --git a/include/loader/ze_loader.h b/include/loader/ze_loader.h index f71f5392..16243a69 100644 --- a/include/loader/ze_loader.h +++ b/include/loader/ze_loader.h @@ -170,6 +170,120 @@ zelLoaderTranslateHandle( void *handleIn, void **handleOut); +/** + * @brief [PROOF OF CONCEPT] Flags controlling zelUnloadDriverExt(). + */ +typedef enum _zel_unload_driver_flag_t +{ + ZEL_UNLOAD_DRIVER_FLAG_NONE = 0, ///< Default: refuse to unload a driver that still owns + ///< live child objects. + ZEL_UNLOAD_DRIVER_FLAG_FORCE = ZE_BIT(0), ///< Unload even if child objects are still live. + ZEL_UNLOAD_DRIVER_FLAG_FORCE_UINT32 = 0x7fffffff +} zel_unload_driver_flag_t; + +typedef uint32_t zel_unload_driver_flags_t; + +/** + * @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 and its DDI tables are cleared, so the driver is no longer + * reported by subsequent enumeration. The driver's slot is emptied but retained, which makes the + * unload reversible via zelReloadDriver(): an unloaded driver is not blacklisted. + * + * Preconditions / limitations (proof of concept): + * - The driver handle must be a loader-issued driver handle for a currently loaded driver. Drivers + * reached through the driver DDI handle path (ZE_DRIVER_DDI_HANDLE_EXT) cannot be unloaded, + * because the application's handle is memory inside the library being unmapped; + * ZE_RESULT_ERROR_UNSUPPORTED_FEATURE is returned for those. + * - 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. Use zelUnloadDriverExt() with ZEL_UNLOAD_DRIVER_FLAG_FORCE to override. + * + * Handle lifetime after a successful unload: + * - The supplied driver handle remains a valid pointer and may be passed to zelReloadDriver(). + * Any Level Zero API called with it returns ZE_RESULT_ERROR_UNINITIALIZED until it is reloaded. + * - Every handle derived from the driver (devices, fabric vertices/edges, sysman and tools objects, + * and any child objects left live by a forced unload) is permanently dead. It remains a valid + * pointer, and every API called with it returns ZE_RESULT_ERROR_UNINITIALIZED forever. These + * handles are never rebound, because a reloaded driver may be a different build. + * + * @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_UNSUPPORTED_FEATURE if hDriver is a driver DDI handle rather than a loader handle. + * - 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 [PROOF OF CONCEPT] Unloads a single Level Zero driver, with flags. + * + * Equivalent to zelUnloadDriver() when flags is ZEL_UNLOAD_DRIVER_FLAG_NONE. + * + * With ZEL_UNLOAD_DRIVER_FLAG_FORCE the live-child-object check is skipped and the library is + * unmapped regardless. This exists for the case where a driver's kernel-mode component has been + * removed from under it: the user-mode driver is already unusable, and the application may hold + * objects it can never cleanly destroy. Any resources still owned by the driver are leaked, which + * is accepted because the library is going away. All of the application's outstanding handles for + * that driver are made permanently dead as described for zelUnloadDriver(). + * + * @param[in] hDriver + * The driver handle to unload, as returned by zeDriverGet() or zeInitDrivers(). + * @param[in] flags + * Combination of ::zel_unload_driver_flag_t. + * + * @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_UNSUPPORTED_FEATURE if hDriver is a driver DDI handle rather than a loader handle. + * - ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE if the driver still owns live child objects and + * ZEL_UNLOAD_DRIVER_FLAG_FORCE was not specified. + * - ZE_RESULT_ERROR_UNINITIALIZED if the loader has not been initialized. + */ +ZE_APIEXPORT ze_result_t ZE_APICALL +zelUnloadDriverExt( + ze_driver_handle_t hDriver, + zel_unload_driver_flags_t flags); + +/** + * @brief [PROOF OF CONCEPT] Reloads a driver previously unloaded by zelUnloadDriver(). + * + * The driver's library is loaded again and every DDI table is rebuilt from scratch, so a newer + * build of the driver, with different interfaces, is picked up in full. No function pointer, + * property or handle from the previous load survives. + * + * The supplied driver handle is rebound onto the freshly loaded driver and is valid again on + * success, so an application that cached it does not have to re-enumerate to keep working. + * Handles below driver level are NOT rebound: the application must call zeDeviceGet() (and any + * other enumeration it depends on) again. Handles obtained before the unload stay permanently + * dead and return ZE_RESULT_ERROR_UNINITIALIZED. + * + * Reload is the only way to bring an unloaded driver back; zeInit(), zeInitDrivers() and + * zeDriverGet() deliberately continue to skip an unloaded slot so that an unrelated component + * cannot silently resurrect a driver the application chose to unload. A failed reload leaves the + * driver unloaded and may be retried. + * + * @param[in] hDriver + * The driver handle that was passed to zelUnloadDriver()/zelUnloadDriverExt(). + * + * @return + * - ZE_RESULT_SUCCESS if the driver was reloaded and hDriver is valid again. + * - ZE_RESULT_ERROR_INVALID_NULL_HANDLE if hDriver is NULL or is not a handle for an unloaded driver. + * - ZE_RESULT_ERROR_INVALID_ARGUMENT if the driver identified by hDriver is not unloaded. + * - ZE_RESULT_ERROR_UNINITIALIZED if the library could not be loaded or failed to initialize. + */ +ZE_APIEXPORT ze_result_t ZE_APICALL +zelReloadDriver( + ze_driver_handle_t hDriver); + /** * @brief Notifies the loader that a driver has been removed and forces prevention of subsequent API calls. * diff --git a/scripts/templates/ldrddi.cpp.mako b/scripts/templates/ldrddi.cpp.mako index e5c3600b..41ff7137 100644 --- a/scripts/templates/ldrddi.cpp.mako +++ b/scripts/templates/ldrddi.cpp.mako @@ -18,6 +18,8 @@ from templates import helper as th */ #include "${x}_loader_internal.h" +#include + using namespace loader_driver_ddi; namespace loader @@ -97,6 +99,9 @@ namespace loader for( auto& drv : *loader::context->sysmanInstanceDrivers ) %endif { + if (drv.slotState == driver_slot_state_t::Unloaded) { + continue; // Skipped, not blacklisted -- only zelReloadDriver brings the slot back. + } %if re.match(r"Init", obj['name']) and namespace == "zes": if(drv.initStatus != ZE_RESULT_SUCCESS || drv.initSysManStatus != ZE_RESULT_SUCCESS) continue; @@ -150,6 +155,12 @@ namespace loader uint32_t total_driver_handle_count = 0; %if re.match(r"\w+InitDrivers$", th.make_func_name(n, tags, obj)): for( auto& drv : loader::context->zeDrivers ) { + if (drv.slotState == driver_slot_state_t::Unloaded) { + // Skipped, not blacklisted. An unrelated component calling zeInitDrivers must not + // silently resurrect a driver the application deliberately unloaded; only an + // explicit zelReloadDriver brings the slot back. + continue; + } if (!drv.handle || !drv.ddiInitialized) { auto res = loader::context->init_driver( drv, 0, desc); if (res != ZE_RESULT_SUCCESS || drv.zeddiInitResult != ZE_RESULT_SUCCESS) { @@ -167,12 +178,16 @@ namespace loader %if not re.match(r"\w+InitDrivers$", th.make_func_name(n, tags, obj)): std::call_once(loader::context->coreDriverSortOnce, []() { loader::context->driverSorting(&loader::context->zeDrivers, nullptr, false); - loader::defaultZerDdiTable = &loader::context->zeDrivers.front().dditable.zer; + // Never read zeDrivers.front() blindly: slot 0 may have been unloaded, and the + // zer entry points dereference this table without a null check. + loader::context->refreshDefaultZerDdiTable(); }); %else: std::call_once(loader::context->coreDriverSortOnce, [desc]() { loader::context->driverSorting(&loader::context->zeDrivers, desc, false); - loader::defaultZerDdiTable = &loader::context->zeDrivers.front().dditable.zer; + // Never read zeDrivers.front() blindly: slot 0 may have been unloaded, and the + // zer entry points dereference this table without a null check. + loader::context->refreshDefaultZerDdiTable(); }); %endif %else: @@ -190,6 +205,9 @@ namespace loader for( auto& drv : *loader::context->sysmanInstanceDrivers ) %endif { + if (drv.slotState == driver_slot_state_t::Unloaded) { + continue; // An unloaded slot is a hole in the list; it is never enumerated. + } %if not (re.match(r"\w+InitDrivers$", th.make_func_name(n, tags, obj))) and namespace != "zes": if(drv.initStatus != ZE_RESULT_SUCCESS || !drv.ddiInitialized) continue; @@ -297,13 +315,25 @@ namespace loader } drv.driverDDIHandleSupportQueried = true; } - if (!(drv.properties.flags & ZE_DRIVER_DDI_HANDLE_EXT_FLAG_DDI_HANDLE_EXT_SUPPORTED) || !loader::context->driverDDIPathDefault) { + // wrapperModePinned keeps a slot that has already handed out wrapper + // handles on the wrapper path. A reloaded, newer UMD may start advertising + // ZE_DRIVER_DDI_HANDLE_EXT; switching to raw handles at that point would + // invalidate the driver handle the application is still holding. + if (!(drv.properties.flags & ZE_DRIVER_DDI_HANDLE_EXT_FLAG_DDI_HANDLE_EXT_SUPPORTED) || !loader::context->driverDDIPathDefault || drv.wrapperModePinned) { if (loader::context->debugTraceEnabled) { std::string message = "Driver DDI Handles Not Supported for " + drv.name; loader::context->debug_trace_message(message, ""); } - ${obj['params'][1]['name']}[ driver_index ] = reinterpret_cast<${n}_driver_handle_t>( - context->${n}_driver_factory.getInstance( ${obj['params'][1]['name']}[ driver_index ], &drv.dditable ) ); + auto driverObject = context->${n}_driver_factory.getInstance( ${obj['params'][1]['name']}[ driver_index ], &drv.dditable ); + // Record every wrapper this slot issues. It is what lets a user-facing + // handle be resolved back to its slot exactly -- the wrapping decision + // is per driver, so intercept_enabled is not a reliable test -- and what + // zelReloadDriver rebinds onto the freshly loaded driver. + if (std::find(drv.${n}DriverObjects.begin(), drv.${n}DriverObjects.end(), driverObject) == drv.${n}DriverObjects.end()) { + drv.${n}DriverObjects.push_back(driverObject); + } + drv.wrapperModePinned = true; + ${obj['params'][1]['name']}[ driver_index ] = reinterpret_cast<${n}_driver_handle_t>( driverObject ); if (drv.zerDriverHandle != nullptr) { drv.zerDriverHandle = ${obj['params'][1]['name']}[ driver_index ]; } @@ -314,8 +344,16 @@ namespace loader } } %else: - ${obj['params'][1]['name']}[ driver_index ] = reinterpret_cast<${n}_driver_handle_t>( - context->${n}_driver_factory.getInstance( ${obj['params'][1]['name']}[ driver_index ], &drv.dditable ) ); + { + auto driverObject = context->${n}_driver_factory.getInstance( ${obj['params'][1]['name']}[ driver_index ], &drv.dditable ); + // See the core path: the slot needs to be able to find its own wrappers + // again to rebind them after zelReloadDriver. + if (std::find(drv.${n}DriverObjects.begin(), drv.${n}DriverObjects.end(), driverObject) == drv.${n}DriverObjects.end()) { + drv.${n}DriverObjects.push_back(driverObject); + } + drv.wrapperModePinned = true; + ${obj['params'][1]['name']}[ driver_index ] = reinterpret_cast<${n}_driver_handle_t>( driverObject ); + } %endif } } @@ -335,10 +373,8 @@ namespace loader result = ${X}_RESULT_SUCCESS; } %if namespace != "zes": - if (loader::context->zeDrivers.front().zerDriverDDISupported) - loader::context->defaultZerDriverHandle = loader::context->zeDrivers.front().zerDriverHandle; - else - loader::context->defaultZerDriverHandle = nullptr; + // Pick the default ZER driver from a slot that is actually loaded; slot 0 may be unloaded. + loader::context->refreshDefaultZerDdiTable(); %endif %else: diff --git a/scripts/templates/ze_loader_internal.h.mako b/scripts/templates/ze_loader_internal.h.mako index 4056e34a..1daa02aa 100644 --- a/scripts/templates/ze_loader_internal.h.mako +++ b/scripts/templates/ze_loader_internal.h.mako @@ -52,6 +52,19 @@ namespace loader ZEL_DRIVER_TYPE_FORCE_UINT32 = 0x7fffffff } zel_driver_type_t; + /////////////////////////////////////////////////////////////////////////////// + /// @brief Lifecycle state of a driver slot. + /// @details A slot is an index into the driver vectors and never moves: object_t + /// instances handed to the application store &driver_t::dditable, so erasing, + /// reallocating or re-sorting the vectors would dangle every live handle of every + /// driver. A slot that is unloaded is therefore emptied in place and reused. + enum class driver_slot_state_t + { + Discovered = 0, ///< Discovered but not yet dlopen'd, or dlopen'd and never initialized. + Loaded, ///< Library is mapped and its DDI tables are populated. + Unloaded ///< Explicitly unloaded via zelUnloadDriver. Skipped by implicit + ///< enumeration, but reloadable via zelReloadDriver -- not blacklisted. + }; ////////////////////////////////////////////////////////////////////////// struct driver_t { @@ -76,6 +89,29 @@ 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; + // Number of LOAD_DRIVER_LIBRARY references this copy owns on `handle`. The module is + // mapped once per copy that loaded it (zeDrivers via zeInit/zeInitDrivers, zesDrivers via + // zesInit); the copies made at discovery time alias the module without owning a reference. + // Unload must release exactly this many or the library stays mapped and "reload to a fresh + // state" is a lie. + uint32_t libraryLoadCount = 0; + // Snapshot taken at unload so reloadDriver restores exactly the copies that were live. + uint32_t reloadLibraryLoadCount = 0; + bool reloadDdiInitialized = false; + // Lifecycle of this slot. Unloaded slots are skipped by zeInit/zeInitDrivers/ + // zeDriverGet/driverSorting/driverOrdering, but remain eligible for zelReloadDriver. + driver_slot_state_t slotState = driver_slot_state_t::Discovered; + // Wrapper objects issued to the application for this slot, in issue order. Used to + // resolve a user-facing handle back to its slot exactly (the wrapping decision is per + // driver, so intercept_enabled is not a reliable test), and to rebind on reload. + // Raw pointers only: driver_t must stay copyable. Ownership while a slot is unloaded + // lives in context_t::parkedZeDriverObjects / parkedZesDriverObjects. + std::vector zeDriverObjects; + std::vector zesDriverObjects; + // Once this slot has issued wrapper handles it keeps issuing wrapper handles, even if a + // reloaded (possibly newer) driver starts advertising ZE_DRIVER_DDI_HANDLE_EXT. Without + // this the application's driver handle could not survive a reload. + bool wrapperModePinned = false; }; using driver_vector_t = std::vector< driver_t >; @@ -126,6 +162,29 @@ 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); + + // Unload a single driver identified by its user-facing handle. With + // ZEL_UNLOAD_DRIVER_FLAG_FORCE the live-child-object check is skipped, which is what a + // driver whose kernel-mode component has been removed requires: the application may hold + // objects it can never cleanly destroy. + ze_result_t unloadDriver(ze_driver_handle_t hDriver, zel_unload_driver_flags_t flags); + // Reload a previously unloaded driver into the same slot and rebind the application's + // original driver handle onto the freshly loaded driver. The reloaded driver may be a + // different (newer) build, so every DDI table is rebuilt from scratch and no handle below + // driver level survives. + ze_result_t reloadDriver(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); + // Locate the slot in zeDrivers backing a user-facing driver handle. Returns nullptr when + // the handle is not a loader wrapper issued for one of our slots. + driver_t *findDriverSlot(ze_driver_handle_t hDriver); + // Move every wrapper object belonging to the given slot out of its factory and onto the + // dead dispatch table, so calls through stale handles fail with UNINITIALIZED instead of + // reaching a reloaded driver with a raw handle from the previous load. + void retireDriverChildObjects(const driver_t &driver); + // Point defaultZerDdiTable/defaultZerDriverHandle at the first loaded slot, or at the dead + // table when nothing is loaded. + void refreshDefaultZerDdiTable(); ~context_t(); bool intercept_enabled = false; bool debugTraceEnabled = false; @@ -141,9 +200,24 @@ namespace loader dditable_t tracing_dditable = {}; std::shared_ptr zel_logger; ze_driver_handle_t defaultZerDriverHandle = nullptr; + // Wrapper objects whose driver has been unloaded. They are kept alive (the application may + // still hold their handles) but point at loader::deadDditable, so every entry point returns + // UNINITIALIZED rather than faulting. Child objects are never resurrected: a reloaded + // driver may be a different build, and its raw handles must never be confused with these. + std::vector> retiredHandleObjects; + // Driver wrappers held across an unload, awaiting rebinding by zelReloadDriver. + std::vector> parkedZeDriverObjects; + std::vector> parkedZesDriverObjects; }; extern ze_handle_t* loaderDispatch; extern zer_dditable_t* defaultZerDdiTable; extern context_t *context; + /////////////////////////////////////////////////////////////////////////////// + /// @brief A permanently zeroed dispatch table. + /// @details Every generated intercept function reads its pfn out of object_t::dditable and + /// returns UNINITIALIZED when that pfn is null, so pointing a retired wrapper here makes all + /// of its entry points fail safely with no generated-code changes. Zero-initialized with + /// static storage duration, so it is valid before and after loader construction/teardown. + extern dditable_t deadDditable; } diff --git a/source/inc/ze_singleton.h b/source/inc/ze_singleton.h index 7f7d09a0..a470fa51 100644 --- a/source/inc/ze_singleton.h +++ b/source/inc/ze_singleton.h @@ -9,6 +9,7 @@ #include #include #include +#include #include ////////////////////////////////////////////////////////////////////////// @@ -70,6 +71,62 @@ 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 + size_t countByDditable( const _dditable_t* dditable ) + { + std::lock_guard lk( mut ); + size_t count = 0; + for( const auto& entry : map ) + { + if( entry.second && entry.second->dditable == dditable ) + ++count; + } + return count; + } + + ////////////////////////////////////////////////////////////////////////// + /// removes and returns every instance whose dditable pointer matches the + /// argument. Used when a driver is unloaded: the objects must outlive the + /// unload (the application may still hold their handles) but they must leave + /// the map so their stale keys cannot collide with the raw handles handed out + /// by a freshly reloaded driver. + template + std::vector extractByDditable( const _dditable_t* dditable ) + { + std::lock_guard lk( mut ); + std::vector extracted; + for( auto iter = map.begin(); iter != map.end(); ) + { + if( iter->second && iter->second->dditable == dditable ) + { + extracted.push_back( std::move( iter->second ) ); + iter = map.erase( iter ); + } + else + { + ++iter; + } + } + return extracted; + } + + ////////////////////////////////////////////////////////////////////////// + /// reinserts an existing instance under a new key. Used to rebind a driver's + /// wrapper object onto the raw handle reported by a reloaded driver. + void adopt( _key_t _key, ptr_t _ptr ) + { + auto key = getKey( _key ); + + if( key == 0 || !_ptr ) + return; + + std::lock_guard lk( mut ); + map[ key ] = std::move( _ptr ); + } + ////////////////////////////////////////////////////////////////////////// /// once the key is no longer valid, release the singleton void release( _key_t _key ) diff --git a/source/lib/ze_lib.cpp b/source/lib/ze_lib.cpp index 14e20511..b6fb0700 100644 --- a/source/lib/ze_lib.cpp +++ b/source/lib/ze_lib.cpp @@ -462,6 +462,61 @@ 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( + 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 +zelUnloadDriverExt( + ze_driver_handle_t hDriver, + zel_unload_driver_flags_t flags) +{ +#ifdef L0_STATIC_LOADER_BUILD + if(nullptr == ze_lib::context->loader) + return ZE_RESULT_ERROR_UNINITIALIZED; + typedef ze_result_t (ZE_APICALL *zelUnloadDriverExtInternal_t)(ze_driver_handle_t hDriver, zel_unload_driver_flags_t flags); + auto unloadDriverExt = reinterpret_cast( + GET_FUNCTION_PTR(ze_lib::context->loader, "zelUnloadDriverExtInternal") ); + if (nullptr == unloadDriverExt) + return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE; + return unloadDriverExt(hDriver, flags); +#else + return zelUnloadDriverExtInternal(hDriver, flags); +#endif +} + +ze_result_t ZE_APICALL +zelReloadDriver( + 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 *zelReloadDriverInternal_t)(ze_driver_handle_t hDriver); + auto reloadDriver = reinterpret_cast( + GET_FUNCTION_PTR(ze_lib::context->loader, "zelReloadDriverInternal") ); + if (nullptr == reloadDriver) + return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE; + return reloadDriver(hDriver); +#else + return zelReloadDriverInternal(hDriver); +#endif +} + ze_result_t ZE_APICALL zelSetDriverTeardown() { diff --git a/source/loader/ze_ldrddi.cpp b/source/loader/ze_ldrddi.cpp index 2b51dad8..a78ddc41 100644 --- a/source/loader/ze_ldrddi.cpp +++ b/source/loader/ze_ldrddi.cpp @@ -9,6 +9,8 @@ */ #include "ze_loader_internal.h" +#include + using namespace loader_driver_ddi; namespace loader @@ -166,6 +168,9 @@ namespace loader bool atLeastOneDriverValid = false; for( auto& drv : loader::context->zeDrivers ) { + if (drv.slotState == driver_slot_state_t::Unloaded) { + continue; // Skipped, not blacklisted -- only zelReloadDriver brings the slot back. + } if(drv.initStatus != ZE_RESULT_SUCCESS) continue; if (!drv.handle || !drv.ddiInitialized) { @@ -214,7 +219,9 @@ namespace loader if (!loader::context->sortingInProgress.exchange(true) && !loader::context->instrumentationEnabled) { std::call_once(loader::context->coreDriverSortOnce, []() { loader::context->driverSorting(&loader::context->zeDrivers, nullptr, false); - loader::defaultZerDdiTable = &loader::context->zeDrivers.front().dditable.zer; + // Never read zeDrivers.front() blindly: slot 0 may have been unloaded, and the + // zer entry points dereference this table without a null check. + loader::context->refreshDefaultZerDdiTable(); }); loader::context->sortingInProgress.store(false); } @@ -222,6 +229,9 @@ namespace loader for( auto& drv : loader::context->zeDrivers ) { + if (drv.slotState == driver_slot_state_t::Unloaded) { + continue; // An unloaded slot is a hole in the list; it is never enumerated. + } if(drv.initStatus != ZE_RESULT_SUCCESS || !drv.ddiInitialized) continue; @@ -301,13 +311,25 @@ namespace loader } drv.driverDDIHandleSupportQueried = true; } - if (!(drv.properties.flags & ZE_DRIVER_DDI_HANDLE_EXT_FLAG_DDI_HANDLE_EXT_SUPPORTED) || !loader::context->driverDDIPathDefault) { + // wrapperModePinned keeps a slot that has already handed out wrapper + // handles on the wrapper path. A reloaded, newer UMD may start advertising + // ZE_DRIVER_DDI_HANDLE_EXT; switching to raw handles at that point would + // invalidate the driver handle the application is still holding. + if (!(drv.properties.flags & ZE_DRIVER_DDI_HANDLE_EXT_FLAG_DDI_HANDLE_EXT_SUPPORTED) || !loader::context->driverDDIPathDefault || drv.wrapperModePinned) { if (loader::context->debugTraceEnabled) { std::string message = "Driver DDI Handles Not Supported for " + drv.name; loader::context->debug_trace_message(message, ""); } - phDrivers[ driver_index ] = reinterpret_cast( - context->ze_driver_factory.getInstance( phDrivers[ driver_index ], &drv.dditable ) ); + auto driverObject = context->ze_driver_factory.getInstance( phDrivers[ driver_index ], &drv.dditable ); + // Record every wrapper this slot issues. It is what lets a user-facing + // handle be resolved back to its slot exactly -- the wrapping decision + // is per driver, so intercept_enabled is not a reliable test -- and what + // zelReloadDriver rebinds onto the freshly loaded driver. + if (std::find(drv.zeDriverObjects.begin(), drv.zeDriverObjects.end(), driverObject) == drv.zeDriverObjects.end()) { + drv.zeDriverObjects.push_back(driverObject); + } + drv.wrapperModePinned = true; + phDrivers[ driver_index ] = reinterpret_cast( driverObject ); if (drv.zerDriverHandle != nullptr) { drv.zerDriverHandle = phDrivers[ driver_index ]; } @@ -334,10 +356,8 @@ namespace loader if (total_driver_handle_count > 0) { result = ZE_RESULT_SUCCESS; } - if (loader::context->zeDrivers.front().zerDriverDDISupported) - loader::context->defaultZerDriverHandle = loader::context->zeDrivers.front().zerDriverHandle; - else - loader::context->defaultZerDriverHandle = nullptr; + // Pick the default ZER driver from a slot that is actually loaded; slot 0 may be unloaded. + loader::context->refreshDefaultZerDdiTable(); return result; } @@ -362,6 +382,12 @@ namespace loader uint32_t total_driver_handle_count = 0; for( auto& drv : loader::context->zeDrivers ) { + if (drv.slotState == driver_slot_state_t::Unloaded) { + // Skipped, not blacklisted. An unrelated component calling zeInitDrivers must not + // silently resurrect a driver the application deliberately unloaded; only an + // explicit zelReloadDriver brings the slot back. + continue; + } if (!drv.handle || !drv.ddiInitialized) { auto res = loader::context->init_driver( drv, 0, desc); if (res != ZE_RESULT_SUCCESS || drv.zeddiInitResult != ZE_RESULT_SUCCESS) { @@ -376,7 +402,9 @@ namespace loader if (!loader::context->sortingInProgress.exchange(true) && !loader::context->instrumentationEnabled) { std::call_once(loader::context->coreDriverSortOnce, [desc]() { loader::context->driverSorting(&loader::context->zeDrivers, desc, false); - loader::defaultZerDdiTable = &loader::context->zeDrivers.front().dditable.zer; + // Never read zeDrivers.front() blindly: slot 0 may have been unloaded, and the + // zer entry points dereference this table without a null check. + loader::context->refreshDefaultZerDdiTable(); }); loader::context->sortingInProgress.store(false); } @@ -384,6 +412,9 @@ namespace loader for( auto& drv : loader::context->zeDrivers ) { + if (drv.slotState == driver_slot_state_t::Unloaded) { + continue; // An unloaded slot is a hole in the list; it is never enumerated. + } if (!drv.ddiInitialized || !drv.dditable.ze.Global.pfnInitDrivers) { drv.initDriversStatus = ZE_RESULT_ERROR_UNINITIALIZED; result = ZE_RESULT_ERROR_UNINITIALIZED; @@ -466,13 +497,25 @@ namespace loader } drv.driverDDIHandleSupportQueried = true; } - if (!(drv.properties.flags & ZE_DRIVER_DDI_HANDLE_EXT_FLAG_DDI_HANDLE_EXT_SUPPORTED) || !loader::context->driverDDIPathDefault) { + // wrapperModePinned keeps a slot that has already handed out wrapper + // handles on the wrapper path. A reloaded, newer UMD may start advertising + // ZE_DRIVER_DDI_HANDLE_EXT; switching to raw handles at that point would + // invalidate the driver handle the application is still holding. + if (!(drv.properties.flags & ZE_DRIVER_DDI_HANDLE_EXT_FLAG_DDI_HANDLE_EXT_SUPPORTED) || !loader::context->driverDDIPathDefault || drv.wrapperModePinned) { if (loader::context->debugTraceEnabled) { std::string message = "Driver DDI Handles Not Supported for " + drv.name; loader::context->debug_trace_message(message, ""); } - phDrivers[ driver_index ] = reinterpret_cast( - context->ze_driver_factory.getInstance( phDrivers[ driver_index ], &drv.dditable ) ); + auto driverObject = context->ze_driver_factory.getInstance( phDrivers[ driver_index ], &drv.dditable ); + // Record every wrapper this slot issues. It is what lets a user-facing + // handle be resolved back to its slot exactly -- the wrapping decision + // is per driver, so intercept_enabled is not a reliable test -- and what + // zelReloadDriver rebinds onto the freshly loaded driver. + if (std::find(drv.zeDriverObjects.begin(), drv.zeDriverObjects.end(), driverObject) == drv.zeDriverObjects.end()) { + drv.zeDriverObjects.push_back(driverObject); + } + drv.wrapperModePinned = true; + phDrivers[ driver_index ] = reinterpret_cast( driverObject ); if (drv.zerDriverHandle != nullptr) { drv.zerDriverHandle = phDrivers[ driver_index ]; } @@ -499,10 +542,8 @@ namespace loader if (total_driver_handle_count > 0) { result = ZE_RESULT_SUCCESS; } - if (loader::context->zeDrivers.front().zerDriverDDISupported) - loader::context->defaultZerDriverHandle = loader::context->zeDrivers.front().zerDriverHandle; - else - loader::context->defaultZerDriverHandle = nullptr; + // Pick the default ZER driver from a slot that is actually loaded; slot 0 may be unloaded. + loader::context->refreshDefaultZerDdiTable(); return result; } diff --git a/source/loader/ze_loader.cpp b/source/loader/ze_loader.cpp index 14a0ef27..f53a8fec 100644 --- a/source/loader/ze_loader.cpp +++ b/source/loader/ze_loader.cpp @@ -8,6 +8,7 @@ #include "ze_loader_utils.h" #include "driver_discovery.h" +#include #include #include @@ -106,6 +107,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.slotState == driver_slot_state_t::Unloaded) { + continue; // Unloaded slots are not eligible for type/index based ordering. + } switch (driver.driverType) { case ZEL_DRIVER_TYPE_DISCRETE_GPU: discreteGPUDrivers.push_back(driver); @@ -147,6 +151,7 @@ namespace loader switch (spec.type) { case DriverOrderSpecType::BY_GLOBAL_INDEX: if (spec.globalIndex < originalDrivers.size() && + originalDrivers[spec.globalIndex].slotState != driver_slot_state_t::Unloaded && usedGlobalIndices.find(spec.globalIndex) == usedGlobalIndices.end()) { orderedDrivers.push_back(originalDrivers[spec.globalIndex]); usedGlobalIndices.insert(spec.globalIndex); @@ -253,6 +258,9 @@ namespace loader for (auto &driver : *drivers) { uint32_t pCount = 0; std::vector driverHandles; + if (driver.slotState == driver_slot_state_t::Unloaded) { + continue; // Unloaded slots 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 +488,13 @@ 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; + // Implicit paths never resurrect an unloaded slot: an unrelated component calling + // zeInitDrivers must not silently bring back a driver the application deliberately + // unloaded. This is a state, not a blacklist -- zelReloadDriver clears it and loads the + // library itself before reaching here. + if (driver.slotState == driver_slot_state_t::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, ""); @@ -523,6 +538,9 @@ namespace loader #endif } driver.handle = handle; + // This copy now owns a dlopen/LoadLibrary reference of its own; unload must + // release it or the module stays mapped. + driver.libraryLoadCount++; } else { std::string loadLibraryErrorValue; GET_LIBRARY_ERROR(loadLibraryErrorValue); @@ -586,6 +604,7 @@ namespace loader } driver.ddiInitialized = true; + driver.slotState = driver_slot_state_t::Loaded; } if (!driver.handle && !driver.ddiInitialized) { @@ -712,6 +731,7 @@ namespace loader } allDrivers.emplace_back(); allDrivers.rbegin()->handle = handle; + allDrivers.rbegin()->libraryLoadCount = 1; allDrivers.rbegin()->name = "ze_null"; } else if (debugTraceEnabled) { GET_LIBRARY_ERROR(loadLibraryErrorValue); @@ -740,6 +760,7 @@ namespace loader } allDrivers.emplace_back(); allDrivers.rbegin()->handle = handle; + allDrivers.rbegin()->libraryLoadCount = 1; allDrivers.rbegin()->name = driverInfo.path; allDrivers.rbegin()->customDriver = driverInfo.customDriver; } else if (debugTraceEnabled) { @@ -766,6 +787,10 @@ namespace loader } std::copy(allDrivers.begin(), allDrivers.end(), std::back_inserter(zeDrivers)); std::copy(allDrivers.begin(), allDrivers.end(), std::back_inserter(zesDrivers)); + // The copies alias allDrivers' module handle but do not own a reference to it. Only the + // copy that actually called LOAD_DRIVER_LIBRARY may release one. + for (auto &drv : zeDrivers) drv.libraryLoadCount = 0; + for (auto &drv : zesDrivers) drv.libraryLoadCount = 0; typedef ze_result_t (ZE_APICALL *getVersion_t)(zel_component_version_t *version); if( getenv_tobool( "ZE_ENABLE_VALIDATION_LAYER" ) ) @@ -949,6 +974,550 @@ namespace loader } }; + // A permanently zeroed dispatch table. Every generated intercept function pulls its function + // pointer out of the object's dditable and returns UNINITIALIZED when that pointer is null, so + // aiming a retired wrapper here makes all ~340 entry points fail safely with no generated-code + // changes. It must be a real object and never nullptr: the generated code dereferences the + // dditable pointer itself without a check. + dditable_t deadDditable = {}; + + bool context_t::isDriverInUse(const dditable_t *dditable) + { + // 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; + } + + driver_t *context_t::findDriverSlot(ze_driver_handle_t hDriver) + { + if (nullptr == hDriver) { + return nullptr; + } + // Compare against the wrapper objects this loader actually issued. Testing intercept_enabled + // and casting is not sufficient: whether a handle is wrapped is decided per driver by + // ZE_DRIVER_DDI_HANDLE_EXT support, so an unwrapped handle would be a wild read. + auto zeObj = reinterpret_cast(hDriver); + auto zesObj = reinterpret_cast(hDriver); + for (auto *vec : { &zeDrivers, &zesDrivers }) { + for (auto &drv : *vec) { + for (auto *obj : drv.zeDriverObjects) { + if (obj == zeObj) return &drv; + } + for (auto *obj : drv.zesDriverObjects) { + if (obj == zesObj) return &drv; + } + } + } + return nullptr; + } + + void context_t::retireDriverChildObjects(const driver_t &driver) + { + const dditable_t *ddi = &driver.dditable; + + // The reverse maps are keyed by the very wrapper objects about to be retired, so purge them + // first; the deleted-entry lookup in zeCommandListAppend* would otherwise read freed keys. + { + std::lock_guard lock(image_handle_map_lock); + for (auto it = image_handle_map.begin(); it != image_handle_map.end(); ) { + if (it->first && it->first->dditable == ddi) it = image_handle_map.erase(it); + else ++it; + } + } + { + std::lock_guard lock(sampler_handle_map_lock); + for (auto it = sampler_handle_map.begin(); it != sampler_handle_map.end(); ) { + if (it->first && it->first->dditable == ddi) it = sampler_handle_map.erase(it); + else ++it; + } + } + + // Everything below driver level is retired permanently -- never rebound. A reloaded driver + // may be a different, newer build with different interfaces, so handing it a device or + // context handle minted by the previous load would be a use-after-free at best. Most of + // these factories (devices, fabric vertices/edges, the zes and zet_metric families) never + // call release(), so without this they would outlive the library holding a stale raw + // pointer and a dditable aimed at a slot that has come back to life. + // + // The wrappers themselves must outlive the unload: the application still holds their + // handles. Extracting them from the factory frees the raw-handle key, so a reloaded library + // that reuses the same pointer values cannot collide with a stale entry. +#define RETIRE( factory ) \ + for (auto &obj : factory.extractByDditable(ddi)) { \ + obj->handle = nullptr; \ + obj->dditable = &deadDditable; \ + retiredHandleObjects.emplace_back(std::move(obj)); \ + } + + RETIRE(ze_command_list_factory); + RETIRE(ze_command_queue_factory); + RETIRE(ze_context_factory); + RETIRE(ze_device_factory); + RETIRE(ze_event_factory); + RETIRE(ze_event_pool_factory); + RETIRE(ze_executable_graph_factory); + RETIRE(ze_external_semaphore_ext_factory); + RETIRE(ze_fabric_edge_factory); + RETIRE(ze_fabric_vertex_factory); + RETIRE(ze_fence_factory); + RETIRE(ze_graph_factory); + RETIRE(ze_image_factory); + RETIRE(ze_kernel_factory); + RETIRE(ze_module_build_log_factory); + RETIRE(ze_module_factory); + RETIRE(ze_physical_mem_factory); + RETIRE(ze_rtas_builder_exp_factory); + RETIRE(ze_rtas_builder_ext_factory); + RETIRE(ze_rtas_parallel_operation_exp_factory); + RETIRE(ze_rtas_parallel_operation_ext_factory); + RETIRE(ze_sampler_factory); + RETIRE(zes_device_factory); + RETIRE(zes_diag_factory); + RETIRE(zes_engine_factory); + RETIRE(zes_fabric_port_factory); + RETIRE(zes_fan_factory); + RETIRE(zes_firmware_factory); + RETIRE(zes_freq_factory); + RETIRE(zes_led_factory); + RETIRE(zes_mem_factory); + RETIRE(zes_overclock_factory); + RETIRE(zes_perf_factory); + RETIRE(zes_psu_factory); + RETIRE(zes_pwr_factory); + RETIRE(zes_ras_factory); + RETIRE(zes_sched_factory); + RETIRE(zes_standby_factory); + RETIRE(zes_temp_factory); + RETIRE(zes_vf_factory); + RETIRE(zet_command_list_factory); + RETIRE(zet_context_factory); + RETIRE(zet_debug_session_factory); + RETIRE(zet_device_factory); + RETIRE(zet_driver_factory); + RETIRE(zet_kernel_factory); + RETIRE(zet_metric_decoder_exp_factory); + RETIRE(zet_metric_factory); + RETIRE(zet_metric_group_factory); + RETIRE(zet_metric_programmable_exp_factory); + RETIRE(zet_metric_query_factory); + RETIRE(zet_metric_query_pool_factory); + RETIRE(zet_metric_streamer_factory); + RETIRE(zet_metric_tracer_exp_factory); + RETIRE(zet_module_factory); + RETIRE(zet_tracer_exp_factory); + +#undef RETIRE + } + + void context_t::refreshDefaultZerDdiTable() + { + for (auto &drv : zeDrivers) { + if (drv.slotState == driver_slot_state_t::Loaded) { + loader::defaultZerDdiTable = &drv.dditable.zer; + defaultZerDriverHandle = drv.zerDriverDDISupported ? drv.zerDriverHandle : nullptr; + return; + } + } + // Nothing is loaded. The zer entry points dereference this table without a null check, so + // it has to point at a real (zeroed) table rather than nullptr. + loader::defaultZerDdiTable = &deadDditable.zer; + defaultZerDriverHandle = nullptr; + } + + ze_result_t context_t::unloadDriver(ze_driver_handle_t hDriver, zel_unload_driver_flags_t flags) + { + if (nullptr == hDriver) { + return ZE_RESULT_ERROR_INVALID_NULL_HANDLE; + } + + std::lock_guard lock(sortMutex); + + driver_t *slot = findDriverSlot(hDriver); + if (!slot) { + // Not one of our wrappers. If it is the raw handle of a driver on the + // ZE_DRIVER_DDI_HANDLE_EXT path, say so explicitly: that handle is memory owned by the + // library we would unmap, so there is nothing the loader could keep valid across the + // unload and nothing to rebind on reload. + for (auto &drv : zeDrivers) { + if (drv.slotState != driver_slot_state_t::Unloaded && drv.zerDriverHandle == hDriver) { + if (debugTraceEnabled) { + debug_trace_message("zelUnloadDriver: driver uses DDI handles, cannot unload ", drv.name); + } + return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE; + } + } + if (debugTraceEnabled) { + debug_trace_message("zelUnloadDriver: driver handle not found", ""); + } + return ZE_RESULT_ERROR_INVALID_NULL_HANDLE; + } + + if (slot->slotState == driver_slot_state_t::Unloaded) { + return ZE_RESULT_SUCCESS; // idempotent + } + + const std::string targetName = slot->name; + + // Safety gate: refuse to unload a driver that still owns live child objects, across every + // copy of the slot. ZEL_UNLOAD_DRIVER_FLAG_FORCE skips it -- which is exactly what a driver + // whose kernel-mode component has been removed needs, since the application is left holding + // objects it can never cleanly destroy. Those objects are retired below, so the calls the + // application makes on them fail with UNINITIALIZED instead of faulting. + if (!(flags & ZEL_UNLOAD_DRIVER_FLAG_FORCE)) { + for (auto *vec : { &zeDrivers, &zesDrivers, &allDrivers }) { + for (auto &drv : *vec) { + if (drv.name == targetName && isDriverInUse(&drv.dditable)) { + if (debugTraceEnabled) { + debug_trace_message("zelUnloadDriver: driver still in use: ", targetName); + } + return ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE; + } + } + } + } + + HMODULE targetModule = nullptr; + uint32_t pendingFrees = 0; + + // Empty every copy of the slot in place. Slots are never erased and never re-sorted: an + // object_t stores &driver_t::dditable, a pointer *into* the driver vector, so relocating an + // element would dangle every live handle of every driver, not just this one. Reusing the + // slot at a fixed index is also what makes same-handle rebinding possible on reload -- + // refilling this dditable gives the parked wrapper an entirely new dispatch table at an + // unchanged address. + for (auto *vec : { &zeDrivers, &zesDrivers, &allDrivers }) { + for (auto &drv : *vec) { + if (drv.name != targetName) { + continue; + } + + retireDriverChildObjects(drv); + + // Park the driver wrappers rather than releasing them: the application still holds + // these handles, and reload rebinds them onto the fresh driver. Releasing here is a + // use-after-free the moment the application touches the handle again. + for (auto &obj : ze_driver_factory.extractByDditable(&drv.dditable)) { + obj->handle = nullptr; + parkedZeDriverObjects.emplace_back(std::move(obj)); + } + for (auto &obj : zes_driver_factory.extractByDditable(&drv.dditable)) { + obj->handle = nullptr; + parkedZesDriverObjects.emplace_back(std::move(obj)); + } + + if (!targetModule && drv.handle) { + targetModule = drv.handle; + } + // Release exactly as many references as this copy took. The module is mapped once + // per copy that loaded it; free too few and it stays resident, and "reload to a + // fresh state" would silently be a no-op. + pendingFrees += drv.libraryLoadCount; + drv.reloadLibraryLoadCount = drv.libraryLoadCount; + drv.reloadDdiInitialized = drv.ddiInitialized; + drv.libraryLoadCount = 0; + + // Zero the dispatch tables in place. The parked wrappers still point here, so every + // call through them now returns UNINITIALIZED until reload refills it. + drv.dditable = {}; + drv.properties = {}; + drv.handle = nullptr; + drv.zerDriverHandle = nullptr; + drv.zerDriverDDISupported = true; + 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; + // name, driverType, customDriver, pciOrderingRequested, wrapperModePinned and the + // wrapper registries are deliberately preserved: this is a state, not a verdict. + drv.slotState = driver_slot_state_t::Unloaded; + } + } + + if (targetModule) { + for (uint32_t i = 0; i < pendingFrees; ++i) { + 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); + } + } + } + } + + refreshDefaultZerDdiTable(); + + if (debugTraceEnabled) { + debug_trace_message("zelUnloadDriver: unloaded driver ", targetName); + } + + return ZE_RESULT_SUCCESS; + } + + ze_result_t context_t::reloadDriver(ze_driver_handle_t hDriver) + { + if (nullptr == hDriver) { + return ZE_RESULT_ERROR_INVALID_NULL_HANDLE; + } + + std::lock_guard lock(sortMutex); + + driver_t *slot = findDriverSlot(hDriver); + if (!slot) { + return ZE_RESULT_ERROR_INVALID_NULL_HANDLE; + } + if (slot->slotState != driver_slot_state_t::Unloaded) { + return ZE_RESULT_ERROR_INVALID_ARGUMENT; + } + + const std::string targetName = slot->name; + + // Take a reference of our own before touching any slot state, so a library that cannot be + // mapped leaves everything exactly as it was. The slot stays Unloaded and reload can be + // retried -- unloading a driver is a state, never a blacklist. + auto probeModule = LOAD_DRIVER_LIBRARY(targetName.c_str()); + if (nullptr == probeModule) { + if (debugTraceEnabled) { + std::string loadLibraryErrorValue; + GET_LIBRARY_ERROR(loadLibraryErrorValue); + debug_trace_message("zelReloadDriver: Load Library of " + targetName + " failed with ", loadLibraryErrorValue); + } + return ZE_RESULT_ERROR_UNINITIALIZED; + } + + ze_result_t result = ZE_RESULT_SUCCESS; + HMODULE reloadedModule = nullptr; + + for (auto *vec : { &zeDrivers, &zesDrivers, &allDrivers }) { + for (auto &drv : *vec) { + if (drv.name != targetName || drv.slotState != driver_slot_state_t::Unloaded) { + continue; + } + + drv.initStatus = ZE_RESULT_SUCCESS; + drv.initSysManStatus = ZE_RESULT_SUCCESS; + drv.initDriversStatus = ZE_RESULT_SUCCESS; + drv.slotState = driver_slot_state_t::Discovered; + + if (drv.reloadLibraryLoadCount == 0 && !drv.reloadDdiInitialized) { + continue; // discovery-time alias; fixed up once the module is known + } + + // Load here rather than letting init_driver do it: init_driver gates the load on + // driverType matching the requested flags, and a slot that was previously loaded + // must come back regardless of its type. With handle already set, init_driver goes + // straight to rebuilding every DDI table from scratch against the new module, so a + // newer UMD's interfaces are picked up wholesale and no pointer from the previous + // load survives. + if (!drv.handle) { + auto handle = LOAD_DRIVER_LIBRARY(drv.name.c_str()); + if (nullptr == handle) { + result = ZE_RESULT_ERROR_UNINITIALIZED; + continue; + } + drv.handle = handle; + drv.libraryLoadCount++; + } + reloadedModule = drv.handle; + + auto res = init_driver(drv, 0, nullptr); + if (res != ZE_RESULT_SUCCESS) { + if (debugTraceEnabled) { + debug_trace_message("zelReloadDriver: init_driver failed for " + targetName + " with ", loader::to_string(res)); + } + result = res; + } + } + } + + // Copies that never owned a reference still need to see the module so their (unused) state + // is consistent with the rest of the slot. + if (reloadedModule) { + for (auto *vec : { &zeDrivers, &zesDrivers, &allDrivers }) { + for (auto &drv : *vec) { + if (drv.name == targetName && !drv.handle) { + drv.handle = reloadedModule; + drv.slotState = driver_slot_state_t::Loaded; + } + } + } + } + + // Drop our own reference now that the slot copies hold theirs. + FREE_DRIVER_LIBRARY(probeModule); + + driver_t *zeSlot = nullptr; + driver_t *zesSlot = nullptr; + for (auto &drv : zeDrivers) { + if (drv.name == targetName) { zeSlot = &drv; break; } + } + for (auto &drv : *sysmanInstanceDrivers) { + if (drv.name == targetName) { zesSlot = &drv; break; } + } + + if (!zeSlot || !zeSlot->ddiInitialized) { + if (debugTraceEnabled) { + debug_trace_message("zelReloadDriver: DDI tables unavailable after reload of ", targetName); + } + return result == ZE_RESULT_SUCCESS ? ZE_RESULT_ERROR_UNINITIALIZED : result; + } + + // Bring the driver up and collect the raw handles it reports now. Preferring + // zeInitDrivers matches how the loader enumerates today; zeInit/zeDriverGet is the fallback + // for a driver that only implements the legacy path. + std::vector freshZeHandles; + { + ze_init_driver_type_desc_t permissiveDesc = {}; + permissiveDesc.stype = ZE_STRUCTURE_TYPE_INIT_DRIVER_TYPE_DESC; + permissiveDesc.pNext = nullptr; + permissiveDesc.flags = UINT32_MAX; + + uint32_t count = 0; + auto pfnInitDrivers = zeSlot->dditable.ze.Global.pfnInitDrivers; + if (pfnInitDrivers && ZE_RESULT_SUCCESS == pfnInitDrivers(&count, nullptr, &permissiveDesc) && count) { + freshZeHandles.resize(count); + if (ZE_RESULT_SUCCESS != pfnInitDrivers(&count, freshZeHandles.data(), &permissiveDesc)) { + freshZeHandles.clear(); + } else { + freshZeHandles.resize(count); + } + } + if (freshZeHandles.empty() && zeSlot->dditable.ze.Global.pfnInit && zeSlot->dditable.ze.Driver.pfnGet) { + if (ZE_RESULT_SUCCESS == zeSlot->dditable.ze.Global.pfnInit(0)) { + zeSlot->legacyInitAttempted = true; + count = 0; + if (ZE_RESULT_SUCCESS == zeSlot->dditable.ze.Driver.pfnGet(&count, nullptr) && count) { + freshZeHandles.resize(count); + if (ZE_RESULT_SUCCESS != zeSlot->dditable.ze.Driver.pfnGet(&count, freshZeHandles.data())) { + freshZeHandles.clear(); + } else { + freshZeHandles.resize(count); + } + } + } + } + } + + std::vector freshZesHandles; + if (zesSlot && zesSlot->ddiInitialized && zesSlot->dditable.zes.Global.pfnInit && + zesSlot->dditable.zes.Driver.pfnGet) { + if (ZE_RESULT_SUCCESS == zesSlot->dditable.zes.Global.pfnInit(0)) { + uint32_t count = 0; + if (ZE_RESULT_SUCCESS == zesSlot->dditable.zes.Driver.pfnGet(&count, nullptr) && count) { + freshZesHandles.resize(count); + if (ZE_RESULT_SUCCESS != zesSlot->dditable.zes.Driver.pfnGet(&count, freshZesHandles.data())) { + freshZesHandles.clear(); + } else { + freshZesHandles.resize(count); + } + } + } + } + + // Rebind the parked wrappers onto the fresh raw handles, in issue order. The wrapper + // pointer -- the value the application holds -- never changes; only what it dispatches to. + // Its dditable already points at this slot's storage, which init_driver has just refilled + // from the new module. Re-keying the factory under the fresh raw handle means a subsequent + // zeInitDrivers/zeDriverGet hands back the identical wrapper. + { + size_t next = 0; + std::vector stillLive; + for (auto *obj : zeSlot->zeDriverObjects) { + auto it = std::find_if(parkedZeDriverObjects.begin(), parkedZeDriverObjects.end(), + [obj](const std::unique_ptr &p) { return p.get() == obj; }); + if (it == parkedZeDriverObjects.end()) { + stillLive.push_back(obj); + continue; + } + if (next < freshZeHandles.size()) { + obj->handle = freshZeHandles[next++]; + obj->dditable = &zeSlot->dditable; + ze_driver_factory.adopt(obj->handle, std::move(*it)); + stillLive.push_back(obj); + } else { + // The reloaded driver reports fewer drivers than the previous build did. There + // is nothing to bind this wrapper to, so retire it permanently rather than + // leave it aimed at a live slot with a handle from the old load. + if (debugTraceEnabled) { + debug_trace_message("zelReloadDriver: retiring surplus driver handle for ", targetName); + } + obj->handle = nullptr; + obj->dditable = &deadDditable; + retiredHandleObjects.emplace_back(std::move(*it)); + } + parkedZeDriverObjects.erase(it); + } + zeSlot->zeDriverObjects = std::move(stillLive); + } + + if (zesSlot) { + size_t next = 0; + std::vector stillLive; + for (auto *obj : zesSlot->zesDriverObjects) { + auto it = std::find_if(parkedZesDriverObjects.begin(), parkedZesDriverObjects.end(), + [obj](const std::unique_ptr &p) { return p.get() == obj; }); + if (it == parkedZesDriverObjects.end()) { + stillLive.push_back(obj); + continue; + } + if (next < freshZesHandles.size()) { + obj->handle = freshZesHandles[next++]; + obj->dditable = &zesSlot->dditable; + zes_driver_factory.adopt(obj->handle, std::move(*it)); + stillLive.push_back(obj); + } else { + obj->handle = nullptr; + obj->dditable = &deadDditable; + retiredHandleObjects.emplace_back(std::move(*it)); + } + parkedZesDriverObjects.erase(it); + } + zesSlot->zesDriverObjects = std::move(stillLive); + } + + // A slot that has ever issued wrapper handles keeps issuing wrapper handles. If the + // reloaded UMD now advertises ZE_DRIVER_DDI_HANDLE_EXT and we switched to raw handles, the + // application's original driver handle could not be preserved across the reload at all. + for (auto *vec : { &zeDrivers, &zesDrivers, &allDrivers }) { + for (auto &drv : *vec) { + if (drv.name == targetName) { + drv.wrapperModePinned = true; + drv.slotState = driver_slot_state_t::Loaded; + } + } + } + + // Leave driverDDIHandleSupportQueried false so the next enumeration re-queries extension + // and driver properties from the new module rather than trusting the old build's answers. + refreshDefaultZerDdiTable(); + + if (debugTraceEnabled) { + debug_trace_message("zelReloadDriver: reloaded driver ", targetName); + } + + return result; + } + 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); diff --git a/source/loader/ze_loader_api.cpp b/source/loader/ze_loader_api.cpp index 514d6d31..48c3b696 100644 --- a/source/loader/ze_loader_api.cpp +++ b/source/loader/ze_loader_api.cpp @@ -292,6 +292,34 @@ zelLoaderTranslateHandleInternal( return ZE_RESULT_SUCCESS; } +ZE_DLLEXPORT ze_result_t ZE_APICALL +zelUnloadDriverExtInternal( + ze_driver_handle_t hDriver, + zel_unload_driver_flags_t flags) +{ + if (!loader::context) { + return ZE_RESULT_ERROR_UNINITIALIZED; + } + return loader::context->unloadDriver(hDriver, flags); +} + +ZE_DLLEXPORT ze_result_t ZE_APICALL +zelUnloadDriverInternal( + ze_driver_handle_t hDriver) +{ + return zelUnloadDriverExtInternal(hDriver, ZEL_UNLOAD_DRIVER_FLAG_NONE); +} + +ZE_DLLEXPORT ze_result_t ZE_APICALL +zelReloadDriverInternal( + ze_driver_handle_t hDriver) +{ + if (!loader::context) { + return ZE_RESULT_ERROR_UNINITIALIZED; + } + return loader::context->reloadDriver(hDriver); +} + #if defined(__cplusplus) } diff --git a/source/loader/ze_loader_api.h b/source/loader/ze_loader_api.h index bbf4c09f..14d0fee2 100644 --- a/source/loader/ze_loader_api.h +++ b/source/loader/ze_loader_api.h @@ -88,6 +88,48 @@ 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_UNSUPPORTED_FEATURE +/// - ::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 + + +/////////////////////////////////////////////////////////////////////////////// +/// @brief Proof-of-concept: unload a single driver, optionally forcing the unload +/// even when the driver still owns live child objects. +/// +/// @returns +/// - ::ZE_RESULT_SUCCESS +/// - ::ZE_RESULT_ERROR_INVALID_NULL_HANDLE +/// - ::ZE_RESULT_ERROR_UNSUPPORTED_FEATURE +/// - ::ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE +ZE_DLLEXPORT ze_result_t ZE_APICALL +zelUnloadDriverExtInternal( + ze_driver_handle_t hDriver, //Input: driver handle to unload + zel_unload_driver_flags_t flags); //Input: combination of ::zel_unload_driver_flag_t + + +/////////////////////////////////////////////////////////////////////////////// +/// @brief Proof-of-concept: reload a previously unloaded driver into its original +/// slot and rebind the supplied driver handle onto the fresh driver. +/// +/// @returns +/// - ::ZE_RESULT_SUCCESS +/// - ::ZE_RESULT_ERROR_INVALID_NULL_HANDLE +/// - ::ZE_RESULT_ERROR_INVALID_ARGUMENT +/// - ::ZE_RESULT_ERROR_UNINITIALIZED +ZE_DLLEXPORT ze_result_t ZE_APICALL +zelReloadDriverInternal( + ze_driver_handle_t hDriver); //Input: driver handle to reload + + #if defined(__cplusplus) } #endif \ No newline at end of file diff --git a/source/loader/ze_loader_internal.h b/source/loader/ze_loader_internal.h index 45dce29c..3cd11467 100644 --- a/source/loader/ze_loader_internal.h +++ b/source/loader/ze_loader_internal.h @@ -43,6 +43,19 @@ namespace loader ZEL_DRIVER_TYPE_FORCE_UINT32 = 0x7fffffff } zel_driver_type_t; + /////////////////////////////////////////////////////////////////////////////// + /// @brief Lifecycle state of a driver slot. + /// @details A slot is an index into the driver vectors and never moves: object_t + /// instances handed to the application store &driver_t::dditable, so erasing, + /// reallocating or re-sorting the vectors would dangle every live handle of every + /// driver. A slot that is unloaded is therefore emptied in place and reused. + enum class driver_slot_state_t + { + Discovered = 0, ///< Discovered but not yet dlopen'd, or dlopen'd and never initialized. + Loaded, ///< Library is mapped and its DDI tables are populated. + Unloaded ///< Explicitly unloaded via zelUnloadDriver. Skipped by implicit + ///< enumeration, but reloadable via zelReloadDriver -- not blacklisted. + }; ////////////////////////////////////////////////////////////////////////// struct driver_t { @@ -67,6 +80,29 @@ 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; + // Number of LOAD_DRIVER_LIBRARY references this copy owns on `handle`. The module is + // mapped once per copy that loaded it (zeDrivers via zeInit/zeInitDrivers, zesDrivers via + // zesInit); the copies made at discovery time alias the module without owning a reference. + // Unload must release exactly this many or the library stays mapped and "reload to a fresh + // state" is a lie. + uint32_t libraryLoadCount = 0; + // Snapshot taken at unload so reloadDriver restores exactly the copies that were live. + uint32_t reloadLibraryLoadCount = 0; + bool reloadDdiInitialized = false; + // Lifecycle of this slot. Unloaded slots are skipped by zeInit/zeInitDrivers/ + // zeDriverGet/driverSorting/driverOrdering, but remain eligible for zelReloadDriver. + driver_slot_state_t slotState = driver_slot_state_t::Discovered; + // Wrapper objects issued to the application for this slot, in issue order. Used to + // resolve a user-facing handle back to its slot exactly (the wrapping decision is per + // driver, so intercept_enabled is not a reliable test), and to rebind on reload. + // Raw pointers only: driver_t must stay copyable. Ownership while a slot is unloaded + // lives in context_t::parkedZeDriverObjects / parkedZesDriverObjects. + std::vector zeDriverObjects; + std::vector zesDriverObjects; + // Once this slot has issued wrapper handles it keeps issuing wrapper handles, even if a + // reloaded (possibly newer) driver starts advertising ZE_DRIVER_DDI_HANDLE_EXT. Without + // this the application's driver handle could not survive a reload. + bool wrapperModePinned = false; }; using driver_vector_t = std::vector< driver_t >; @@ -164,6 +200,29 @@ 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); + + // Unload a single driver identified by its user-facing handle. With + // ZEL_UNLOAD_DRIVER_FLAG_FORCE the live-child-object check is skipped, which is what a + // driver whose kernel-mode component has been removed requires: the application may hold + // objects it can never cleanly destroy. + ze_result_t unloadDriver(ze_driver_handle_t hDriver, zel_unload_driver_flags_t flags); + // Reload a previously unloaded driver into the same slot and rebind the application's + // original driver handle onto the freshly loaded driver. The reloaded driver may be a + // different (newer) build, so every DDI table is rebuilt from scratch and no handle below + // driver level survives. + ze_result_t reloadDriver(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); + // Locate the slot in zeDrivers backing a user-facing driver handle. Returns nullptr when + // the handle is not a loader wrapper issued for one of our slots. + driver_t *findDriverSlot(ze_driver_handle_t hDriver); + // Move every wrapper object belonging to the given slot out of its factory and onto the + // dead dispatch table, so calls through stale handles fail with UNINITIALIZED instead of + // reaching a reloaded driver with a raw handle from the previous load. + void retireDriverChildObjects(const driver_t &driver); + // Point defaultZerDdiTable/defaultZerDriverHandle at the first loaded slot, or at the dead + // table when nothing is loaded. + void refreshDefaultZerDdiTable(); ~context_t(); bool intercept_enabled = false; bool debugTraceEnabled = false; @@ -179,9 +238,24 @@ namespace loader dditable_t tracing_dditable = {}; std::shared_ptr zel_logger; ze_driver_handle_t defaultZerDriverHandle = nullptr; + // Wrapper objects whose driver has been unloaded. They are kept alive (the application may + // still hold their handles) but point at loader::deadDditable, so every entry point returns + // UNINITIALIZED rather than faulting. Child objects are never resurrected: a reloaded + // driver may be a different build, and its raw handles must never be confused with these. + std::vector> retiredHandleObjects; + // Driver wrappers held across an unload, awaiting rebinding by zelReloadDriver. + std::vector> parkedZeDriverObjects; + std::vector> parkedZesDriverObjects; }; extern ze_handle_t* loaderDispatch; extern zer_dditable_t* defaultZerDdiTable; extern context_t *context; + /////////////////////////////////////////////////////////////////////////////// + /// @brief A permanently zeroed dispatch table. + /// @details Every generated intercept function reads its pfn out of object_t::dditable and + /// returns UNINITIALIZED when that pfn is null, so pointing a retired wrapper here makes all + /// of its entry points fail safely with no generated-code changes. Zero-initialized with + /// static storage duration, so it is valid before and after loader construction/teardown. + extern dditable_t deadDditable; } diff --git a/source/loader/zes_ldrddi.cpp b/source/loader/zes_ldrddi.cpp index 99b19adf..6aa97767 100644 --- a/source/loader/zes_ldrddi.cpp +++ b/source/loader/zes_ldrddi.cpp @@ -9,6 +9,8 @@ */ #include "ze_loader_internal.h" +#include + using namespace loader_driver_ddi; namespace loader @@ -142,6 +144,9 @@ namespace loader bool atLeastOneDriverValid = false; for( auto& drv : *loader::context->sysmanInstanceDrivers ) { + if (drv.slotState == driver_slot_state_t::Unloaded) { + continue; // Skipped, not blacklisted -- only zelReloadDriver brings the slot back. + } if(drv.initStatus != ZE_RESULT_SUCCESS || drv.initSysManStatus != ZE_RESULT_SUCCESS) continue; if (!drv.handle || !drv.ddiInitialized) { @@ -197,6 +202,9 @@ namespace loader for( auto& drv : *loader::context->sysmanInstanceDrivers ) { + if (drv.slotState == driver_slot_state_t::Unloaded) { + continue; // An unloaded slot is a hole in the list; it is never enumerated. + } if(drv.initStatus != ZE_RESULT_SUCCESS || drv.initSysManStatus != ZE_RESULT_SUCCESS || !drv.ddiInitialized) continue; @@ -228,8 +236,16 @@ namespace loader { for( uint32_t i = 0; i < library_driver_handle_count; ++i ) { uint32_t driver_index = total_driver_handle_count + i; - phDrivers[ driver_index ] = reinterpret_cast( - context->zes_driver_factory.getInstance( phDrivers[ driver_index ], &drv.dditable ) ); + { + auto driverObject = context->zes_driver_factory.getInstance( phDrivers[ driver_index ], &drv.dditable ); + // See the core path: the slot needs to be able to find its own wrappers + // again to rebind them after zelReloadDriver. + if (std::find(drv.zesDriverObjects.begin(), drv.zesDriverObjects.end(), driverObject) == drv.zesDriverObjects.end()) { + drv.zesDriverObjects.push_back(driverObject); + } + drv.wrapperModePinned = true; + phDrivers[ driver_index ] = reinterpret_cast( driverObject ); + } } } catch( std::bad_alloc& ) diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index a1fa8dd6..26ef9cbc 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -225,6 +225,69 @@ else() set_property(TEST tests_multi_driver_zeandzesdriverget_sort APPEND PROPERTY ENVIRONMENT "ZE_ENABLE_LOADER_DEBUG_TRACE=1;ZE_ENABLE_ALT_DRIVERS=$,$") 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=$/ze_null_test1.dll,$/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=$,$") +endif() + +# Unload/reload lifecycle. Each of these mutates process-global loader state, so +# each one needs its own process. + +# The app's original driver handle survives unload and is rebound by zelReloadDriver. +add_test(NAME tests_unload_driver_reload_same_handle COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenUnloadedDriverWhenReloadedThenSameHandleIsValidAgain) +if (MSVC) + set_property(TEST tests_unload_driver_reload_same_handle 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=$/ze_null_test1.dll,$/ze_null_test2.dll") +else() + set_property(TEST tests_unload_driver_reload_same_handle 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=$,$") +endif() + +# Handles issued below the driver are permanently retired -- a newer UMD never +# sees a device handle minted by the load it replaced. +add_test(NAME tests_unload_driver_stale_devices COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenUnloadedDriverWhenReloadedThenStaleDeviceHandlesAreDead) +if (MSVC) + set_property(TEST tests_unload_driver_stale_devices 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=$/ze_null_test1.dll,$/ze_null_test2.dll") +else() + set_property(TEST tests_unload_driver_stale_devices 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=$,$") +endif() + +# The replacement driver advertises a different DDI-handle capability; the slot +# stays pinned to wrapper mode so the app's handle keeps working. +add_test(NAME tests_unload_driver_ddi_ext_change COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenReloadedDriverWhenNewUmdChangesDdiSupportThenHandleStillWorks) +if (MSVC) + set_property(TEST tests_unload_driver_ddi_ext_change 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=$/ze_null_test1.dll,$/ze_null_test2.dll") +else() + set_property(TEST tests_unload_driver_ddi_ext_change 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=$,$") +endif() + +# The kernel-driver-yanked-out-from-under-us case: force past the in-use gate. +add_test(NAME tests_unload_driver_force COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenDriverWithLiveObjectsWhenForceUnloadedThenSucceedsAndReloads) +if (MSVC) + set_property(TEST tests_unload_driver_force 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=$/ze_null_test1.dll,$/ze_null_test2.dll") +else() + set_property(TEST tests_unload_driver_force 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=$,$") +endif() + +# Implicit init paths never resurrect an unloaded driver, but it is not blacklisted: +# an explicit zelReloadDriver restores it. +add_test(NAME tests_unload_driver_no_resurrect COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenUnloadedDriverWhenInitDriversCalledThenNotResurrected) +if (MSVC) + set_property(TEST tests_unload_driver_no_resurrect 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=$/ze_null_test1.dll,$/ze_null_test2.dll") +else() + set_property(TEST tests_unload_driver_no_resurrect 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=$,$") +endif() + +# A reload that cannot open the library leaves the slot unloaded and retryable. +add_test(NAME tests_unload_driver_failed_reload_retry COMMAND tests --gtest_filter=*LoaderUnloadDriver.GivenFailedReloadWhenRetriedThenSucceeds) +if (MSVC) + set_property(TEST tests_unload_driver_failed_reload_retry 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=$/ze_null_test1.dll,$/ze_null_test2.dll") +else() + set_property(TEST tests_unload_driver_failed_reload_retry 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=$,$") +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") diff --git a/test/loader_api.cpp b/test/loader_api.cpp index cb0d0bcf..586ce92b 100644 --- a/test/loader_api.cpp +++ b/test/loader_api.cpp @@ -14,16 +14,20 @@ #include "zes_api.h" #include "zer_api.h" +#include #include +#include +#include +#include #if defined(_WIN32) #include - #include #include #include #define putenv_safe _putenv #else #include + #include #include #include #define _dup dup @@ -3997,4 +4001,406 @@ 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 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 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); +} + +// --------------------------------------------------------------------------- +// Shared helpers for the unload/reload tests below. Every one of these tests +// mutates process-global loader state, so each gets its own add_test entry. +// --------------------------------------------------------------------------- + +// Bring up every driver. All of these tests need at least two, so the survivor +// can be checked while the other one is unloaded. +void initAllDrivers(std::vector &drivers) { + ze_init_driver_type_desc_t desc = {ZE_STRUCTURE_TYPE_INIT_DRIVER_TYPE_DESC}; + desc.flags = UINT32_MAX; + desc.pNext = nullptr; + + uint32_t count = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeInitDrivers(&count, nullptr, &desc)); + ASSERT_GE(count, 2u) + << "This test requires at least two drivers via ZE_ENABLE_ALT_DRIVERS"; + drivers.resize(count); + ASSERT_EQ(ZE_RESULT_SUCCESS, zeInitDrivers(&count, drivers.data(), &desc)); +} + +// Drive a driver end-to-end -- context, device, module -- then tear it down. +void 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 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)); +} + +// The library paths the loader was told to use, in discovery order. +std::vector altDriverPaths() { + std::vector paths; + std::stringstream ss(getenv_string("ZE_ENABLE_ALT_DRIVERS")); + std::string entry; + while (std::getline(ss, entry, ',')) { + if (!entry.empty()) + paths.push_back(entry); + } + return paths; +} + +#if !defined(_WIN32) +// How many of the alt driver libraries are still mapped into this process. +// RTLD_NOLOAD returns a handle only for an already-resident library, so this is +// the direct test of whether unload really unmapped the module -- glibc will +// silently pin a library that exports STB_GNU_UNIQUE symbols, and if it does, +// "reload to a fresh state" is a lie. +size_t residentAltDriverCount() { + size_t resident = 0; + for (const auto &path : altDriverPaths()) { + void *handle = dlopen(path.c_str(), RTLD_LAZY | RTLD_NOLOAD); + if (handle != nullptr) { + ++resident; + dlclose(handle); + } + } + return resident; +} +#endif + +// Requirement 1: unload is a reversible state, not a verdict. The driver handle +// the app holds stays a valid pointer across unload, answers UNINITIALIZED while +// the slot is dead, and is rebound onto the fresh driver by zelReloadDriver. +TEST( + LoaderUnloadDriver, + GivenUnloadedDriverWhenReloadedThenSameHandleIsValidAgain) { + + std::vector drivers; + ASSERT_NO_FATAL_FAILURE(initAllDrivers(drivers)); + ze_driver_handle_t survivor = drivers[0]; + ze_driver_handle_t target = drivers[1]; + + ASSERT_NO_FATAL_FAILURE(executeOnDriver(target)); + +#if !defined(_WIN32) + const size_t residentBefore = residentAltDriverCount(); +#endif + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelUnloadDriver(target)); + +#if !defined(_WIN32) + EXPECT_EQ(residentAltDriverCount(), residentBefore - 1) + << "the unloaded driver library is still mapped -- reload cannot produce " + "a fresh driver state"; +#endif + + // The handle is dead but not dangling: calls through it fail cleanly. + uint32_t deadCount = 0; + EXPECT_EQ(ZE_RESULT_ERROR_UNINITIALIZED, zeDeviceGet(target, &deadCount, nullptr)); + ze_context_desc_t contextDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC}; + ze_context_handle_t deadContext = nullptr; + EXPECT_EQ(ZE_RESULT_ERROR_UNINITIALIZED, + zeContextCreate(target, &contextDesc, &deadContext)); + + // The survivor is untouched while its neighbour is unloaded. + ASSERT_NO_FATAL_FAILURE(executeOnDriver(survivor)); + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelReloadDriver(target)); + + // Same ze_driver_handle_t, fresh driver underneath it. + ASSERT_NO_FATAL_FAILURE(executeOnDriver(target)); + ASSERT_NO_FATAL_FAILURE(executeOnDriver(survivor)); +} + +// Requirement 3: a reloaded driver may be a newer UMD with different interfaces, +// so nothing below the driver handle survives. Handles issued by the old load +// are permanently dead -- they must never be handed to the new driver. +TEST( + LoaderUnloadDriver, + GivenUnloadedDriverWhenReloadedThenStaleDeviceHandlesAreDead) { + + std::vector drivers; + ASSERT_NO_FATAL_FAILURE(initAllDrivers(drivers)); + ze_driver_handle_t target = drivers[1]; + + uint32_t deviceCount = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeDeviceGet(target, &deviceCount, nullptr)); + ASSERT_GT(deviceCount, 0u); + std::vector staleDevices(deviceCount); + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeDeviceGet(target, &deviceCount, staleDevices.data())); + ze_device_handle_t staleDevice = staleDevices[0]; + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelUnloadDriver(target)); + ASSERT_EQ(ZE_RESULT_SUCCESS, zelReloadDriver(target)); + + // The stale device handle is retired for good, even though its driver is back. + ze_device_properties_t staleProps = {ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES}; + EXPECT_EQ(ZE_RESULT_ERROR_UNINITIALIZED, + zeDeviceGetProperties(staleDevice, &staleProps)); + uint32_t staleSubCount = 0; + EXPECT_EQ(ZE_RESULT_ERROR_UNINITIALIZED, + zeDeviceGetSubDevices(staleDevice, &staleSubCount, nullptr)); + + // Re-enumerating through the rebound driver handle yields working devices. + uint32_t freshCount = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeDeviceGet(target, &freshCount, nullptr)); + ASSERT_GT(freshCount, 0u); + std::vector freshDevices(freshCount); + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeDeviceGet(target, &freshCount, freshDevices.data())); + EXPECT_NE(freshDevices[0], staleDevice); + + ze_device_properties_t freshProps = {ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES}; + EXPECT_EQ(ZE_RESULT_SUCCESS, zeDeviceGetProperties(freshDevices[0], &freshProps)); +} + +// A reloaded UMD may advertise a different feature set than the one it replaced. +// The null driver re-reads ZEL_TEST_NULL_DRIVER_DISABLE_DDI_EXT in its global +// context_t constructor, which re-runs on the fresh load, so flipping the +// variable between unload and reload genuinely simulates a different-generation +// driver. The slot stays in wrapper mode regardless, which is what keeps the +// app's original handle valid. +TEST( + LoaderUnloadDriver, + GivenReloadedDriverWhenNewUmdChangesDdiSupportThenHandleStillWorks) { + + // The ctest entry starts this process with DISABLE_DDI_EXT=3 (both null + // drivers hide the extension), so the loader wraps every handle it issues. + ASSERT_EQ("3", getenv_string("ZEL_TEST_NULL_DRIVER_DISABLE_DDI_EXT")); + + std::vector drivers; + ASSERT_NO_FATAL_FAILURE(initAllDrivers(drivers)); + ze_driver_handle_t target = drivers[1]; + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelUnloadDriver(target)); + + // The replacement UMD now advertises ZE_DRIVER_DDI_HANDLE_EXT. + putenv_safe(const_cast("ZEL_TEST_NULL_DRIVER_DISABLE_DDI_EXT=")); + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelReloadDriver(target)); + + // The handle was pinned to wrapper mode at first load, so it still resolves. + ASSERT_NO_FATAL_FAILURE(executeOnDriver(target)); + ASSERT_NO_FATAL_FAILURE(executeOnDriver(drivers[0])); +} + +// The Microsoft scenario: the kernel-mode driver is yanked out from under a live +// user-mode driver, so the app is left holding objects it can never cleanly +// destroy. The default in-use gate blocks recovery entirely; FORCE overrides it. +TEST( + LoaderUnloadDriver, + GivenDriverWithLiveObjectsWhenForceUnloadedThenSucceedsAndReloads) { + + std::vector drivers; + ASSERT_NO_FATAL_FAILURE(initAllDrivers(drivers)); + ze_driver_handle_t target = drivers[1]; + + ze_context_desc_t contextDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC}; + ze_context_handle_t strandedContext = nullptr; + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeContextCreate(target, &contextDesc, &strandedContext)); + + // Safe by default: a driver with live children refuses to unload. + EXPECT_EQ(ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE, zelUnloadDriver(target)); + EXPECT_EQ(ZE_RESULT_ERROR_HANDLE_OBJECT_IN_USE, + zelUnloadDriverExt(target, ZEL_UNLOAD_DRIVER_FLAG_NONE)); + + // FORCE accepts the driver-side leak and unloads anyway. + ASSERT_EQ(ZE_RESULT_SUCCESS, + zelUnloadDriverExt(target, ZEL_UNLOAD_DRIVER_FLAG_FORCE)); + + // The stranded context fails cleanly instead of faulting into a freed module. + EXPECT_EQ(ZE_RESULT_ERROR_UNINITIALIZED, zeContextDestroy(strandedContext)); + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelReloadDriver(target)); + ASSERT_NO_FATAL_FAILURE(executeOnDriver(target)); + ASSERT_NO_FATAL_FAILURE(executeOnDriver(drivers[0])); +} + +// Requirement 1, stated negatively and then positively: the implicit paths never +// resurrect an unloaded driver (an unrelated component calling zeInitDrivers must +// not undo a deliberate unload), but the slot is not blacklisted either -- an +// explicit zelReloadDriver brings it back and the counts return to normal. +TEST( + LoaderUnloadDriver, + GivenUnloadedDriverWhenInitDriversCalledThenNotResurrected) { + + ze_init_driver_type_desc_t desc = {ZE_STRUCTURE_TYPE_INIT_DRIVER_TYPE_DESC}; + desc.flags = UINT32_MAX; + desc.pNext = nullptr; + + std::vector drivers; + ASSERT_NO_FATAL_FAILURE(initAllDrivers(drivers)); + const uint32_t originalCount = static_cast(drivers.size()); + ze_driver_handle_t target = drivers[1]; + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelUnloadDriver(target)); + + // Neither zeInitDrivers nor zeDriverGet may bring it back. + uint32_t reinitCount = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeInitDrivers(&reinitCount, nullptr, &desc)); + EXPECT_EQ(reinitCount, originalCount - 1) + << "zeInitDrivers should not resurrect the unloaded driver"; + + uint32_t getCount = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeDriverGet(&getCount, nullptr)); + EXPECT_EQ(getCount, originalCount - 1) + << "zeDriverGet should not resurrect the unloaded driver"; + + // The drivers still standing keep working. + std::vector remaining(reinitCount); + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeInitDrivers(&reinitCount, remaining.data(), &desc)); + for (auto driver : remaining) { + ze_context_desc_t contextDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC}; + ze_context_handle_t context = nullptr; + EXPECT_EQ(ZE_RESULT_SUCCESS, zeContextCreate(driver, &contextDesc, &context)); + EXPECT_EQ(ZE_RESULT_SUCCESS, zeContextDestroy(context)); + } + + // Not blacklisted: the explicit reload restores the slot and the counts. + ASSERT_EQ(ZE_RESULT_SUCCESS, zelReloadDriver(target)); + + uint32_t afterReloadInit = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeInitDrivers(&afterReloadInit, nullptr, &desc)); + EXPECT_EQ(afterReloadInit, originalCount); + + uint32_t afterReloadGet = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeDriverGet(&afterReloadGet, nullptr)); + EXPECT_EQ(afterReloadGet, originalCount); + + ASSERT_NO_FATAL_FAILURE(executeOnDriver(target)); +} + +// A reload that cannot open the library must leave the slot unloaded and +// retryable rather than tombstoned -- exactly the case where the replacement +// driver package is mid-install. +TEST( + LoaderUnloadDriver, + GivenFailedReloadWhenRetriedThenSucceeds) { + +#if defined(_WIN32) + GTEST_SKIP() << "Relies on renaming the driver library out from under the loader"; +#else + auto paths = altDriverPaths(); + ASSERT_GE(paths.size(), 2u) + << "This test requires two drivers via ZE_ENABLE_ALT_DRIVERS"; + + // Work against a private copy of the second driver so the test can move the + // file without disturbing the build tree. + const std::string copyPath = "/tmp/zel_reload_test_driver.so.1"; + const std::string hiddenPath = copyPath + ".hidden"; + std::remove(hiddenPath.c_str()); + { + std::ifstream src(paths[1], std::ios::binary); + ASSERT_TRUE(src.good()) << "cannot read " << paths[1]; + std::ofstream dst(copyPath, std::ios::binary | std::ios::trunc); + ASSERT_TRUE(dst.good()) << "cannot write " << copyPath; + dst << src.rdbuf(); + } + + // putenv() keeps the caller's buffer, so this string must outlive the test. + static std::string altDriverEnv = + "ZE_ENABLE_ALT_DRIVERS=" + paths[0] + "," + copyPath; + ASSERT_EQ(0, putenv_safe(const_cast(altDriverEnv.c_str()))); + + std::vector drivers; + ASSERT_NO_FATAL_FAILURE(initAllDrivers(drivers)); + ze_driver_handle_t target = drivers[1]; + + ASSERT_EQ(ZE_RESULT_SUCCESS, zelUnloadDriver(target)); + + // The library disappears before the reload lands. + ASSERT_EQ(0, std::rename(copyPath.c_str(), hiddenPath.c_str())); + EXPECT_EQ(ZE_RESULT_ERROR_UNINITIALIZED, zelReloadDriver(target)); + + // Still unloaded, still skipped by the implicit paths, still not blacklisted. + uint32_t getCount = 0; + ASSERT_EQ(ZE_RESULT_SUCCESS, zeDriverGet(&getCount, nullptr)); + EXPECT_EQ(getCount, drivers.size() - 1); + + // The package finishes installing; the retry succeeds. + ASSERT_EQ(0, std::rename(hiddenPath.c_str(), copyPath.c_str())); + ASSERT_EQ(ZE_RESULT_SUCCESS, zelReloadDriver(target)); + + ASSERT_NO_FATAL_FAILURE(executeOnDriver(target)); + ASSERT_NO_FATAL_FAILURE(executeOnDriver(drivers[0])); + + std::remove(copyPath.c_str()); +#endif +} + } // namespace