Initial Feature commit for Extension Callback Tracing

include/loader/ze_loader.h

@@ -13,6 +13,7 @@
 #endif
 
 #include "../ze_api.h"
+#include "../layers/zel_tracing_register_cb.h"
 
 #if !defined(__cplusplus)
 #include <stdbool.h>
@@ -565,6 +566,106 @@ 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 Signature of the per-driver hook the loader/tracing-layer uses to
+///        install its extension-function interception wrappers on a driver.
+///
+/// A driver that supports extension-function tracing exposes this by name
+/// ("zelDriverSetLoaderCallbackForExtension") via
+/// zeDriverGetExtensionFunctionAddress. The tracing layer calls it to register a
+/// single loader-owned prologue/epilogue wrapper (plus an opaque loader context)
+/// for the named extension function. The driver invokes @p loaderPrologue before,
+/// and @p loaderEpilogue after, the body of the extension function named
+/// @p functionName, forwarding @p pLoaderContext back unchanged. Passing null for
+/// both wrappers unregisters. The loader owns the fan-out to any number of
+/// registered tracers, so the driver stores at most one wrapper per function.
+typedef ze_result_t (ZE_APICALL *zel_pfnDriverSetLoaderCallbackForExtension_t)(
+    ze_driver_handle_t hDriver,                      // [in] handle of the driver instance
+    const char* functionName,                        // [in] extension function name to intercept
+    zel_pfnDriverExtensionFunctionCb_t loaderPrologue, // [in][optional] loader prologue wrapper
+    zel_pfnDriverExtensionFunctionCb_t loaderEpilogue, // [in][optional] loader epilogue wrapper
+    void* pLoaderContext                             // [in][optional] loader context echoed to wrappers
+    );
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Registers a prologue or epilogue callback on a tracer for a named
+///        extension function of a specific driver.
+///
+/// Extension functions obtained by string name via
+/// zeDriverGetExtensionFunctionAddress() return a raw driver pointer that the
+/// application calls directly, bypassing the loader and therefore the per-API
+/// tracing interceptors. This API routes such functions through the same tracer
+/// (::zel_tracer_handle_t) infrastructure used for core APIs: the tracing layer
+/// installs a loader-owned wrapper on @p hDriver (via the driver's
+/// zelDriverSetLoaderCallbackForExtension hook) and fans out to every enabled
+/// tracer that registered @p functionName for @p hDriver.
+///
+/// Registration is keyed by (@p hDriver, @p functionName) and is order-independent
+/// relative to zeDriverGetExtensionFunctionAddress() — it takes effect on the next
+/// invocation even if the application already cached the function pointer. The
+/// callback receives the tracer's pUserData (from ::zelTracerCreate) as
+/// pTracerUserData. Multiple tracers may register the same function to stack
+/// callbacks. The callbacks fire only when the tracing layer is enabled for the
+/// driver and the tracer is enabled.
+///
+/// @param[in] hTracer        handle of the tracer to register the callback on
+/// @param[in] hDriver        handle of the driver whose extension function to trace
+/// @param[in] functionName   name of the extension function to intercept
+/// @param[in] callback_type  ::ZEL_REGISTER_PROLOGUE or ::ZEL_REGISTER_EPILOGUE
+/// @param[in] pCallback      handler to register (null clears that slot)
+///
+/// @return
+///   - ZE_RESULT_SUCCESS on success (including clearing a slot).
+///   - ZE_RESULT_ERROR_UNINITIALIZED if the loader/tracing layer is not initialized.
+///   - ZE_RESULT_ERROR_UNSUPPORTED_FEATURE if the driver does not implement the hook.
+///   - ZE_RESULT_ERROR_INVALID_NULL_HANDLE if @p hTracer or @p hDriver is null.
+///   - ZE_RESULT_ERROR_INVALID_NULL_POINTER if @p functionName is null.
+///   - ZE_RESULT_ERROR_INVALID_ARGUMENT if the tracer is not in the disabled state.
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelTracerDriverExtensionRegisterCallback(
+    zel_tracer_handle_t hTracer,                     // [in] handle of the tracer
+    ze_driver_handle_t hDriver,                      // [in] handle of the driver instance
+    const char* functionName,                        // [in] extension function name to intercept
+    zel_tracer_reg_t callback_type,                  // [in] prologue or epilogue
+    zel_pfnDriverExtensionFunctionCb_t pCallback     // [in][optional] handler (null clears slot)
+    );
+
 #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, "zelDriverSetLoaderCallbackForExtension" ) ) {
+                *ppFunctionAddress = reinterpret_cast<void*>( &driver::zelDriverSetLoaderCallbackForExtension );
+                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,82 @@ 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 the single loader wrapper registered for this function.
+        context_t::loader_extension_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: the wrapper fires only when tracing is globally enabled
+        // AND a loader wrapper 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.loaderPrologue )
+            cbs.loaderPrologue( &params, result, cbs.pLoaderContext, &pInstanceData );
+
+        // The (trivial) work of the extension function.
+        if( nullptr != pOutput )
+            *pOutput = input * 2;
+
+        if( fire && nullptr != cbs.loaderEpilogue )
+            cbs.loaderEpilogue( &params, result, cbs.pLoaderContext, &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 loader-callback registration entry (resolved by name
+    ///        from the tracing layer). Stores the single loader wrapper (+ opaque
+    ///        context) per function name; null+null unregisters.
+    ze_result_t ZE_APICALL zelDriverSetLoaderCallbackForExtension(
+        ze_driver_handle_t, const char* functionName,
+        zel_pfnDriverExtensionFunctionCb_t loaderPrologue,
+        zel_pfnDriverExtensionFunctionCb_t loaderEpilogue,
+        void* pLoaderContext )
+    {
+        if( nullptr == functionName )
+            return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
+
+        std::lock_guard<std::mutex> lock( context.extensionCallbackMutex );
+        if( nullptr == loaderPrologue && nullptr == loaderEpilogue ) {
+            context.extensionCallbacks.erase( functionName );
+        } else {
+            auto& entry = context.extensionCallbacks[ functionName ];
+            entry.loaderPrologue = loaderPrologue;
+            entry.loaderEpilogue = loaderEpilogue;
+            entry.pLoaderContext = pLoaderContext;
+        }
+        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,26 @@ namespace driver
         std::vector<BaseNullHandle*> globalBaseNullHandle;
 	bool ddiExtensionSupported = false;
 	std::vector<char *> env_vars{};
+
+        // Registry for zelDriverSetLoaderCallbackForExtension: maps an extension
+        // function name to the single loader-owned wrapper the driver invokes
+        // from that function's body. The loader/tracing-layer owns the fan-out to
+        // any number of tracers, so the driver stores at most one wrapper (plus
+        // an opaque loader context) per function. Keyed by name (order-
+        // independent vs fetch).
+        struct loader_extension_callbacks_t {
+            zel_pfnDriverExtensionFunctionCb_t loaderPrologue = nullptr;
+            zel_pfnDriverExtensionFunctionCb_t loaderEpilogue = nullptr;
+            void* pLoaderContext = nullptr;
+        };
+        std::mutex extensionCallbackMutex;
+        std::map<std::string, loader_extension_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 +93,38 @@ 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 the single
+    // loader-owned wrapper registered through zelDriverSetLoaderCallbackForExtension,
+    // 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 loader-callback registration entry, resolved by name from the
+    // tracing layer. Stores the single loader wrapper (+ context) per function;
+    // null+null unregisters.
+    ze_result_t ZE_APICALL zelDriverSetLoaderCallbackForExtension(
+        ze_driver_handle_t hDriver, const char* functionName,
+        zel_pfnDriverExtensionFunctionCb_t loaderPrologue,
+        zel_pfnDriverExtensionFunctionCb_t loaderEpilogue,
+        void* pLoaderContext );
+
+    // 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/layers/tracing/tracing.h

@@ -10,6 +10,7 @@
 #include "ze_api.h"
 #include "layers/zel_tracing_api.h"
 #include "layers/zel_tracing_register_cb.h"
+#include "loader/ze_loader.h"
 #include "ze_tracing_cb_structs.h"
 #include "zer_tracing_cb_structs.h"
 
@@ -30,6 +31,10 @@ struct APITracer : _zel_tracer_handle_t {
     virtual zel_zer_all_callbacks_t& getZerProEpilogues(zel_tracer_reg_t callback_type, ze_result_t& result) = 0;
     virtual ze_result_t resetAllCallbacks() = 0;
     virtual ze_result_t enableTracer(ze_bool_t enable) = 0;
+    virtual ze_result_t registerExtensionCallback(ze_driver_handle_t hDriver,
+                                                  const char *functionName,
+                                                  zel_tracer_reg_t callback_type,
+                                                  zel_pfnDriverExtensionFunctionCb_t pCallback) = 0;
 };
 
 ze_result_t createAPITracer(const zel_tracer_desc_t *desc, zel_tracer_handle_t *phTracer);