Add tracing layer feature to modify return values

include/layers/zel_tracing_api.h

@@ -187,6 +187,37 @@ zelTracerSetEnabled(
     ze_bool_t enable                                ///< [in] enable the tracer if true; disable if false
     );
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Overrides the value returned to the application by the API call
+///        currently being traced.
+/// 
+/// @details
+///     - Must be called from within an epilogue callback of an API that
+///       returns ::ze_result_t.
+///     - Epilogues run in tracer registration order. Each later epilogue
+///       receives the overridden value as its `result` parameter; the last
+///       override wins.
+///     - Loaders that predate this function do not provide it. A static
+///       loader built with this function returns
+///       ::ZE_RESULT_ERROR_UNSUPPORTED_FEATURE when the runtime loader lacks it.
+///     - The value is not checked. Replacing an error, including one from the
+///       validation layer, hides it from the application; override only the
+///       results you handle. See source/layers/tracing/README.md.
+///     - Other return types (for example ::ze_context_handle_t) are reserved
+///       for future type-specific zelTracerSet<Type>ReturnValue functions.
+/// 
+/// @returns
+///     - ::ZE_RESULT_SUCCESS
+///     - ::ZE_RESULT_ERROR_UNINITIALIZED
+///     - ::ZE_RESULT_ERROR_UNSUPPORTED_FEATURE
+///         + the tracing layer in use does not support this function
+///     - ::ZE_RESULT_ERROR_INVALID_ARGUMENT
+///         + not called from an epilogue of an API that returns ::ze_result_t
+ZE_APIEXPORT ze_result_t ZE_APICALL
+zelTracerSetResultReturnValue(
+    ze_result_t value                           ///< [in] value the traced API returns to the application
+    );
+
 #if !defined(__GNUC__)
 #pragma endregion
 #endif

source/layers/tracing/README.md

@@ -175,6 +175,63 @@ Callback handlers are functions that are implemented by the application, and reg
 
     - __ppTracerInstanceUserData__ : a per-tracer, per-instance, per-thread storage location; typically used for passing data from the prologue to the epilogue. See example below.
 
+## Overriding the Return Value
+
+An epilogue callback can replace the value that the traced API returns to the application by calling __zelTracerSetResultReturnValue(ze_result_t value)__:
+
+```c
+void OnExitEventHostSynchronize(
+    ze_event_host_synchronize_params_t* params,
+    ze_result_t result,
+    void* pTracerUserData,
+    void** ppTracerInstanceUserData )
+{
+    // Only handle NOT_READY; every other result passes through unchanged.
+    if (result != ZE_RESULT_NOT_READY)
+        return;
+    while (result == ZE_RESULT_NOT_READY) {
+        run_progress();
+        // Calls made from inside a callback are not traced again.
+        result = zeEventHostSynchronize(*params->phEvent, *params->ptimeout);
+    }
+    // Report the final result of the wait, which may itself be an error.
+    zelTracerSetResultReturnValue(result);
+}
+```
+
+Notes:
+- It is only valid from an epilogue of an API that returns `ze_result_t`. Anywhere else it returns `ZE_RESULT_ERROR_INVALID_ARGUMENT` and changes nothing.
+- Epilogues run in tracer registration order. Each later epilogue receives the overridden value as its `result` parameter. The last override wins.
+- APIs that return other types (`zeDriverGetDefaultContext`, `zerGetDefaultContext`, `zerTranslateIdentifierToDeviceHandle`, `zerTranslateDeviceHandleToIdentifier`) are reserved for future type-specific `zelTracerSet<Type>ReturnValue` functions.
+- Older loaders do not provide this function:
+    - Applications that link the static loader get `ZE_RESULT_ERROR_UNSUPPORTED_FEATURE` and the original return value is left unchanged.
+    - Applications that link the dynamic loader fail with an undefined symbol when the function is first called. To support older loaders, look the function up at runtime (`dlsym` / `GetProcAddress`) instead of calling it directly.
+
+### Warning: an override can hide real errors
+
+The value set by __zelTracerSetResultReturnValue__ is what the application sees. The loader does not check it and does not log it. If you replace an error with `ZE_RESULT_SUCCESS`, the application believes the call worked when it did not.
+
+The tracing layer sits above the validation layer (application → tracing layer → validation layer → driver). This means an epilogue can also hide errors that come from the validation layer:
+
+- The validation layer still runs first and still stops an invalid call before it reaches the driver.
+- The validation layer's own checks still see the real result.
+- The epilogue still receives the real error in its `result` parameter.
+- The application only sees the value from the last override. For example, an epilogue that always sets `ZE_RESULT_SUCCESS` turns `ZE_RESULT_ERROR_INVALID_NULL_POINTER` from a bad argument into `ZE_RESULT_SUCCESS`.
+
+When an error is hidden:
+
+- Output parameters are not written, but the application reads them as if they were. For example, a count stays 0, a handle stays null, or a buffer holds old data.
+- Handles and objects may not exist, so later calls fail far from the real cause, or crash.
+- Device loss and out-of-memory errors are masked, so the application keeps going on a device or allocation that is not usable.
+- Tools that depend on return codes, such as the validation layer's error reporting to the application or retry logic, no longer work.
+
+To use overrides safely:
+
+- Only replace the specific results you handle. For example, replace `ZE_RESULT_NOT_READY` after you have waited for completion. Do not override unconditionally.
+- Pass every other result through unchanged. Do not call __zelTracerSetResultReturnValue__ for results you do not handle.
+- Do not turn an error into `ZE_RESULT_SUCCESS` unless your callback has completed the work the application asked for, including all output parameters.
+- When you debug a failure with tracing enabled, first disable any tracer that overrides results, then run with `ZE_ENABLE_VALIDATION_LAYER=1` and `ZE_ENABLE_PARAMETER_VALIDATION=1`.
+
 ##  __zeInit__ is traceable for all calls subsequent from the creation and enabling of the tracer itself.
 
 ## Enabling, Disabling and Destruction

source/layers/tracing/tracing_imp.cpp

@@ -12,6 +12,7 @@
 namespace tracing_layer {
 
 thread_local ze_bool_t tracingInProgress = 0;
+thread_local ze_result_t *pEpilogueResultReturnValue = nullptr;
 
 struct APITracerContextImp *pGlobalAPITracerContextImp;
 

source/layers/tracing/tracing_imp.h

@@ -33,6 +33,9 @@
 namespace tracing_layer {
 
 extern thread_local ze_bool_t tracingInProgress;
+// Points at the in-flight ze_result_t return value while epilogues run, so
+// zelTracerSetResultReturnValue can override it; nullptr at all other times.
+extern thread_local ze_result_t *pEpilogueResultReturnValue;
 extern struct APITracerContextImp *pGlobalAPITracerContextImp;
 
 // Keys registration and per-call fan-out by (hDriver, functionName) so a callback
@@ -279,6 +282,12 @@ class APITracerCallbackDataImp {
         }                                                                            \
     }
 
+// Only APIs returning ze_result_t expose their return value to
+// zelTracerSetResultReturnValue; other return types get no slot.
+template <typename TRet>
+inline ze_result_t *resultReturnValueSlot(TRet *) { return nullptr; }
+inline ze_result_t *resultReturnValueSlot(ze_result_t *pRet) { return pRet; }
+
 template <typename TRet, typename TFunction_pointer, typename TParams, typename TTracer,
           typename TTracerPrologCallbacks, typename TTracerEpilogCallbacks,
           typename... Args>
@@ -310,12 +319,14 @@ APITracerWrapperImp(TFunction_pointer zeApiPtr, TParams paramsStruct,
                 &ppTracerInstanceUserData[i]);
     }
     ret = zeApiPtr(args...);
+    tracing_layer::pEpilogueResultReturnValue = resultReturnValueSlot(&ret);
     for (size_t i = 0; i < callbacksEpilogs->size(); i++) {
         if (callbacksEpilogs->at(i).current_api_callback != nullptr)
             callbacksEpilogs->at(i).current_api_callback(
                 paramsStruct, ret, callbacksEpilogs->at(i).pUserData,
                 &ppTracerInstanceUserData[i]);
     }
+    tracing_layer::pEpilogueResultReturnValue = nullptr;
     tracing_layer::tracingInProgress = 0;
     tracing_layer::pGlobalAPITracerContextImp->releaseActivetracersList();
     return ret;

source/layers/tracing/ze_tracing.cpp

@@ -192,6 +192,21 @@ zelTracerDriverExtensionRegisterCallback(
     return applyLoaderExtensionInstall( hDriver, functionName, install );
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Overrides the value the traced API returns to the application. Only
+///        valid from an epilogue of an API that returns ze_result_t.
+ZE_DLLEXPORT ze_result_t ZE_APICALL
+zelTracerSetResultReturnValue(
+    ze_result_t value
+    )
+{
+    if( nullptr == tracing_layer::pEpilogueResultReturnValue )
+        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
+
+    *tracing_layer::pEpilogueResultReturnValue = value;
+    return ZE_RESULT_SUCCESS;
+}
+
 ZE_DLLEXPORT ze_result_t ZE_APICALL
 zelLoaderGetVersion(zel_component_version_t *version)    
 {

source/lib/zel_tracing_libapi.cpp

@@ -177,6 +177,37 @@ zelTracerSetEnabled(
     return pfnSetEnabled( hTracer, enable );
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Overrides the value returned by the API call currently being traced.
+///        Only valid from an epilogue of an API that returns ze_result_t.
+/// 
+/// @returns
+///     - ::ZE_RESULT_SUCCESS
+///     - ::ZE_RESULT_ERROR_UNINITIALIZED
+///     - ::ZE_RESULT_ERROR_UNSUPPORTED_FEATURE
+///     - ::ZE_RESULT_ERROR_INVALID_ARGUMENT
+ze_result_t ZE_APICALL
+zelTracerSetResultReturnValue(
+    ze_result_t value                           ///< [in] value the traced API returns to the application
+    )
+{
+    if(ze_lib::destruction)
+        return ZE_RESULT_ERROR_UNINITIALIZED;
+    if(!ze_lib::context->tracing_lib)
+        return ZE_RESULT_ERROR_UNINITIALIZED;
+
+    typedef ze_result_t (ZE_APICALL *ze_pfnSetResultReturnValue_t)( ze_result_t );
+
+    // Looked up by name so an older tracing layer without it reports unsupported.
+    auto func = reinterpret_cast<ze_pfnSetResultReturnValue_t>(
+        GET_FUNCTION_PTR(ze_lib::context->tracing_lib,
+                         "zelTracerSetResultReturnValue") );
+    if(!func)
+        return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+
+    return func( value );
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Registers a prologue/epilogue callback on a tracer for a named
 ///        extension function of a specific driver. See loader/ze_loader.h.

test/CMakeLists.txt

@@ -1158,6 +1158,22 @@ set_property(TEST test_ze_and_zer_tracing_dynamic PROPERTY ENVIRONMENT "ZE_ENABL
 add_test(NAME test_zer_unsupported_and_ze_tracing_dynamic COMMAND tests --gtest_filter=*TracingParameterizedTest*GivenLoaderWithDynamicTracingEnabledAndZerApisUnsupportedAndBothZeAndZerCallbacksRegisteredWhenCallingBothApisThenTracingWorksForZeAndZerCallbacksAreStillInvoked*)
 set_property(TEST test_zer_unsupported_and_ze_tracing_dynamic PROPERTY ENVIRONMENT "ZE_ENABLE_NULL_DRIVER=1")
 
+# zelTracerSetResultReturnValue tests
+add_test(NAME test_tracing_set_result_return_value_static COMMAND tests --gtest_filter=*TracingParameterizedTest*GivenEpilogueCallbackWhenSettingResultReturnValueThenCallerReceivesOverriddenValue*)
+set_property(TEST test_tracing_set_result_return_value_static PROPERTY ENVIRONMENT "ZE_ENABLE_NULL_DRIVER=1")
+
+add_test(NAME test_tracing_set_result_return_value_dynamic COMMAND tests --gtest_filter=*TracingParameterizedTest*GivenDynamicTracingAndEpilogueCallbackWhenSettingResultReturnValueThenCallerReceivesOverriddenValue*)
+set_property(TEST test_tracing_set_result_return_value_dynamic PROPERTY ENVIRONMENT "ZE_ENABLE_NULL_DRIVER=1")
+
+add_test(NAME test_tracing_set_result_return_value_multiple_tracers COMMAND tests --gtest_filter=*TracingParameterizedTest*GivenMultipleTracersWhenEarlierEpilogueSetsResultReturnValueThenLaterEpilogueSeesItAndLastWriteWins*)
+set_property(TEST test_tracing_set_result_return_value_multiple_tracers PROPERTY ENVIRONMENT "ZE_ENABLE_NULL_DRIVER=1")
+
+add_test(NAME test_tracing_set_result_return_value_outside_epilogue COMMAND tests --gtest_filter=*TracingParameterizedTest*GivenNoEpilogueInProgressWhenSettingResultReturnValueThenInvalidArgumentIsReturnedAndResultIsUnchanged*)
+set_property(TEST test_tracing_set_result_return_value_outside_epilogue PROPERTY ENVIRONMENT "ZE_ENABLE_NULL_DRIVER=1")
+
+add_test(NAME test_tracing_set_result_return_value_non_result_api COMMAND tests --gtest_filter=*TracingParameterizedTest*GivenEpilogueOfApiNotReturningZeResultWhenSettingResultReturnValueThenInvalidArgumentIsReturned*)
+set_property(TEST test_tracing_set_result_return_value_non_result_api PROPERTY ENVIRONMENT "ZE_ENABLE_NULL_DRIVER=1")
+
 
 
 # ZER API Validation Layer Tests