[UR] Add spec and implementation for urQueueGetGraphExp (#21817)

Introduces UR graph spec and implementation for `urQueueGetGraphExp`
which returns graph handles from a recording queue. The following
changes are made:

- A map of L0 to UR graph handles is made to fetch UR graph handles from
across command lists. This allows UR to fetch the correct graph handle
if created by a separate command list manager due to fork-join. Map is
stored in the context and protected by a lock.
- Map entries are added and removed in graph creation and destruction
respectively.
- In implicit recording mode (where the user makes no `urGraphCreateExp`
call), graph handle creation is done in the first instance of
`urQueueGetGraphExp` (if it occurs) instead of `urQueueEndCaptureExp` as
the user must be returned a ur_exp_graph_handle_t.
- L0 experimental headers are bumped to
https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4
where `zeCommandListGetGraphExp` is supported.

Author: mmichel11

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

@@ -509,6 +509,8 @@ typedef enum ur_function_t {
   UR_FUNCTION_USM_HOST_ALLOC_REGISTER_EXP = 312,
   /// Enumerator for ::urUSMHostAllocUnregisterExp
   UR_FUNCTION_USM_HOST_ALLOC_UNREGISTER_EXP = 313,
+  /// Enumerator for ::urQueueGetGraphExp
+  UR_FUNCTION_QUEUE_GET_GRAPH_EXP = 314,
   /// @cond
   UR_FUNCTION_FORCE_UINT32 = 0x7fffffff
   /// @endcond
@@ -13807,6 +13809,27 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueIsGraphCaptureEnabledExp(
     /// [out] Pointer to a boolean where the result will be stored.
     bool *pResult);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Get the graph handle currently being captured on the specified queue.
+///
+/// @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 == hQueue`
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+///         + `NULL == phGraph`
+///     - ::UR_RESULT_ERROR_INVALID_OPERATION
+///         + The queue is not in graph capture mode.
+UR_APIEXPORT ur_result_t UR_APICALL urQueueGetGraphExp(
+    /// [in] Handle of the queue to query.
+    ur_queue_handle_t hQueue,
+    /// [out] Pointer to the handle of the graph being captured. Set to
+    /// nullptr if queue is not in capture mode.
+    ur_exp_graph_handle_t *phGraph);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Return whether the given recorded graph contains any nodes.
 ///
@@ -14657,6 +14680,15 @@ typedef struct ur_queue_is_graph_capture_enabled_exp_params_t {
   bool **ppResult;
 } ur_queue_is_graph_capture_enabled_exp_params_t;
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function parameters for urQueueGetGraphExp
+/// @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_queue_get_graph_exp_params_t {
+  ur_queue_handle_t *phQueue;
+  ur_exp_graph_handle_t **pphGraph;
+} ur_queue_get_graph_exp_params_t;
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function parameters for urSamplerCreate
 /// @details Each entry is a pointer to the parameter passed to the function;

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

@@ -88,6 +88,7 @@ _UR_API(urQueueBeginGraphCaptureExp)
 _UR_API(urQueueBeginCaptureIntoGraphExp)
 _UR_API(urQueueEndGraphCaptureExp)
 _UR_API(urQueueIsGraphCaptureEnabledExp)
+_UR_API(urQueueGetGraphExp)
 _UR_API(urSamplerCreate)
 _UR_API(urSamplerRetain)
 _UR_API(urSamplerRelease)

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

@@ -677,13 +677,19 @@ typedef ur_result_t(UR_APICALL *ur_pfnQueueEndGraphCaptureExp_t)(
 typedef ur_result_t(UR_APICALL *ur_pfnQueueIsGraphCaptureEnabledExp_t)(
     ur_queue_handle_t, bool *);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function-pointer for urQueueGetGraphExp
+typedef ur_result_t(UR_APICALL *ur_pfnQueueGetGraphExp_t)(
+    ur_queue_handle_t, ur_exp_graph_handle_t *);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Table of QueueExp functions pointers
 typedef struct ur_queue_exp_dditable_t {
   ur_pfnQueueBeginGraphCaptureExp_t pfnBeginGraphCaptureExp;
   ur_pfnQueueBeginCaptureIntoGraphExp_t pfnBeginCaptureIntoGraphExp;
   ur_pfnQueueEndGraphCaptureExp_t pfnEndGraphCaptureExp;
   ur_pfnQueueIsGraphCaptureEnabledExp_t pfnIsGraphCaptureEnabledExp;
+  ur_pfnQueueGetGraphExp_t pfnGetGraphExp;
 } ur_queue_exp_dditable_t;
 
 ///////////////////////////////////////////////////////////////////////////////

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

@@ -2354,6 +2354,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urPrintQueueIsGraphCaptureEnabledExpParams(
     const struct ur_queue_is_graph_capture_enabled_exp_params_t *params,
     char *buffer, const size_t buff_size, size_t *out_size);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Print ur_queue_get_graph_exp_params_t struct
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_SIZE
+///         - `buff_size < out_size`
+UR_APIEXPORT ur_result_t UR_APICALL urPrintQueueGetGraphExpParams(
+    const struct ur_queue_get_graph_exp_params_t *params, char *buffer,
+    const size_t buff_size, size_t *out_size);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Print ur_sampler_create_params_t struct
 /// @returns

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

@@ -1375,6 +1375,9 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_function_t value) {
   case UR_FUNCTION_USM_HOST_ALLOC_UNREGISTER_EXP:
     os << "UR_FUNCTION_USM_HOST_ALLOC_UNREGISTER_EXP";
     break;
+  case UR_FUNCTION_QUEUE_GET_GRAPH_EXP:
+    os << "UR_FUNCTION_QUEUE_GET_GRAPH_EXP";
+    break;
   default:
     os << "unknown enumerator";
     break;
@@ -15490,6 +15493,26 @@ inline std::ostream &operator<<(
   return os;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Print operator for the ur_queue_get_graph_exp_params_t type
+/// @returns
+///     std::ostream &
+inline std::ostream &operator<<(
+    std::ostream &os,
+    [[maybe_unused]] const struct ur_queue_get_graph_exp_params_t *params) {
+
+  os << ".hQueue = ";
+
+  ur::details::printPtr(os, *(params->phQueue));
+
+  os << ", ";
+  os << ".phGraph = ";
+
+  ur::details::printPtr(os, *(params->pphGraph));
+
+  return os;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Print operator for the ur_sampler_create_params_t type
 /// @returns
@@ -22656,6 +22679,9 @@ inline ur_result_t UR_APICALL printFunctionParams(std::ostream &os,
   case UR_FUNCTION_QUEUE_IS_GRAPH_CAPTURE_ENABLED_EXP: {
     os << (const struct ur_queue_is_graph_capture_enabled_exp_params_t *)params;
   } break;
+  case UR_FUNCTION_QUEUE_GET_GRAPH_EXP: {
+    os << (const struct ur_queue_get_graph_exp_params_t *)params;
+  } break;
   case UR_FUNCTION_SAMPLER_CREATE: {
     os << (const struct ur_sampler_create_params_t *)params;
   } break;

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

@@ -49,6 +49,7 @@ Functions
    * ${x}QueueBeginCaptureIntoGraphExp
    * ${x}QueueEndGraphCaptureExp
    * ${x}QueueIsGraphCaptureEnabledExp
+   * ${x}QueueGetGraphExp
    * ${x}EnqueueGraphExp
 
 Changelog
@@ -68,6 +69,9 @@ Changelog
 |           | Rename QueueAppendGraphExp into             |
 |           | EnqueueGraphExp.                            |
 +-----------+---------------------------------------------+
+| 1.3       | Add ${x}QueueGetGraphExp to retrieve graph  |
+|           | handle from queue in capture mode.          |
++-----------+---------------------------------------------+
 
 Support
 --------------------------------------------------------------------------------

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

@@ -160,6 +160,21 @@ params:
       desc: "[out] Pointer to a boolean where the result will be stored."
 --- #--------------------------------------------------------------------------
 type: function
+desc: "Get the graph handle currently being captured on the specified queue."
+class: $xQueue
+name: GetGraphExp
+params:
+    - type: $x_queue_handle_t
+      name: hQueue
+      desc: "[in] Handle of the queue to query."
+    - type: $x_exp_graph_handle_t*
+      name: phGraph
+      desc: "[out] Pointer to the handle of the graph being captured. Set to nullptr if queue is not in capture mode."
+returns:
+    - $X_RESULT_ERROR_INVALID_OPERATION:
+        - "The queue is not in graph capture mode."
+--- #--------------------------------------------------------------------------
+type: function
 desc: "Return whether the given recorded graph contains any nodes."
 class: $xGraph
 name: IsEmptyExp

unified-runtime/scripts/core/registry.yml

@@ -721,7 +721,10 @@ etors:
 - name: USM_HOST_ALLOC_UNREGISTER_EXP
   desc: Enumerator for $xUSMHostAllocUnregisterExp
   value: '313'
-max_id: '313'
+- name: QUEUE_GET_GRAPH_EXP
+  desc: Enumerator for $xQueueGetGraphExp
+  value: '314'
+max_id: '314'
 ---
 type: enum
 desc: Defines structure types

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

@@ -289,6 +289,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueIsGraphCaptureEnabledExp(
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urQueueGetGraphExp(
+    ur_queue_handle_t /* hQueue */, ur_exp_graph_handle_t * /* phGraph */) {
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urEnqueueHostTaskExp(
     ur_queue_handle_t /* hQueue */,
     ur_exp_host_task_function_t /* pfnHostTask */, void * /* data */,

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

@@ -256,6 +256,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urGetQueueExpProcAddrTable(
   pDdiTable->pfnBeginCaptureIntoGraphExp = urQueueBeginCaptureIntoGraphExp;
   pDdiTable->pfnEndGraphCaptureExp = urQueueEndGraphCaptureExp;
   pDdiTable->pfnIsGraphCaptureEnabledExp = urQueueIsGraphCaptureEnabledExp;
+  pDdiTable->pfnGetGraphExp = urQueueGetGraphExp;
 
   return UR_RESULT_SUCCESS;
 }