Reference count command handles

include/ur_api.h

@@ -216,6 +216,8 @@ typedef enum ur_function_t {
     UR_FUNCTION_COMMAND_BUFFER_APPEND_USM_ADVISE_EXP = 213,                    ///< Enumerator for ::urCommandBufferAppendUSMAdviseExp
     UR_FUNCTION_ENQUEUE_COOPERATIVE_KERNEL_LAUNCH_EXP = 214,                   ///< Enumerator for ::urEnqueueCooperativeKernelLaunchExp
     UR_FUNCTION_KERNEL_SUGGEST_MAX_COOPERATIVE_GROUP_COUNT_EXP = 215,          ///< Enumerator for ::urKernelSuggestMaxCooperativeGroupCountExp
+    UR_FUNCTION_COMMAND_BUFFER_RETAIN_COMMAND_EXP = 216,                       ///< Enumerator for ::urCommandBufferRetainCommandExp
+    UR_FUNCTION_COMMAND_BUFFER_RELEASE_COMMAND_EXP = 217,                      ///< Enumerator for ::urCommandBufferReleaseCommandExp
     /// @cond
     UR_FUNCTION_FORCE_UINT32 = 0x7fffffff
     /// @endcond
@@ -8435,6 +8437,43 @@ urCommandBufferEnqueueExp(
                                                    ///< command-buffer execution instance.
 );
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Increment the command object's reference count.
+///
+/// @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 == hCommand`
+///     - ::UR_RESULT_ERROR_INVALID_COMMAND_BUFFER_COMMAND_HANDLE_EXP
+///     - ::UR_RESULT_ERROR_OUT_OF_RESOURCES
+///     - ::UR_RESULT_ERROR_OUT_OF_HOST_MEMORY
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferRetainCommandExp(
+    ur_exp_command_buffer_command_handle_t hCommand ///< [in] Handle of the command-buffer command.
+);
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Decrement the command object's reference count and delete the command
+///        object if the reference count becomes zero.
+///
+/// @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 == hCommand`
+///     - ::UR_RESULT_ERROR_INVALID_COMMAND_BUFFER_COMMAND_HANDLE_EXP
+///     - ::UR_RESULT_ERROR_OUT_OF_RESOURCES
+///     - ::UR_RESULT_ERROR_OUT_OF_HOST_MEMORY
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferReleaseCommandExp(
+    ur_exp_command_buffer_command_handle_t hCommand ///< [in] Handle of the command-buffer command.
+);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Update a kernel launch command.
 ///
@@ -10819,6 +10858,22 @@ typedef struct ur_command_buffer_enqueue_exp_params_t {
     ur_event_handle_t **pphEvent;
 } ur_command_buffer_enqueue_exp_params_t;
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function parameters for urCommandBufferRetainCommandExp
+/// @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_command_buffer_retain_command_exp_params_t {
+    ur_exp_command_buffer_command_handle_t *phCommand;
+} ur_command_buffer_retain_command_exp_params_t;
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function parameters for urCommandBufferReleaseCommandExp
+/// @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_command_buffer_release_command_exp_params_t {
+    ur_exp_command_buffer_command_handle_t *phCommand;
+} ur_command_buffer_release_command_exp_params_t;
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function parameters for urCommandBufferUpdateKernelLaunchExp
 /// @details Each entry is a pointer to the parameter passed to the function;

include/ur_ddi.h

@@ -2011,6 +2011,16 @@ typedef ur_result_t(UR_APICALL *ur_pfnCommandBufferEnqueueExp_t)(
     const ur_event_handle_t *,
     ur_event_handle_t *);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function-pointer for urCommandBufferRetainCommandExp
+typedef ur_result_t(UR_APICALL *ur_pfnCommandBufferRetainCommandExp_t)(
+    ur_exp_command_buffer_command_handle_t);
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function-pointer for urCommandBufferReleaseCommandExp
+typedef ur_result_t(UR_APICALL *ur_pfnCommandBufferReleaseCommandExp_t)(
+    ur_exp_command_buffer_command_handle_t);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function-pointer for urCommandBufferUpdateKernelLaunchExp
 typedef ur_result_t(UR_APICALL *ur_pfnCommandBufferUpdateKernelLaunchExp_t)(
@@ -2037,6 +2047,8 @@ typedef struct ur_command_buffer_exp_dditable_t {
     ur_pfnCommandBufferAppendUSMPrefetchExp_t pfnAppendUSMPrefetchExp;
     ur_pfnCommandBufferAppendUSMAdviseExp_t pfnAppendUSMAdviseExp;
     ur_pfnCommandBufferEnqueueExp_t pfnEnqueueExp;
+    ur_pfnCommandBufferRetainCommandExp_t pfnRetainCommandExp;
+    ur_pfnCommandBufferReleaseCommandExp_t pfnReleaseCommandExp;
     ur_pfnCommandBufferUpdateKernelLaunchExp_t pfnUpdateKernelLaunchExp;
 } ur_command_buffer_exp_dditable_t;
 

include/ur_print.hpp

@@ -889,6 +889,12 @@ inline std::ostream &operator<<(std::ostream &os, ur_function_t value) {
     case UR_FUNCTION_KERNEL_SUGGEST_MAX_COOPERATIVE_GROUP_COUNT_EXP:
         os << "UR_FUNCTION_KERNEL_SUGGEST_MAX_COOPERATIVE_GROUP_COUNT_EXP";
         break;
+    case UR_FUNCTION_COMMAND_BUFFER_RETAIN_COMMAND_EXP:
+        os << "UR_FUNCTION_COMMAND_BUFFER_RETAIN_COMMAND_EXP";
+        break;
+    case UR_FUNCTION_COMMAND_BUFFER_RELEASE_COMMAND_EXP:
+        os << "UR_FUNCTION_COMMAND_BUFFER_RELEASE_COMMAND_EXP";
+        break;
     default:
         os << "unknown enumerator";
         break;
@@ -15595,6 +15601,34 @@ inline std::ostream &operator<<(std::ostream &os, [[maybe_unused]] const struct
     return os;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Print operator for the ur_command_buffer_retain_command_exp_params_t type
+/// @returns
+///     std::ostream &
+inline std::ostream &operator<<(std::ostream &os, [[maybe_unused]] const struct ur_command_buffer_retain_command_exp_params_t *params) {
+
+    os << ".hCommand = ";
+
+    ur::details::printPtr(os,
+                          *(params->phCommand));
+
+    return os;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Print operator for the ur_command_buffer_release_command_exp_params_t type
+/// @returns
+///     std::ostream &
+inline std::ostream &operator<<(std::ostream &os, [[maybe_unused]] const struct ur_command_buffer_release_command_exp_params_t *params) {
+
+    os << ".hCommand = ";
+
+    ur::details::printPtr(os,
+                          *(params->phCommand));
+
+    return os;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Print operator for the ur_command_buffer_update_kernel_launch_exp_params_t type
 /// @returns
@@ -16749,6 +16783,12 @@ inline ur_result_t UR_APICALL printFunctionParams(std::ostream &os, ur_function_
     case UR_FUNCTION_COMMAND_BUFFER_ENQUEUE_EXP: {
         os << (const struct ur_command_buffer_enqueue_exp_params_t *)params;
     } break;
+    case UR_FUNCTION_COMMAND_BUFFER_RETAIN_COMMAND_EXP: {
+        os << (const struct ur_command_buffer_retain_command_exp_params_t *)params;
+    } break;
+    case UR_FUNCTION_COMMAND_BUFFER_RELEASE_COMMAND_EXP: {
+        os << (const struct ur_command_buffer_release_command_exp_params_t *)params;
+    } break;
     case UR_FUNCTION_COMMAND_BUFFER_UPDATE_KERNEL_LAUNCH_EXP: {
         os << (const struct ur_command_buffer_update_kernel_launch_exp_params_t *)params;
     } break;

scripts/core/EXP-COMMAND-BUFFER.rst

@@ -95,9 +95,10 @@ queue handle, and the dependencies and return parameters are sync-points instead
 of event handles.
 
 The entry-point for appending a kernel launch command also returns an optional
-handle to the command being appended. Returning this handle does not extend the
-lifetime of the parent command-buffer, and using the handle after the
-command-buffer has been destroyed is invalid behaviour.
+handle to the command being appended. Returning this handle extends the
+lifetime of the parent command-buffer by incrementing its reference count, and
+the parent command-buffer's reference count is decremented again when the command
+handle is destroyed.
 
 Currently only the following commands are supported:
 
@@ -195,17 +196,17 @@ parameters to the kernel and the execution ND-Range.
     // Define kernel argument at index 0 to be a new input buffer object
     ${x}_exp_command_buffer_update_memobj_arg_desc_t newInputArg {
         UR_STRUCTURE_TYPE_EXP_COMMAND_BUFFER_UPDATE_MEMOBJ_ARG_DESC, // stype
-        nullptr, // pNext,
-        0, // argIndex,
+        nullptr, // pNext
+        0, // argIndex
         nullptr, // pProperties
         newInputBuffer, // hArgValue
     };
 
     // Define kernel argument at index 1 to be a new output buffer object
     ${x}_exp_command_buffer_update_memobj_arg_desc_t newOutputArg {
         UR_STRUCTURE_TYPE_EXP_COMMAND_BUFFER_UPDATE_MEMOBJ_ARG_DESC, // stype
-        nullptr, // pNext,
-        1, // argIndex,
+        nullptr, // pNext
+        1, // argIndex
         nullptr, // pProperties
         newOutputBuffer, // hArgValue
     };
@@ -219,7 +220,7 @@ parameters to the kernel and the execution ND-Range.
         0, // numPointerArgs
         0, // numValueArgs
         0, // numExecInfos
-        0, // workDim;
+        0, // workDim
         new_args, // pArgMemobjList
         nullptr, // pArgPointerList
         nullptr, // pArgValueList
@@ -310,6 +311,8 @@ Functions
 * ${x}CommandBufferAppendUSMPrefetchExp
 * ${x}CommandBufferAppendUSMAdviseExp
 * ${x}CommandBufferEnqueueExp
+* ${x}CommandBufferRetainCommandExp
+* ${x}CommandBufferReleaseCommandExp
 * ${x}CommandBufferUpdateKernelLaunchExp
 
 Changelog

scripts/core/exp-command-buffer.yml

@@ -863,6 +863,32 @@ returns:
     - $X_RESULT_ERROR_OUT_OF_RESOURCES
 --- #--------------------------------------------------------------------------
 type: function
+desc: "Increment the command object's reference count."
+class: $xCommandBuffer
+name: RetainCommandExp
+params:
+    - type: $x_exp_command_buffer_command_handle_t
+      name: hCommand
+      desc: "[in] Handle of the command-buffer command."
+returns:
+    - $X_RESULT_ERROR_INVALID_COMMAND_BUFFER_COMMAND_HANDLE_EXP
+    - $X_RESULT_ERROR_OUT_OF_RESOURCES
+    - $X_RESULT_ERROR_OUT_OF_HOST_MEMORY
+--- #--------------------------------------------------------------------------
+type: function
+desc: "Decrement the command object's reference count and delete the command object if the reference count becomes zero."
+class: $xCommandBuffer
+name: ReleaseCommandExp
+params:
+    - type: $x_exp_command_buffer_command_handle_t
+      name: hCommand
+      desc: "[in] Handle of the command-buffer command."
+returns:
+    - $X_RESULT_ERROR_INVALID_COMMAND_BUFFER_COMMAND_HANDLE_EXP
+    - $X_RESULT_ERROR_OUT_OF_RESOURCES
+    - $X_RESULT_ERROR_OUT_OF_HOST_MEMORY
+--- #--------------------------------------------------------------------------
+type: function
 desc: "Update a kernel launch command."
 class: $xCommandBuffer
 name: UpdateKernelLaunchExp

scripts/core/registry.yml

@@ -562,6 +562,12 @@ etors:
 - name: KERNEL_SUGGEST_MAX_COOPERATIVE_GROUP_COUNT_EXP
   desc: Enumerator for $xKernelSuggestMaxCooperativeGroupCountExp
   value: '215'
+- name: COMMAND_BUFFER_RETAIN_COMMAND_EXP
+  desc: Enumerator for $xCommandBufferRetainCommandExp
+  value: '216'
+- name: COMMAND_BUFFER_RELEASE_COMMAND_EXP
+  desc: Enumerator for $xCommandBufferReleaseCommandExp
+  value: '217'
 ---
 type: enum
 desc: Defines structure types

source/adapters/cuda/command_buffer.cpp

@@ -772,6 +772,21 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferEnqueueExp(
   return Result;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferRetainCommandExp(
+    ur_exp_command_buffer_command_handle_t hCommand) {
+  hCommand->incrementReferenceCount();
+  return UR_RESULT_SUCCESS;
+}
+
+UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferReleaseCommandExp(
+    ur_exp_command_buffer_command_handle_t hCommand) {
+  if (hCommand->decrementReferenceCount() != 0)
+    return UR_RESULT_SUCCESS;
+
+  delete hCommand;
+  return UR_RESULT_SUCCESS;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferUpdateKernelLaunchExp(
     ur_exp_command_buffer_command_handle_t hCommand,
     const ur_exp_command_buffer_update_kernel_launch_desc_t *pKernelLaunch) {

source/adapters/cuda/command_buffer.hpp

@@ -187,7 +187,9 @@ struct ur_exp_command_buffer_command_handle_t_ {
       uint32_t WorkDim, const size_t *pGlobalWorkOffset,
       const size_t *pGlobalWorkSize, const size_t *pLocalWorkSize)
       : CommandBuffer(CommandBuffer), Kernel(Kernel), Node(Node),
-        Params(Params), WorkDim(WorkDim) {
+        Params(Params), WorkDim(WorkDim), RefCount(1) {
+    UR_TRACE(urCommandBufferRetainExp(CommandBuffer));
+
     const size_t CopySize = sizeof(size_t) * WorkDim;
     std::memcpy(GlobalWorkOffset, pGlobalWorkOffset, CopySize);
     std::memcpy(GlobalWorkSize, pGlobalWorkSize, CopySize);
@@ -201,6 +203,10 @@ struct ur_exp_command_buffer_command_handle_t_ {
     }
   }
 
+  ~ur_exp_command_buffer_command_handle_t_() {
+    UR_TRACE(urCommandBufferReleaseExp(CommandBuffer));
+  }
+
   void SetGlobalOffset(const size_t *pGlobalWorkOffset) {
     const size_t CopySize = sizeof(size_t) * WorkDim;
     std::memcpy(GlobalWorkOffset, pGlobalWorkOffset, CopySize);
@@ -228,6 +234,12 @@ struct ur_exp_command_buffer_command_handle_t_ {
     }
   }
 
+  // Used when retaining an object.
+  uint32_t incrementReferenceCount() noexcept { return ++RefCount; }
+  // Used when releasing an object.
+  uint32_t decrementReferenceCount() noexcept { return --RefCount; }
+  uint32_t getReferenceCount() const noexcept { return RefCount; }
+
   ur_exp_command_buffer_handle_t CommandBuffer;
   ur_kernel_handle_t Kernel;
   std::shared_ptr<CUgraphNode> Node;
@@ -237,6 +249,8 @@ struct ur_exp_command_buffer_command_handle_t_ {
   size_t GlobalWorkOffset[3];
   size_t GlobalWorkSize[3];
   size_t LocalWorkSize[3];
+
+  std::atomic_uint32_t RefCount;
 };
 
 struct ur_exp_command_buffer_handle_t_ {

source/adapters/hip/command_buffer.cpp

@@ -164,6 +164,20 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferEnqueueExp(
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferRetainCommandExp(ur_exp_command_buffer_command_handle_t) {
+  detail::ur::die("Experimental Command-buffer feature is not "
+                  "implemented for the HIP adapter.");
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferReleaseCommandExp(ur_exp_command_buffer_command_handle_t) {
+  detail::ur::die("Experimental Command-buffer feature is not "
+                  "implemented for the HIP adapter.");
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferUpdateKernelLaunchExp(
     ur_exp_command_buffer_command_handle_t,
     const ur_exp_command_buffer_update_kernel_launch_desc_t *) {

source/adapters/level_zero/command_buffer.cpp

@@ -959,6 +959,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferEnqueueExp(
   return UR_RESULT_SUCCESS;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferRetainCommandExp(ur_exp_command_buffer_command_handle_t) {
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferReleaseCommandExp(ur_exp_command_buffer_command_handle_t) {
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferUpdateKernelLaunchExp(
     ur_exp_command_buffer_command_handle_t,
     const ur_exp_command_buffer_update_kernel_launch_desc_t *) {

source/adapters/native_cpu/command_buffer.cpp

@@ -135,6 +135,20 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferEnqueueExp(
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferRetainCommandExp(ur_exp_command_buffer_command_handle_t) {
+  detail::ur::die("Experimental Command-buffer feature is not "
+                  "implemented for the NativeCPU adapter.");
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
+UR_APIEXPORT ur_result_t UR_APICALL
+urCommandBufferReleaseCommandExp(ur_exp_command_buffer_command_handle_t) {
+  detail::ur::die("Experimental Command-buffer feature is not "
+                  "implemented for the NativeCPU adapter.");
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferUpdateKernelLaunchExp(
     ur_exp_command_buffer_command_handle_t,
     const ur_exp_command_buffer_update_kernel_launch_desc_t *) {