[UR][L0v2] Add spec and implementation for urGraphSetDestructionCallback

unified-runtime/include/unified-runtime/ur_api.h

@@ -510,6 +510,8 @@ typedef enum ur_function_t {
   UR_FUNCTION_USM_HOST_ALLOC_UNREGISTER_EXP = 313,
   /// Enumerator for ::urQueueGetGraphExp
   UR_FUNCTION_QUEUE_GET_GRAPH_EXP = 314,
+  /// Enumerator for ::urGraphSetDestructionCallbackExp
+  UR_FUNCTION_GRAPH_SET_DESTRUCTION_CALLBACK_EXP = 315,
   /// @cond
   UR_FUNCTION_FORCE_UINT32 = 0x7fffffff
   /// @endcond
@@ -13874,6 +13876,34 @@ UR_APIEXPORT ur_result_t UR_APICALL urGraphIsEmptyExp(
     /// [out] Pointer to a boolean where the result will be stored.
     bool *pResult);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Callback function invoked when a graph is destroyed.
+typedef void (*ur_exp_graph_destruction_callback_t)(
+    /// [in] pointer to user data to be passed to the callback
+    void *pUserData);
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Register a callback to be invoked when the graph is destroyed.
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_UNINITIALIZED
+///     - ::UR_RESULT_ERROR_DEVICE_LOST
+///     - ::UR_RESULT_ERROR_ADAPTER_SPECIFIC
+///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
+///         + `NULL == hGraph`
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+///         + `NULL == pfnCallback`
+UR_APIEXPORT ur_result_t UR_APICALL urGraphSetDestructionCallbackExp(
+    /// [in] Handle of the graph to register the callback for.
+    ur_exp_graph_handle_t hGraph,
+    /// [in] Function pointer to the callback. The callback must not access
+    /// hGraph.
+    ur_exp_graph_destruction_callback_t pfnCallback,
+    /// [in][optional] Pointer to user data to be passed to the callback. The
+    /// user data must not reference hGraph.
+    void *pUserData);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Dump the contents of the recorded graph to the provided file path.
 ///
@@ -16414,6 +16444,16 @@ typedef struct ur_graph_is_empty_exp_params_t {
   bool **ppResult;
 } ur_graph_is_empty_exp_params_t;
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function parameters for urGraphSetDestructionCallbackExp
+/// @details Each entry is a pointer to the parameter passed to the function;
+///     allowing the callback the ability to modify the parameter's value
+typedef struct ur_graph_set_destruction_callback_exp_params_t {
+  ur_exp_graph_handle_t *phGraph;
+  ur_exp_graph_destruction_callback_t *ppfnCallback;
+  void **ppUserData;
+} ur_graph_set_destruction_callback_exp_params_t;
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function parameters for urGraphDumpContentsExp
 /// @details Each entry is a pointer to the parameter passed to the function;

unified-runtime/include/unified-runtime/ur_api_funcs.def

@@ -217,6 +217,7 @@ _UR_API(urGraphInstantiateGraphExp)
 _UR_API(urGraphDestroyExp)
 _UR_API(urGraphExecutableGraphDestroyExp)
 _UR_API(urGraphIsEmptyExp)
+_UR_API(urGraphSetDestructionCallbackExp)
 _UR_API(urGraphDumpContentsExp)
 _UR_API(urIPCGetMemHandleExp)
 _UR_API(urIPCPutMemHandleExp)

unified-runtime/include/unified-runtime/ur_ddi.h

@@ -1887,6 +1887,11 @@ typedef ur_result_t(UR_APICALL *ur_pfnGraphExecutableGraphDestroyExp_t)(
 typedef ur_result_t(UR_APICALL *ur_pfnGraphIsEmptyExp_t)(ur_exp_graph_handle_t,
                                                          bool *);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function-pointer for urGraphSetDestructionCallbackExp
+typedef ur_result_t(UR_APICALL *ur_pfnGraphSetDestructionCallbackExp_t)(
+    ur_exp_graph_handle_t, ur_exp_graph_destruction_callback_t, void *);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function-pointer for urGraphDumpContentsExp
 typedef ur_result_t(UR_APICALL *ur_pfnGraphDumpContentsExp_t)(
@@ -1900,6 +1905,7 @@ typedef struct ur_graph_exp_dditable_t {
   ur_pfnGraphDestroyExp_t pfnDestroyExp;
   ur_pfnGraphExecutableGraphDestroyExp_t pfnExecutableGraphDestroyExp;
   ur_pfnGraphIsEmptyExp_t pfnIsEmptyExp;
+  ur_pfnGraphSetDestructionCallbackExp_t pfnSetDestructionCallbackExp;
   ur_pfnGraphDumpContentsExp_t pfnDumpContentsExp;
 } ur_graph_exp_dditable_t;
 

unified-runtime/include/unified-runtime/ur_print.h

@@ -3725,6 +3725,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urPrintGraphIsEmptyExpParams(
     const struct ur_graph_is_empty_exp_params_t *params, char *buffer,
     const size_t buff_size, size_t *out_size);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Print ur_graph_set_destruction_callback_exp_params_t struct
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_SIZE
+///         - `buff_size < out_size`
+UR_APIEXPORT ur_result_t UR_APICALL urPrintGraphSetDestructionCallbackExpParams(
+    const struct ur_graph_set_destruction_callback_exp_params_t *params,
+    char *buffer, const size_t buff_size, size_t *out_size);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Print ur_graph_dump_contents_exp_params_t struct
 /// @returns

unified-runtime/include/unified-runtime/ur_print.hpp

@@ -1377,6 +1377,9 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_function_t value) {
   case UR_FUNCTION_QUEUE_GET_GRAPH_EXP:
     os << "UR_FUNCTION_QUEUE_GET_GRAPH_EXP";
     break;
+  case UR_FUNCTION_GRAPH_SET_DESTRUCTION_CALLBACK_EXP:
+    os << "UR_FUNCTION_GRAPH_SET_DESTRUCTION_CALLBACK_EXP";
+    break;
   default:
     os << "unknown enumerator";
     break;
@@ -21568,6 +21571,33 @@ inline std::ostream &operator<<(
   return os;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Print operator for the ur_graph_set_destruction_callback_exp_params_t
+/// type
+/// @returns
+///     std::ostream &
+inline std::ostream &operator<<(
+    std::ostream &os,
+    [[maybe_unused]] const struct ur_graph_set_destruction_callback_exp_params_t
+        *params) {
+
+  os << ".hGraph = ";
+
+  ur::details::printPtr(os, *(params->phGraph));
+
+  os << ", ";
+  os << ".pfnCallback = ";
+
+  os << reinterpret_cast<void *>(*(params->ppfnCallback));
+
+  os << ", ";
+  os << ".pUserData = ";
+
+  ur::details::printPtr(os, *(params->ppUserData));
+
+  return os;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Print operator for the ur_graph_dump_contents_exp_params_t type
 /// @returns
@@ -23232,6 +23262,9 @@ inline ur_result_t UR_APICALL printFunctionParams(std::ostream &os,
   case UR_FUNCTION_GRAPH_IS_EMPTY_EXP: {
     os << (const struct ur_graph_is_empty_exp_params_t *)params;
   } break;
+  case UR_FUNCTION_GRAPH_SET_DESTRUCTION_CALLBACK_EXP: {
+    os << (const struct ur_graph_set_destruction_callback_exp_params_t *)params;
+  } break;
   case UR_FUNCTION_GRAPH_DUMP_CONTENTS_EXP: {
     os << (const struct ur_graph_dump_contents_exp_params_t *)params;
   } break;

unified-runtime/scripts/core/EXP-GRAPH.rst

@@ -72,6 +72,10 @@ Changelog
 | 1.3       | Add ${x}QueueGetGraphExp to retrieve graph  |
 |           | handle from queue in capture mode.          |
 +-----------+---------------------------------------------+
+| 1.4       | Add ${x}GraphSetDestructionCallbackExp to   |
+|           | register user callbacks invoked on graph    |
+|           | destruction.                                |
++-----------+---------------------------------------------+
 
 Support
 --------------------------------------------------------------------------------
@@ -84,3 +88,4 @@ Contributors
 
 * Krzysztof, Filipek `krzysztof.filipek@intel.com <krzysztof.filipek@intel.com>`_
 * Krzysztof, Swiecicki `krzysztof.swiecicki@intel.com <krzysztof.swiecicki@intel.com>`_
+* Matthew, Michel `matthew.michel@intel.com <matthew.michel@intel.com>`_

unified-runtime/scripts/core/exp-graph.yml

@@ -187,6 +187,30 @@ params:
 returns:
     - $X_RESULT_ERROR_INVALID_GRAPH
 --- #--------------------------------------------------------------------------
+type: fptr_typedef
+desc: "Callback function invoked when a graph is destroyed."
+name: $x_exp_graph_destruction_callback_t
+return: void
+params:
+    - type: void*
+      name: pUserData
+      desc: "[in] pointer to user data to be passed to the callback"
+--- #--------------------------------------------------------------------------
+type: function
+desc: "Register a callback to be invoked when the graph is destroyed."
+class: $xGraph
+name: SetDestructionCallbackExp
+params:
+    - type: $x_exp_graph_handle_t
+      name: hGraph
+      desc: "[in] Handle of the graph to register the callback for."
+    - type: $x_exp_graph_destruction_callback_t
+      name: pfnCallback
+      desc: "[in] Function pointer to the callback. The callback must not access hGraph."
+    - type: void*
+      name: pUserData
+      desc: "[in][optional] Pointer to user data to be passed to the callback. The user data must not reference hGraph."
+--- #--------------------------------------------------------------------------
 type: function
 desc: "Dump the contents of the recorded graph to the provided file path."
 class: $xGraph

unified-runtime/scripts/core/registry.yml

@@ -724,7 +724,10 @@ etors:
 - name: QUEUE_GET_GRAPH_EXP
   desc: Enumerator for $xQueueGetGraphExp
   value: '314'
-max_id: '314'
+- name: GRAPH_SET_DESTRUCTION_CALLBACK_EXP
+  desc: Enumerator for $xGraphSetDestructionCallbackExp
+  value: '315'
+max_id: '315'
 ---
 type: enum
 desc: Defines structure types

unified-runtime/source/adapters/cuda/graph.cpp

@@ -29,6 +29,13 @@ urGraphIsEmptyExp(ur_exp_graph_handle_t /* hGraph */, bool * /* pIsEmpty */) {
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urGraphSetDestructionCallbackExp(
+    ur_exp_graph_handle_t /* hGraph */,
+    ur_exp_graph_destruction_callback_t /* pfnCallback */,
+    void * /* pUserData */) {
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urGraphDumpContentsExp(
     ur_exp_graph_handle_t /* hGraph */, const char * /* pDotFilePath */) {
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;

unified-runtime/source/adapters/cuda/ur_interface_loader.cpp

@@ -89,6 +89,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urGetGraphExpProcAddrTable(
   pDdiTable->pfnDestroyExp = urGraphDestroyExp;
   pDdiTable->pfnExecutableGraphDestroyExp = urGraphExecutableGraphDestroyExp;
   pDdiTable->pfnIsEmptyExp = urGraphIsEmptyExp;
+  pDdiTable->pfnSetDestructionCallbackExp = urGraphSetDestructionCallbackExp;
   pDdiTable->pfnDumpContentsExp = urGraphDumpContentsExp;
 
   return UR_RESULT_SUCCESS;

unified-runtime/source/adapters/hip/graph.cpp

@@ -29,6 +29,13 @@ urGraphIsEmptyExp(ur_exp_graph_handle_t /* hGraph */, bool * /* pIsEmpty */) {
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urGraphSetDestructionCallbackExp(
+    ur_exp_graph_handle_t /* hGraph */,
+    ur_exp_graph_destruction_callback_t /* pfnCallback */,
+    void * /* pUserData */) {
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urGraphDumpContentsExp(
     ur_exp_graph_handle_t /* hGraph */, const char * /* pDotFilePath */) {
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;

unified-runtime/source/adapters/hip/ur_interface_loader.cpp

@@ -89,6 +89,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urGetGraphExpProcAddrTable(
   pDdiTable->pfnDestroyExp = urGraphDestroyExp;
   pDdiTable->pfnExecutableGraphDestroyExp = urGraphExecutableGraphDestroyExp;
   pDdiTable->pfnIsEmptyExp = urGraphIsEmptyExp;
+  pDdiTable->pfnSetDestructionCallbackExp = urGraphSetDestructionCallbackExp;
   pDdiTable->pfnDumpContentsExp = urGraphDumpContentsExp;
 
   return UR_RESULT_SUCCESS;

unified-runtime/source/adapters/level_zero/graph.cpp

@@ -44,6 +44,16 @@ ur_result_t urGraphIsEmptyExp(ur_exp_graph_handle_t /* hGraph */,
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+ur_result_t urGraphSetDestructionCallbackExp(
+    ur_exp_graph_handle_t /* hGraph */,
+    ur_exp_graph_destruction_callback_t /* pfnCallback */,
+    void * /* pUserData */) {
+  UR_LOG_LEGACY(ERR,
+                logger::LegacyMessage("[UR][L0] {} function not implemented!"),
+                "{} function not implemented!", __FUNCTION__);
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
 ur_result_t urGraphDumpContentsExp(ur_exp_graph_handle_t /* hGraph */,
                                    const char * /* pDotFilePath */) {
   UR_LOG_LEGACY(ERR,

unified-runtime/source/adapters/level_zero/platform.cpp

@@ -613,6 +613,9 @@ ur_result_t ur_platform_handle_t_::initialize() {
         {
             {"zeCommandListGetGraphExp",
              reinterpret_cast<void **>(&ZeGraphExt.zeCommandListGetGraphExp)},
+            {"zeGraphSetDestructionCallbackExp",
+             reinterpret_cast<void **>(
+                 &ZeGraphExt.zeGraphSetDestructionCallbackExp)},
         };
 
     for (auto &[funcName, funcAddr] : ZeGraphOptionalFuncNameToAddrMap) {

unified-runtime/source/adapters/level_zero/platform.hpp

@@ -203,6 +203,9 @@ struct ur_platform_handle_t_ : ur::handle_base<ur::level_zero::ddi_getter>,
     ze_result_t (*zeGraphIsEmptyExp)(ze_graph_handle_t hGraph);
     ze_result_t (*zeGraphDumpContentsExp)(ze_graph_handle_t hGraph,
                                           const char *filePath, void *pNext);
+    ze_result_t (*zeGraphSetDestructionCallbackExp)(
+        ze_graph_handle_t hGraph, zex_mem_graph_free_callback_fn_t pfnCallback,
+        void *pUserData, void *pNext);
   } ZeGraphExt;
 
   struct ZeHostTaskExtension {

unified-runtime/source/adapters/level_zero/v2/graph.cpp

@@ -13,6 +13,8 @@
 #include "common.hpp"
 #include "context.hpp"
 
+#include <memory>
+
 ur_exp_graph_handle_t_::ur_exp_graph_handle_t_(ur_context_handle_t hContext)
     : hContext(hContext) {
   ZE2UR_CALL_THROWS(hContext->getPlatform()->ZeGraphExt.zeGraphCreateExp,
@@ -53,6 +55,23 @@ ur_exp_executable_graph_handle_t_::~ur_exp_executable_graph_handle_t_() {
   }
 }
 
+namespace {
+
+// L0 imposes specific callback conventions. We must wrap the UR callback in
+// order to not violate this requirement.
+struct DestructionCallbackContext {
+  ur_exp_graph_destruction_callback_t callback;
+  void *userData;
+};
+
+void ZE_CALLBACK destructionCallbackWrapper(void *pUserData) {
+  auto *CbData = static_cast<DestructionCallbackContext *>(pUserData);
+  CbData->callback(CbData->userData);
+  delete CbData;
+}
+
+} // namespace
+
 namespace ur::level_zero {
 
 ur_result_t urGraphCreateExp(ur_context_handle_t hContext,
@@ -129,6 +148,32 @@ ur_result_t urGraphIsEmptyExp(ur_exp_graph_handle_t hGraph, bool *pIsEmpty) {
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t urGraphSetDestructionCallbackExp(
+    ur_exp_graph_handle_t hGraph,
+    ur_exp_graph_destruction_callback_t pfnCallback, void *pUserData) {
+  ur_context_handle_t hContext = hGraph->getContext();
+  auto ZeSetCallback =
+      hContext->getPlatform()->ZeGraphExt.zeGraphSetDestructionCallbackExp;
+  if (!checkGraphExtensionSupport(hContext) || !ZeSetCallback) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  auto CbData = std::make_unique<DestructionCallbackContext>(
+      DestructionCallbackContext{pfnCallback, pUserData});
+
+  ze_result_t ZeResult = ZE_CALL_NOCHECK(
+      ZeSetCallback, (hGraph->getZeHandle(), destructionCallbackWrapper,
+                      static_cast<void *>(CbData.get()), nullptr));
+
+  if (ZeResult != ZE_RESULT_SUCCESS) {
+    return ze2urResult(ZeResult);
+  }
+
+  // Ownership is transfered to the graph destruction callback for deletion
+  CbData.release();
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t urGraphDumpContentsExp(ur_exp_graph_handle_t hGraph,
                                    const char *filePath) {
   ur_context_handle_t hContext = hGraph->getContext();