[L0] Implement Support for zeInitDrivers

- As of v1.10 of the L0 spec, zeInit and zeDriverGet is the old init
  pathway and the desired init api is zeInitDrivers. This new api allows
for multi heterogenous drivers to coexist in a single L0 Process.

Signed-off-by: Neil R. Spruit <[email protected]>

Author: nrspruit

source/adapters/level_zero/adapter.cpp

@@ -78,18 +78,25 @@ ur_result_t getZesDeviceHandle(zes_uuid_t coreDeviceUuid,
 
 ur_result_t initPlatforms(PlatformVec &platforms,
                           ze_result_t ZesResult) noexcept try {
-  uint32_t ZeDriverCount = 0;
-  ZE2UR_CALL(zeDriverGet, (&ZeDriverCount, nullptr));
-  if (ZeDriverCount == 0) {
+  if (!GlobalAdapter->initDriversFunctionPtr) {
+    ZE2UR_CALL(zeDriverGet, (&GlobalAdapter->ZeDriverCount, nullptr));
+  }
+  if (GlobalAdapter->ZeDriverCount == 0) {
     return UR_RESULT_SUCCESS;
   }
 
   std::vector<ze_driver_handle_t> ZeDrivers;
   std::vector<ze_device_handle_t> ZeDevices;
-  ZeDrivers.resize(ZeDriverCount);
-
-  ZE2UR_CALL(zeDriverGet, (&ZeDriverCount, ZeDrivers.data()));
-  for (uint32_t I = 0; I < ZeDriverCount; ++I) {
+  ZeDrivers.resize(GlobalAdapter->ZeDriverCount);
+
+  if (GlobalAdapter->initDriversFunctionPtr) {
+    ZE2UR_CALL(GlobalAdapter->initDriversFunctionPtr,
+               (&GlobalAdapter->ZeDriverCount, ZeDrivers.data(),
+                &GlobalAdapter->InitDriversDesc));
+  } else {
+    ZE2UR_CALL(zeDriverGet, (&GlobalAdapter->ZeDriverCount, ZeDrivers.data()));
+  }
+  for (uint32_t I = 0; I < GlobalAdapter->ZeDriverCount; ++I) {
     // Keep track of the first platform init for this Driver
     bool DriverPlatformInit = false;
     ze_device_properties_t device_properties{};
@@ -214,6 +221,15 @@ ur_adapter_handle_t_::ur_adapter_handle_t_()
       return std::atoi(UrRet);
     }();
 
+    // Dynamically load the new L0 apis separately.
+    // This must be done to avoid attempting to use symbols that do
+    // not exist in older loader runtimes.
+#ifdef _WIN32
+    HMODULE processHandle = GetModuleHandle(NULL);
+#else
+    HMODULE processHandle = nullptr;
+#endif
+
     // initialize level zero only once.
     if (GlobalAdapter->ZeResult == std::nullopt) {
       // Setting these environment variables before running zeInit will enable
@@ -235,20 +251,37 @@ ur_adapter_handle_t_::ur_adapter_handle_t_()
       // called multiple times.  Declaring the return value as "static" ensures
       // it's only called once.
 
-      // Init with all flags set to enable for all driver types to be init in
-      // the application.
-      ze_init_flags_t L0InitFlags = ZE_INIT_FLAG_GPU_ONLY;
-      if (UrL0InitAllDrivers) {
-        L0InitFlags |= ZE_INIT_FLAG_VPU_ONLY;
-      }
+      GlobalAdapter->initDriversFunctionPtr =
+          (ze_pfnInitDrivers_t)ur_loader::LibLoader::getFunctionPtr(
+              processHandle, "zeInitDrivers");
 
       // Set ZES_ENABLE_SYSMAN by default if the user has not set it.
       if (UrSysManEnvInitEnabled) {
         setEnvVar("ZES_ENABLE_SYSMAN", "1");
       }
-      logger::debug("\nzeInit with flags value of {}\n",
-                    static_cast<int>(L0InitFlags));
-      GlobalAdapter->ZeResult = ZE_CALL_NOCHECK(zeInit, (L0InitFlags));
+
+      if (GlobalAdapter->initDriversFunctionPtr) {
+        GlobalAdapter->InitDriversDesc.pNext = nullptr;
+        GlobalAdapter->InitDriversDesc.flags = ZE_INIT_DRIVER_TYPE_FLAG_GPU;
+        logger::debug("\nzeInitDrivers with flags value of {}\n",
+                      static_cast<int>(GlobalAdapter->InitDriversDesc.flags));
+        GlobalAdapter->ZeResult =
+            ZE_CALL_NOCHECK(GlobalAdapter->initDriversFunctionPtr,
+                            (&GlobalAdapter->ZeDriverCount, nullptr,
+                             &GlobalAdapter->InitDriversDesc));
+      }
+
+      if (!GlobalAdapter->ZeResult) {
+        // Init with all flags set to enable for all driver types to be init in
+        // the application.
+        ze_init_flags_t L0InitFlags = ZE_INIT_FLAG_GPU_ONLY;
+        if (UrL0InitAllDrivers) {
+          L0InitFlags |= ZE_INIT_FLAG_VPU_ONLY;
+        }
+        logger::debug("\nzeInit with flags value of {}\n",
+                      static_cast<int>(L0InitFlags));
+        GlobalAdapter->ZeResult = ZE_CALL_NOCHECK(zeInit, (L0InitFlags));
+      }
     }
     assert(GlobalAdapter->ZeResult !=
            std::nullopt); // verify that level-zero is initialized
@@ -265,14 +298,6 @@ ur_adapter_handle_t_::ur_adapter_handle_t_()
 
       return;
     }
-    // Dynamically load the new L0 SysMan separate init and new EXP apis
-    // separately. This must be done to avoid attempting to use symbols that do
-    // not exist in older loader runtimes.
-#ifdef _WIN32
-    HMODULE processHandle = GetModuleHandle(NULL);
-#else
-    HMODULE processHandle = nullptr;
-#endif
 
     // Check if the user has enabled the default L0 SysMan initialization.
     const int UrSysmanZesinitEnable = [] {

source/adapters/level_zero/adapter.hpp

@@ -17,6 +17,7 @@
 #include <optional>
 #include <ur/ur.hpp>
 #include <ze_api.h>
+#include <ze_ddi.h>
 #include <zes_ddi.h>
 
 using PlatformVec = std::vector<std::unique_ptr<ur_platform_handle_t_>>;
@@ -31,6 +32,10 @@ struct ur_adapter_handle_t_ {
   zes_pfnDriverGetDeviceByUuidExp_t getDeviceByUUIdFunctionPtr = nullptr;
   zes_pfnDriverGet_t getSysManDriversFunctionPtr = nullptr;
   zes_pfnInit_t sysManInitFunctionPtr = nullptr;
+  ze_pfnInitDrivers_t initDriversFunctionPtr = nullptr;
+  ze_init_driver_type_desc_t InitDriversDesc = {
+      ZE_STRUCTURE_TYPE_INIT_DRIVER_TYPE_DESC, nullptr, 0};
+  uint32_t ZeDriverCount = 0;
 
   std::optional<ze_result_t> ZeResult;
   std::optional<ze_result_t> ZesResult;