Add Dynamic Tracing Support for Zer API's (#387)

* Add Dynamic Tracing Support for Zer API's

Signed-off-by: Vishnu Khanth <[email protected]>

Author: vishnu-khanth

source/lib/ze_lib.cpp

@@ -177,14 +177,21 @@ namespace ze_lib
             return ZE_RESULT_ERROR_UNINITIALIZED;
         }
         tracing_lib = getTracing();
-        typedef ze_result_t (ZE_APICALL *zelLoaderTracingLayerInit_t)(std::atomic<ze_dditable_t *> &zeDdiTable);
-        auto loaderTracingLayerInit = reinterpret_cast<zelLoaderTracingLayerInit_t>(
+        typedef ze_result_t (ZE_APICALL *zelLoaderZeTracingLayerInit_t)(std::atomic<ze_dditable_t *> &zeDdiTable);
+        auto loaderZeTracingLayerInit = reinterpret_cast<zelLoaderZeTracingLayerInit_t>(
                 GET_FUNCTION_PTR(loader, "zelLoaderTracingLayerInit") );
-        if (loaderTracingLayerInit == nullptr) {
+        if (loaderZeTracingLayerInit == nullptr) {
             std::string message = "ze_lib Context Init() zelLoaderTracingLayerInit missing, disabling dynamic tracer support ";
             debug_trace_message(message, "");
             this->dynamicTracingSupported = false;
         }
+        typedef ze_result_t (ZE_APICALL *zelLoaderZerTracingLayerInit_t)(std::atomic<zer_dditable_t *> &zerDdiTable);
+        auto loaderZerTracingLayerInit = reinterpret_cast<zelLoaderZerTracingLayerInit_t>(
+                GET_FUNCTION_PTR(loader, "zelLoaderZerTracingLayerInit") );
+        if (loaderZerTracingLayerInit == nullptr) {
+            std::string message = "ze_lib Context Init() zelLoaderZerTracingLayerInit missing, dynamic tracing support for Zer API's will be unavailable ";
+            debug_trace_message(message, "");
+        }
         typedef loader::context_t * (ZE_APICALL *zelLoaderGetContext_t)();
         auto loaderGetContext = reinterpret_cast<zelLoaderGetContext_t>(
                 GET_FUNCTION_PTR(loader, "zelLoaderGetContext") );
@@ -280,11 +287,17 @@ namespace ze_lib
         if( ZE_RESULT_SUCCESS == result )
         {
             #ifdef L0_STATIC_LOADER_BUILD
-            if (loaderTracingLayerInit) {
-                result = loaderTracingLayerInit(this->pTracingZeDdiTable);
+            if (loaderZeTracingLayerInit) {
+                result = loaderZeTracingLayerInit(this->pTracingZeDdiTable);
+                if (result == ZE_RESULT_SUCCESS && loaderZerTracingLayerInit) {
+                    result = loaderZerTracingLayerInit(this->pTracingZerDdiTable);
+                }
             }
             #else
             result = zelLoaderTracingLayerInit(this->pTracingZeDdiTable);
+            if (result == ZE_RESULT_SUCCESS) {
+                result = zelLoaderZerTracingLayerInit(this->pTracingZerDdiTable);
+            }
             #endif
         }
         // End DDI Table Inits
@@ -638,6 +651,9 @@ zelEnableTracingLayer()
     }
     if (ze_lib::context->tracingLayerEnableCounter.fetch_add(1) == 0) {
         ze_lib::context->zeDdiTable.exchange(ze_lib::context->pTracingZeDdiTable);
+        if (ze_lib::context->pTracingZerDdiTable != nullptr) {
+            ze_lib::context->zerDdiTable.exchange(ze_lib::context->pTracingZerDdiTable);
+        }
     }
     #endif
     return ZE_RESULT_SUCCESS;
@@ -659,6 +675,9 @@ zelDisableTracingLayer()
     }
     if (ze_lib::context->tracingLayerEnableCounter.fetch_sub(1) <= 1) {
         ze_lib::context->zeDdiTable.exchange(&ze_lib::context->initialzeDdiTable);
+        if (ze_lib::context->pTracingZerDdiTable != nullptr) {
+            ze_lib::context->zerDdiTable.exchange(&ze_lib::context->initialzerDdiTable);
+        }
     }
     #endif
     return ZE_RESULT_SUCCESS;

source/lib/ze_lib.h

@@ -169,6 +169,7 @@ namespace ze_lib
         ze_result_t zelTracingDdiTableInit(ze_api_version_t version);
         zel_tracing_dditable_t  zelTracingDdiTable = {};
         std::atomic<ze_dditable_t *> pTracingZeDdiTable = {nullptr};
+        std::atomic<zer_dditable_t *> pTracingZerDdiTable = {nullptr};
         ze_dditable_t initialzeDdiTable;
         zet_dditable_t initialzetDdiTable;
         zes_dditable_t initialzesDdiTable;

source/loader/ze_loader_api.cpp

@@ -60,14 +60,23 @@ zelLoaderGetContext() {
 }
 
 ///////////////////////////////////////////////////////////////////////////////
-/// @brief Internal function for Setting the ddi table for the Tracing Layer.
+/// @brief Internal function for Setting the ZE ddi table for the Tracing Layer.
 ///
 ZE_DLLEXPORT ze_result_t ZE_APICALL
 zelLoaderTracingLayerInit(std::atomic<ze_dditable_t *> &zeDdiTable) {
     zeDdiTable.store(&loader::context->tracing_dditable.ze);
     return ZE_RESULT_SUCCESS;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Internal function for Setting the ZER ddi table for the Tracing Layer.
+///
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelLoaderZerTracingLayerInit(std::atomic<zer_dditable_t *> &zerDdiTable) {
+    zerDdiTable.store(&loader::context->tracing_dditable.zer);
+    return ZE_RESULT_SUCCESS;
+}
+
 ZE_DLLEXPORT ze_result_t ZE_APICALL
 zelLoaderGetVersionsInternal(
    size_t *num_elems,                     //Pointer to num versions to get.  

source/loader/ze_loader_api.h

@@ -1,6 +1,6 @@
 /*
  *
- * Copyright (C) 2020-2021 Intel Corporation
+ * Copyright (C) 2020-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -37,11 +37,17 @@ zelLoaderDriverCheck(ze_init_flags_t flags, ze_init_driver_type_desc_t* desc, ze
 
 
 ///////////////////////////////////////////////////////////////////////////////
-/// @brief Internal function for Setting the ddi table for the Tracing Layer.
+/// @brief Internal function for Setting the ZE ddi table for the Tracing Layer.
 ///
 ZE_DLLEXPORT ze_result_t ZE_APICALL
 zelLoaderTracingLayerInit(std::atomic<ze_dditable_t *> &zeDdiTable);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Internal function for Setting the ZER ddi table for the Tracing Layer.
+///
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelLoaderZerTracingLayerInit(std::atomic<zer_dditable_t *> &zerDdiTable);
+
 
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for getting tracing lib handle