From 2603dbb8b6133a1357e1861349ff1df7bd6a714c Mon Sep 17 00:00:00 2001
From: "Neil R. Spruit" <neil.r.spruit@intel.com>
Date: Fri, 28 Jun 2024 16:56:50 -0700
Subject: [PATCH] Fix layer init with multiple drivers and ensure init for
 instrumentation support (#164)

-Fix potential race condition when initializing the layers
 with multiple drivers by verifying the usable drivers and layers
 before init of the final ddi tables.
-Document the reasons for requiring ddi table init before driver init
 checks can be done.

Signed-off-by: Neil R. Spruit <neil.r.spruit@intel.com>
---
 source/lib/ze_lib.cpp | 34 +++++++++++++++++++++-------------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/source/lib/ze_lib.cpp b/source/lib/ze_lib.cpp
index cfa21f5b..4e26d29a 100644
--- a/source/lib/ze_lib.cpp
+++ b/source/lib/ze_lib.cpp
@@ -74,40 +74,48 @@ namespace ze_lib
             ze_lib::context->zesDdiTable.exchange(&ze_lib::context->initialzesDdiTable);
         }
 
+        // Always call the inits for all the ddi tables before checking which drivers are usable to enable Instrumentation correctly.
 
-        // Check which drivers and layers are functional by checking which can be loaded.
-        // This check needs to be done before all inits to avoid overwriting the layer ddi table pointers.
-        if( ZE_RESULT_SUCCESS == result && !sysmanOnly)
-        {
-            //Check which drivers support the ze_driver_flag_t specified
-            //No need to check if only initializing sysman
-            result = zelLoaderDriverCheck(flags);
-        }
-
+        // Init the ZE DDI Tables
         if( ZE_RESULT_SUCCESS == result )
         {
             result = zeInit();
         }
-
+        // Init the ZET DDI Tables
         if( ZE_RESULT_SUCCESS == result )
         {
             result = zetInit();
         }
-
+        // Init the ZES DDI Tables
         if( ZE_RESULT_SUCCESS == result )
         {
             result = zesInit();
         }
-
+        // Init the Tracing API DDI Tables
         if( ZE_RESULT_SUCCESS == result )
         {
             result = zelTracingInit();
         }
-
+        // Init the stored ddi tables for the tracing layer
         if( ZE_RESULT_SUCCESS == result )
         {
             result = zelLoaderTracingLayerInit(this->pTracingZeDdiTable, this->pTracingZetDdiTable, this->pTracingZesDdiTable);
         }
+        // End DDI Table Inits
+
+        // Check which drivers and layers can be init on this system.
+        if( ZE_RESULT_SUCCESS == result && !sysmanOnly)
+        {
+            // Check which drivers support the ze_driver_flag_t specified
+            // No need to check if only initializing sysman
+            result = zelLoaderDriverCheck(flags);
+            // reInit the ze ddi tables after verifying the zeInit() with dummy tables.
+            // This ensures the tracing and validation layers are pointing to the correct function pointers after init.
+            if( ZE_RESULT_SUCCESS == result )
+            {
+                result = zeInit();
+            }
+        }
 
         if( ZE_RESULT_SUCCESS == result )
         {