Initial Feature commit for Extension Callback Tracing

include/loader/ze_loader.h

@@ -565,6 +565,78 @@ zelDisableTracingLayer(void);
 ZE_DLLEXPORT ze_result_t ZE_APICALL
 zelGetTracingLayerState(bool* enabled); // Pointer to bool to receive tracing layer state
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Callback signature for extension-function prologue/epilogue handlers.
+///
+/// This intentionally mirrors the established per-API tracing callback shape
+/// (see the ze_pfnXCb_t typedefs in ze_api.h) so tools can reuse their existing
+/// callback infrastructure. Because an arbitrary extension function has no
+/// generated params struct, @p pParams is passed as an opaque void* whose layout
+/// is defined by the driver for the named function (may be null for pure-vendor
+/// functions). The identity of the fired function is carried via
+/// @p pTracerUserData (set at registration time).
+///
+/// @param[in] pParams                    driver-populated parameter block (opaque)
+/// @param[in] result                     epilogue only: the function's return value
+/// @param[in] pTracerUserData            per-registration user data
+/// @param[in,out] ppTracerInstanceUserData  per-call scratch for prologue->epilogue handoff
+typedef void (ZE_APICALL *zel_pfnDriverExtensionFunctionCb_t)(
+    void* pParams,
+    ze_result_t result,
+    void* pTracerUserData,
+    void** ppTracerInstanceUserData
+    );
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Signature of the per-driver hook that enables or disables the driver's
+///        extension-function callbacks.
+///
+/// A driver that supports extension-function tracing exposes this by name
+/// ("zelDriverEnableTracing") via zeDriverGetExtensionFunctionAddress. The loader
+/// calls it on each active driver when the tracing layer is enabled/disabled
+/// (including static ZE_ENABLE_TRACING_LAYER enablement and late-loaded drivers).
+/// When disabled, the driver must not invoke any registered prologue/epilogue.
+typedef ze_result_t (ZE_APICALL *zel_pfnDriverEnableTracing_t)(
+    ze_driver_handle_t hDriver,
+    ze_bool_t enable
+    );
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Registers prologue/epilogue callbacks for a named extension function.
+///
+/// Extension functions obtained by string name via
+/// zeDriverGetExtensionFunctionAddress() return a raw driver pointer that the
+/// application calls directly, bypassing the loader and therefore the tracing
+/// layer. This API provides a driver-side interception hook: the specified
+/// driver invokes the registered @p prologue before, and @p epilogue after, the
+/// body of the extension function named @p functionName.
+///
+/// Registration is keyed by @p functionName and is order-independent relative to
+/// zeDriverGetExtensionFunctionAddress() — it takes effect on the next invocation
+/// of the function, even if the application already cached the function pointer.
+/// Passing null for both @p prologue and @p epilogue unregisters the callbacks.
+///
+/// @param[in] hDriver        handle of the driver instance
+/// @param[in] functionName   name of the extension function to intercept
+/// @param[in] pUserData      user data passed to the callbacks as pTracerUserData
+/// @param[in] prologue       handler invoked before the function body (may be null)
+/// @param[in] epilogue       handler invoked after the function body (may be null)
+///
+/// @return
+///   - ZE_RESULT_SUCCESS on success (including unregister).
+///   - ZE_RESULT_ERROR_UNINITIALIZED if the loader is not initialized.
+///   - ZE_RESULT_ERROR_UNSUPPORTED_FEATURE if the driver does not implement this hook.
+///   - ZE_RESULT_ERROR_INVALID_NULL_HANDLE if @p hDriver is null.
+///   - ZE_RESULT_ERROR_INVALID_NULL_POINTER if @p functionName is null.
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelDriverSetExtensionFunctionCallback(
+    ze_driver_handle_t hDriver,                      // [in] handle of the driver instance
+    const char* functionName,                        // [in] extension function name to intercept
+    void* pUserData,                                 // [in][optional] user data passed to callbacks
+    zel_pfnDriverExtensionFunctionCb_t prologue,     // [in][optional] prologue handler
+    zel_pfnDriverExtensionFunctionCb_t epilogue      // [in][optional] epilogue handler
+    );
+
 #if defined(__cplusplus)
 } // extern "C"
 #endif

source/drivers/null/ze_null.cpp

@@ -47,6 +47,33 @@ namespace driver
             return ZE_RESULT_SUCCESS;
         };
 
+        //////////////////////////////////////////////////////////////////////////
+        // Custom extension-function resolver. Returns real driver pointers by name
+        // for the setter and the sample extension function (the generic intercept
+        // in ze_nullddi.cpp defers to this hook and forwards *ppFunctionAddress).
+        zeDdiTable.Driver.pfnGetExtensionFunctionAddress = [](
+            ze_driver_handle_t,
+            const char* name,
+            void** ppFunctionAddress )
+        {
+            if( nullptr == name || nullptr == ppFunctionAddress )
+                return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
+            if( 0 == strcmp( name, "zelDriverSetExtensionFunctionCallback" ) ) {
+                *ppFunctionAddress = reinterpret_cast<void*>( &driver::zelDriverSetExtensionFunctionCallback );
+                return ZE_RESULT_SUCCESS;
+            }
+            if( 0 == strcmp( name, "zelDriverEnableTracing" ) ) {
+                *ppFunctionAddress = reinterpret_cast<void*>( &driver::zelDriverEnableTracing );
+                return ZE_RESULT_SUCCESS;
+            }
+            if( 0 == strcmp( name, "zeSampleExtFunc" ) ) {
+                *ppFunctionAddress = reinterpret_cast<void*>( &driver::zeSampleExtFunc );
+                return ZE_RESULT_SUCCESS;
+            }
+            *ppFunctionAddress = nullptr;
+            return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+        };
+
         //////////////////////////////////////////////////////////////////////////
         zeDdiTable.Device.pfnGet = [](
             ze_driver_handle_t,
@@ -680,6 +707,80 @@ namespace driver
         pRuntime.version = ZE_API_VERSION_CURRENT;
     }
 
+    ///////////////////////////////////////////////////////////////////////////
+    /// @brief Sample extension function reachable only by name. Its body invokes
+    ///        any registered prologue/epilogue with a typed params block.
+    ze_result_t ZE_APICALL zeSampleExtFunc(
+        ze_driver_handle_t hDriver, uint32_t input, uint32_t* pOutput )
+    {
+        // Snapshot any registered callbacks for this function (name-keyed).
+        context_t::extension_function_callbacks_t cbs;
+        bool haveCbs = false;
+        {
+            std::lock_guard<std::mutex> lock( context.extensionCallbackMutex );
+            auto it = context.extensionCallbacks.find( "zeSampleExtFunc" );
+            if( it != context.extensionCallbacks.end() ) {
+                cbs = it->second;
+                haveCbs = true;
+            }
+        }
+
+        // Two-level gate: callbacks fire only when tracing is globally enabled
+        // AND a callback is registered for this function.
+        const bool fire = haveCbs && context.extensionCallbacksEnabled.load();
+
+        // Typed parameter block the driver exposes to the callbacks.
+        ze_sample_ext_func_params_t params = { &hDriver, &input, &pOutput };
+        void* pInstanceData = nullptr;
+        ze_result_t result = ZE_RESULT_SUCCESS;
+
+        if( fire && nullptr != cbs.prologue )
+            cbs.prologue( &params, result, cbs.pUserData, &pInstanceData );
+
+        // The (trivial) work of the extension function.
+        if( nullptr != pOutput )
+            *pOutput = input * 2;
+
+        if( fire && nullptr != cbs.epilogue )
+            cbs.epilogue( &params, result, cbs.pUserData, &pInstanceData );
+
+        return result;
+    }
+
+    ///////////////////////////////////////////////////////////////////////////
+    /// @brief Enable/disable this driver's extension-function callbacks (the
+    ///        global gate). Called by the loader when the tracing layer is
+    ///        enabled/disabled.
+    ze_result_t ZE_APICALL zelDriverEnableTracing(
+        ze_driver_handle_t /*hDriver*/, ze_bool_t enable )
+    {
+        context.extensionCallbacksEnabled.store( enable != 0 );
+        return ZE_RESULT_SUCCESS;
+    }
+
+    ///////////////////////////////////////////////////////////////////////////
+    /// @brief Driver-side registration entry (resolved by name from the loader).
+    ///        Permissive and name-keyed: any name registers; null+null unregisters.
+    ze_result_t ZE_APICALL zelDriverSetExtensionFunctionCallback(
+        ze_driver_handle_t, const char* functionName, void* pUserData,
+        zel_pfnDriverExtensionFunctionCb_t prologue,
+        zel_pfnDriverExtensionFunctionCb_t epilogue )
+    {
+        if( nullptr == functionName )
+            return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
+
+        std::lock_guard<std::mutex> lock( context.extensionCallbackMutex );
+        if( nullptr == prologue && nullptr == epilogue ) {
+            context.extensionCallbacks.erase( functionName );
+        } else {
+            auto& entry = context.extensionCallbacks[ functionName ];
+            entry.pUserData = pUserData;
+            entry.prologue = prologue;
+            entry.epilogue = epilogue;
+        }
+        return ZE_RESULT_SUCCESS;
+    }
+
     char *context_t::setenv_var_with_driver_id(const std::string &key, uint32_t driverId)
     {
         std::string env = key + "=" + std::to_string(driverId);

source/drivers/null/ze_null.h

@@ -10,11 +10,16 @@
 #pragma once
 #include <stdlib.h>
 #include <vector>
+#include <map>
+#include <mutex>
+#include <string>
+#include <atomic>
 #include "ze_ddi.h"
 #include "zet_ddi.h"
 #include "zes_ddi.h"
 #include "ze_util.h"
 #include "ze_ddi_common.h"
+#include "loader/ze_loader.h"
 
 #ifndef ZEL_NULL_DRIVER_ID
 #define ZEL_NULL_DRIVER_ID 1
@@ -47,6 +52,23 @@ namespace driver
         std::vector<BaseNullHandle*> globalBaseNullHandle;
 	bool ddiExtensionSupported = false;
 	std::vector<char *> env_vars{};
+
+        // Registry for zelDriverSetExtensionFunctionCallback: maps an extension
+        // function name to the prologue/epilogue the driver invokes from inside
+        // that function's body. Keyed by name (order-independent vs fetch).
+        struct extension_function_callbacks_t {
+            void* pUserData = nullptr;
+            zel_pfnDriverExtensionFunctionCb_t prologue = nullptr;
+            zel_pfnDriverExtensionFunctionCb_t epilogue = nullptr;
+        };
+        std::mutex extensionCallbackMutex;
+        std::map<std::string, extension_function_callbacks_t> extensionCallbacks;
+
+        // Global gate for extension-function callbacks, toggled by the loader via
+        // zelDriverEnableTracing. Callbacks fire only when this is set AND a
+        // callback is registered for the function (two-level gate).
+        std::atomic<bool> extensionCallbacksEnabled{false};
+
         context_t();
         ~context_t();
 
@@ -68,7 +90,36 @@ namespace driver
     uint32_t ZE_APICALL zerTranslateDeviceHandleToIdentifier(ze_device_handle_t hDevice);
     ze_device_handle_t ZE_APICALL zerTranslateIdentifierToDeviceHandle(uint32_t identifier);
     ze_context_handle_t ZE_APICALL zerGetDefaultContext(void);
-
+
+    ///////////////////////////////////////////////////////////////////////////
+    // Extension-function callback prototype demonstration.
+    //
+    // "zeSampleExtFunc" is a stand-in vendor extension function reachable only by
+    // name via zeDriverGetExtensionFunctionAddress. Its body invokes any
+    // prologue/epilogue registered through zelDriverSetExtensionFunctionCallback,
+    // passing a typed params block (the driver knows its own signature).
+    typedef struct _ze_sample_ext_func_params_t
+    {
+        ze_driver_handle_t* phDriver;
+        uint32_t* pinput;
+        uint32_t** ppOutput;
+    } ze_sample_ext_func_params_t;
+
+    ze_result_t ZE_APICALL zeSampleExtFunc(
+        ze_driver_handle_t hDriver, uint32_t input, uint32_t* pOutput );
+
+    // Driver-side registration entry, resolved by name from the loader's
+    // zelDriverSetExtensionFunctionCallback forward.
+    ze_result_t ZE_APICALL zelDriverSetExtensionFunctionCallback(
+        ze_driver_handle_t hDriver, const char* functionName, void* pUserData,
+        zel_pfnDriverExtensionFunctionCb_t prologue,
+        zel_pfnDriverExtensionFunctionCb_t epilogue );
+
+    // Driver-side enable/disable of extension-function callbacks, resolved by
+    // name from the loader when the tracing layer is enabled/disabled.
+    ze_result_t ZE_APICALL zelDriverEnableTracing(
+        ze_driver_handle_t hDriver, ze_bool_t enable );
+
     extern context_t context;
 } // namespace driver
 

source/lib/ze_lib.cpp

@@ -205,6 +205,10 @@ namespace ze_lib
         if (loaderGetContext == nullptr) {
             std::string message = "ze_lib Context Init() zelLoaderGetContext missing";
             debug_trace_message(message, "");
+        } else {
+            // Cache the loader context portably (the loader symbols are not
+            // link-time visible in the static-loader build).
+            ze_lib::context->loaderContext = loaderGetContext();
         }
 
         std::string version_message = "Loader API Version to be requested is v" + std::to_string(ZE_MAJOR_VERSION(version)) + "." + std::to_string(ZE_MINOR_VERSION(version));
@@ -215,6 +219,7 @@ namespace ze_lib
         if( ZE_RESULT_SUCCESS == result ) {
             tracing_lib = zeLoaderGetTracingHandle();
         }
+        ze_lib::context->loaderContext = loader::context;
 
 #endif
 
@@ -635,6 +640,12 @@ zelEnableTracingLayer()
         if (ze_lib::context->pTracingZerDdiTable != nullptr) {
             ze_lib::context->zerDdiTable.exchange(ze_lib::context->pTracingZerDdiTable);
         }
+        // Propagate the enable to each active driver's extension-function tracing.
+        if (loader::context) {
+            for (auto &drv : loader::context->zeDrivers) {
+                loader::enableDriverExtensionTracing(drv, true);
+            }
+        }
     }
     #endif
     return ZE_RESULT_SUCCESS;
@@ -684,14 +695,91 @@ zelDisableTracingLayer()
     if (ze_lib::context->dynamicTracingSupported == false) {
         return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
     }
-    if (ze_lib::context->tracingLayerEnableCounter.fetch_sub(1) <= 1) {
+    // Guard against underflow: decrement only when the counter is > 0, so a
+    // disable with no matching enable (e.g. under ZE_ENABLE_TRACING_LAYER=1) is a
+    // safe no-op rather than wrapping the unsigned counter and corrupting state.
+    uint32_t prev = ze_lib::context->tracingLayerEnableCounter.load();
+    while (prev > 0 &&
+           !ze_lib::context->tracingLayerEnableCounter.compare_exchange_weak(prev, prev - 1)) {
+        // prev is reloaded by compare_exchange_weak on failure
+    }
+    if (prev == 1) {
+        // 1 -> 0 transition: tear down the dynamic tracing DDI tables.
         ze_lib::context->zeDdiTable.exchange(&ze_lib::context->initialzeDdiTable);
         if (ze_lib::context->pTracingZerDdiTable != nullptr) {
             ze_lib::context->zerDdiTable.exchange(&ze_lib::context->initialzerDdiTable);
         }
+        // Disable per-driver extension tracing, unless tracing was enabled
+        // statically via ZE_ENABLE_TRACING_LAYER (documented to stay on for the
+        // whole application) - respect that sticky state.
+        if (loader::context && !loader::context->tracingLayerEnabled) {
+            for (auto &drv : loader::context->zeDrivers) {
+                loader::enableDriverExtensionTracing(drv, false);
+            }
+        }
     }
     #endif
     return ZE_RESULT_SUCCESS;
 }
 
+ze_result_t ZE_APICALL
+zelDriverSetExtensionFunctionCallback(
+    ze_driver_handle_t hDriver,
+    const char* functionName,
+    void* pUserData,
+    zel_pfnDriverExtensionFunctionCb_t prologue,
+    zel_pfnDriverExtensionFunctionCb_t epilogue
+    )
+{
+    if( nullptr == hDriver )
+        return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
+    if( nullptr == functionName )
+        return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
+    if( ze_lib::destruction )
+        return ZE_RESULT_ERROR_UNINITIALIZED;
+
+    // Type of the driver-side registration entry point, discovered by name.
+    typedef ze_result_t (ZE_APICALL *zelDriverSetExtensionFunctionCallback_t)(
+        ze_driver_handle_t, const char*, void*,
+        zel_pfnDriverExtensionFunctionCb_t, zel_pfnDriverExtensionFunctionCb_t );
+
+    // Resolve the driver's registration entry via the standard extension-address
+    // lookup on this specific driver. The driver owns the registry and the
+    // invocation of the callbacks from inside the extension-function body.
+    auto pfnGetExtensionFunctionAddress =
+        ze_lib::context->zeDdiTable.load()->Driver.pfnGetExtensionFunctionAddress;
+    if( nullptr == pfnGetExtensionFunctionAddress ) {
+        if( !ze_lib::context->isInitialized )
+            return ZE_RESULT_ERROR_UNINITIALIZED;
+        return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    }
+
+    void* pfnRaw = nullptr;
+    ze_result_t result = pfnGetExtensionFunctionAddress(
+        hDriver, "zelDriverSetExtensionFunctionCallback", &pfnRaw );
+    if( result != ZE_RESULT_SUCCESS )
+        return result;
+    if( nullptr == pfnRaw )
+        return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+
+    // Sync this driver's global extension-tracing gate to the current tracing
+    // state. Covers static ZE_ENABLE_TRACING_LAYER enablement (the driver becomes
+    // usable only by the time the app registers) and drivers registered on after
+    // a dynamic enable. Disabling is handled centrally by zelDisableTracingLayer
+    // (respecting sticky env), so only propagate the enable here.
+    bool tracingOn = ze_lib::context->tracingLayerEnableCounter.load() > 0;
+    if( !tracingOn && ze_lib::context->loaderContext )
+        tracingOn = ze_lib::context->loaderContext->tracingLayerEnabled;
+    if( tracingOn ) {
+        void* pfnEnableRaw = nullptr;
+        if( ZE_RESULT_SUCCESS == pfnGetExtensionFunctionAddress(
+                hDriver, "zelDriverEnableTracing", &pfnEnableRaw ) && pfnEnableRaw ) {
+            reinterpret_cast<zel_pfnDriverEnableTracing_t>(pfnEnableRaw)( hDriver, true );
+        }
+    }
+
+    auto pfnSet = reinterpret_cast<zelDriverSetExtensionFunctionCallback_t>( pfnRaw );
+    return pfnSet( hDriver, functionName, pUserData, prologue, epilogue );
+}
+
 } //extern "c"