[SYCL][Graph][UR] Propagate graph update list to UR (#17019)

Update the `urCommandBufferUpdateKernelLaunchExp` API for updating
commands in a command-buffer to take a list of commands.

The current API operates on a single command, this means that the
SYCL-Graph `update(std::vector<nodes>)` API needs to serialize the list
into N calls to the UR API. Given that both OpenCL
`clUpdateMutableCommandsKHR` and Level-Zero
`zeCommandListUpdateMutableCommandsExp` can operate on a list of
commands, this serialization at the UR layer of the stack introduces
extra host API calls.

This PR improves the `urCommandBufferUpdateKernelLaunchExp` API so that
a list of commands is passed all the way from SYCL to the native backend
API.

As highlighted in
oneapi-src/unified-runtime#2671 this patch
requires the handle translation auto generated code to be able to handle
a list of structs, which is not currently the case. This is PR includes
a API specific temporary workaround in the mako file which will unblock
this PR until a more permanent solution is completed.

---------

Co-authored-by: Ross Brunton <[email protected]>

Author: Ewan Crawford

sycl/doc/extensions/experimental/sycl_ext_oneapi_graph.asciidoc

@@ -1431,6 +1431,7 @@ Exceptions:
   created.
 * Throws with error code `invalid` if `node` is not part of the
   graph.
+* If any other exception is thrown the state of the graph node is undefined.
 
 |
 [source,c++]
@@ -1465,6 +1466,7 @@ Exceptions:
   `property::graph::updatable` was not set when the executable graph was created.
 * Throws with error code `invalid` if any node in `nodes` is not part of the
   graph.
+* If any other exception is thrown the state of the graph nodes is undefined.
 
 |
 [source, c++]
@@ -1517,6 +1519,8 @@ Exceptions:
 * Throws synchronously with error code `invalid` if
   `property::graph::updatable` was not set when the executable graph was
   created.
+
+* If any other exception is thrown the state of the graph nodes is undefined.
 |===
 
 Table {counter: tableNumber}. Member functions of the `command_graph` class for

sycl/source/detail/graph_impl.cpp

@@ -1303,7 +1303,30 @@ class exec_graph_impl {
   void update(std::shared_ptr<node_impl> Node);
   void update(const std::vector<std::shared_ptr<node_impl>> &Nodes);
 
-  void updateImpl(std::shared_ptr<node_impl> NodeImpl);
+  /// Calls UR entry-point to update nodes in command-buffer.
+  /// @param CommandBuffer The UR command-buffer to update commands in.
+  /// @param Nodes List of nodes to update. Only nodes which can be updated
+  /// through UR should be included in this list, currently this is only
+  /// nodes of kernel type.
+  void updateURImpl(ur_exp_command_buffer_handle_t CommandBuffer,
+                    const std::vector<std::shared_ptr<node_impl>> &Nodes) const;
+
+  /// Update host-task nodes
+  /// @param Nodes List of nodes to update, any node that is not a host-task
+  /// will be ignored.
+  void updateHostTasksImpl(
+      const std::vector<std::shared_ptr<node_impl>> &Nodes) const;
+
+  /// Splits a list of nodes into separate lists of nodes for each
+  /// command-buffer partition.
+  ///
+  /// Only nodes that can be updated through the UR interface are included
+  /// in the list. Currently this is only kernel node types.
+  ///
+  /// @param Nodes List of nodes to split
+  /// @return Map of partition indexes to nodes
+  std::map<int, std::vector<std::shared_ptr<node_impl>>> getURUpdatableNodes(
+      const std::vector<std::shared_ptr<node_impl>> &Nodes) const;
 
   unsigned long long getID() const { return MID; }
 
@@ -1373,6 +1396,34 @@ class exec_graph_impl {
     Stream.close();
   }
 
+  /// Determines if scheduler needs to be used for node update.
+  /// @param[in] Nodes List of nodes to be updated
+  /// @param[out] UpdateRequirements Accessor requirements found in /p Nodes.
+  /// return True if update should be done through the scheduler.
+  bool needsScheduledUpdate(
+      const std::vector<std::shared_ptr<node_impl>> &Nodes,
+      std::vector<sycl::detail::AccessorImplHost *> &UpdateRequirements);
+
+  /// Sets the UR struct values required to update a graph node.
+  /// @param[in] Node The node to be updated.
+  /// @param[out] BundleObjs UR objects created from kernel bundle.
+  /// Responsibility of the caller to release.
+  /// @param[out] MemobjDescs Memory object arguments to update.
+  /// @param[out] MemobjProps Properties used in /p MemobjDescs structs.
+  /// @param[out] PtrDescs Pointer arguments to update.
+  /// @param[out] ValueDescs Value arguments to update.
+  /// @param[out] NDRDesc ND-Range to update.
+  /// @param[out] UpdateDesc Base struct in the pointer chain.
+  void populateURKernelUpdateStructs(
+      const std::shared_ptr<node_impl> &Node,
+      std::pair<ur_program_handle_t, ur_kernel_handle_t> &BundleObjs,
+      std::vector<ur_exp_command_buffer_update_memobj_arg_desc_t> &MemobjDescs,
+      std::vector<ur_kernel_arg_mem_obj_properties_t> &MemobjProps,
+      std::vector<ur_exp_command_buffer_update_pointer_arg_desc_t> &PtrDescs,
+      std::vector<ur_exp_command_buffer_update_value_arg_desc_t> &ValueDescs,
+      sycl::detail::NDRDescT &NDRDesc,
+      ur_exp_command_buffer_update_kernel_launch_desc_t &UpdateDesc) const;
+
   /// Execution schedule of nodes in the graph.
   std::list<std::shared_ptr<node_impl>> MSchedule;
   /// Pointer to the modifiable graph impl associated with this executable

sycl/source/detail/graph_impl.hpp

@@ -3707,7 +3707,17 @@ ur_result_t UpdateCommandBufferCommand::enqueueImp() {
     default:
       break;
     }
-    MGraph->updateImpl(Node);
+  }
+
+  // Split list of nodes into nodes per UR command-buffer partition, then
+  // call UR update on each command-buffer partition with those updatable
+  // nodes.
+  auto PartitionedNodes = MGraph->getURUpdatableNodes(MNodes);
+  auto Device = MQueue->get_device();
+  auto &Partitions = MGraph->getPartitions();
+  for (auto &[PartitionIndex, NodeImpl] : PartitionedNodes) {
+    auto CommandBuffer = Partitions[PartitionIndex]->MCommandBuffers[Device];
+    MGraph->updateURImpl(CommandBuffer, NodeImpl);
   }
 
   return UR_RESULT_SUCCESS;

sycl/source/detail/scheduler/commands.cpp

@@ -10446,6 +10446,21 @@ typedef struct ur_exp_command_buffer_desc_t {
 
 } ur_exp_command_buffer_desc_t;
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief A value that identifies a command inside of a command-buffer, used
+/// for
+///        defining dependencies between commands in the same command-buffer.
+typedef uint32_t ur_exp_command_buffer_sync_point_t;
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Handle of Command-Buffer object
+typedef struct ur_exp_command_buffer_handle_t_ *ur_exp_command_buffer_handle_t;
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Handle of a Command-Buffer command
+typedef struct ur_exp_command_buffer_command_handle_t_
+    *ur_exp_command_buffer_command_handle_t;
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Descriptor type for updating a kernel command memobj argument.
 typedef struct ur_exp_command_buffer_update_memobj_arg_desc_t {
@@ -10509,6 +10524,8 @@ typedef struct ur_exp_command_buffer_update_kernel_launch_desc_t {
   ur_structure_type_t stype;
   /// [in][optional] pointer to extension-specific structure
   const void *pNext;
+  /// [in] Handle of the command-buffer kernel command to update.
+  ur_exp_command_buffer_command_handle_t hCommand;
   /// [in][optional] The new kernel handle. If this parameter is nullptr,
   /// the current kernel handle in `hCommand`
   /// will be used. If a kernel handle is passed, it must be a valid kernel
@@ -10558,21 +10575,6 @@ typedef struct ur_exp_command_buffer_update_kernel_launch_desc_t {
 
 } ur_exp_command_buffer_update_kernel_launch_desc_t;
 
-///////////////////////////////////////////////////////////////////////////////
-/// @brief A value that identifies a command inside of a command-buffer, used
-/// for
-///        defining dependencies between commands in the same command-buffer.
-typedef uint32_t ur_exp_command_buffer_sync_point_t;
-
-///////////////////////////////////////////////////////////////////////////////
-/// @brief Handle of Command-Buffer object
-typedef struct ur_exp_command_buffer_handle_t_ *ur_exp_command_buffer_handle_t;
-
-///////////////////////////////////////////////////////////////////////////////
-/// @brief Handle of a Command-Buffer command
-typedef struct ur_exp_command_buffer_command_handle_t_
-    *ur_exp_command_buffer_command_handle_t;
-
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Create a Command-Buffer object
 ///
@@ -11520,74 +11522,84 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferEnqueueExp(
 ///
 /// @details
 /// This entry-point is synchronous and may block if the command-buffer is
-/// executing when the entry-point is called.
+/// executing when the entry-point is called. On error, the state of the
+/// command-buffer commands being updated is undefined.
 ///
 /// @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`
+///         + `NULL == hCommandBuffer`
+///         + `NULL == pUpdateKernelLaunch->hCommand`
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
 ///         + `NULL == pUpdateKernelLaunch`
+///     - ::UR_RESULT_ERROR_INVALID_COMMAND_BUFFER_EXP
+///     - ::UR_RESULT_ERROR_INVALID_SIZE
+///         + `numKernelUpdates == 0`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_FEATURE
 ///         + If
 ///         ::UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_KERNEL_ARGUMENTS
-///         is not supported by the device, but any of
-///         `pUpdateKernelLaunch->numNewMemObjArgs`,
-///         `pUpdateKernelLaunch->numNewPointerArgs`, or
-///         `pUpdateKernelLaunch->numNewValueArgs` are not zero.
+///         is not supported by the device, and for any of any element of
+///         `pUpdateKernelLaunch` the `numNewMemObjArgs`, `numNewPointerArgs`,
+///         or `numNewValueArgs` members are not zero.
 ///         + If
 ///         ::UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_LOCAL_WORK_SIZE is
-///         not supported by the device but
-///         `pUpdateKernelLaunch->pNewLocalWorkSize` is not nullptr.
+///         not supported by the device, and for any element of
+///         `pUpdateKernelLaunch` the `pNewLocalWorkSize` member is not nullptr.
 ///         + If
 ///         ::UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_LOCAL_WORK_SIZE is
-///         not supported by the device but
-///         `pUpdateKernelLaunch->pNewLocalWorkSize` is nullptr and
-///         `pUpdateKernelLaunch->pNewGlobalWorkSize` is not nullptr.
+///         not supported by the device, and for any element of
+///         `pUpdateKernelLaunch` the `pNewLocalWorkSize` member is nullptr and
+///         `pNewGlobalWorkSize` is not nullptr.
 ///         + If
 ///         ::UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_GLOBAL_WORK_SIZE
-///         is not supported by the device but
-///         `pUpdateKernelLaunch->pNewGlobalWorkSize` is not nullptr
+///         is not supported by the device, and for any element of
+///         `pUpdateKernelLaunch` the `pNewGlobalWorkSize` member is not nullptr
 ///         + If
 ///         ::UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_GLOBAL_WORK_OFFSET
-///         is not supported by the device but
-///         `pUpdateKernelLaunch->pNewGlobalWorkOffset` is not nullptr.
+///         is not supported by the device, and for any element of
+///         `pUpdateKernelLaunch` the `pNewGlobalWorkOffset` member is not
+///         nullptr.
 ///         + If ::UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_KERNEL_HANDLE
-///         is not supported by the device but `pUpdateKernelLaunch->hNewKernel`
-///         is not nullptr.
+///         is not supported by the device, and for any element of
+///         `pUpdateKernelLaunch` the `hNewKernel` member is not nullptr.
 ///     - ::UR_RESULT_ERROR_INVALID_OPERATION
 ///         + If ::ur_exp_command_buffer_desc_t::isUpdatable was not set to true
-///         on creation of the command-buffer `hCommand` belongs to.
-///         + If the command-buffer `hCommand` belongs to has not been
-///         finalized.
+///         on creation of the `hCommandBuffer`.
+///         + If `hCommandBuffer`  has not been finalized.
 ///     - ::UR_RESULT_ERROR_INVALID_COMMAND_BUFFER_COMMAND_HANDLE_EXP
-///         + If `hCommand` is not a kernel execution command.
+///         + If for any element of `pUpdateKernelLaunch` the `hCommand` member
+///         is not a kernel execution command.
+///         + If for any element of `pUpdateKernelLaunch` the `hCommand` member
+///         was not created from `hCommandBuffer`.
 ///     - ::UR_RESULT_ERROR_INVALID_MEM_OBJECT
 ///     - ::UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX
 ///     - ::UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_SIZE
 ///     - ::UR_RESULT_ERROR_INVALID_ENUMERATION
 ///     - ::UR_RESULT_ERROR_INVALID_WORK_DIMENSION
-///         + `pUpdateKernelLaunch->newWorkDim < 1 ||
-///         pUpdateKernelLaunch->newWorkDim > 3`
+///         + If for any element of `pUpdateKernelLaunch` the `newWorkDim`
+///         member is less than 1 or greater than 3.
 ///     - ::UR_RESULT_ERROR_INVALID_WORK_GROUP_SIZE
 ///     - ::UR_RESULT_ERROR_INVALID_VALUE
-///         + If `pUpdateKernelLaunch->hNewKernel` was not passed to the
-///         `hKernel` or `phKernelAlternatives` parameters of
-///         ::urCommandBufferAppendKernelLaunchExp when this command was
-///         created.
-///         + If `pUpdateKernelLaunch->newWorkDim` is different from the current
-///         workDim in `hCommand` and,
-///         `pUpdateKernelLaunch->pNewGlobalWorkSize`, or
-///         `pUpdateKernelLaunch->pNewGlobalWorkOffset` are nullptr.
+///         + If for any element of `pUpdateKernelLaunch` the `hNewKernel`
+///         member was not passed to the `hKernel` or `phKernelAlternatives`
+///         parameters of ::urCommandBufferAppendKernelLaunchExp when the
+///         command was created.
+///         + If for any element of `pUpdateKernelLaunch` the `newWorkDim`
+///         member is different from the current workDim in the `hCommand`
+///         member, and `pNewGlobalWorkSize` or `pNewGlobalWorkOffset` are
+///         nullptr.
 ///     - ::UR_RESULT_ERROR_OUT_OF_HOST_MEMORY
 ///     - ::UR_RESULT_ERROR_OUT_OF_RESOURCES
 UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferUpdateKernelLaunchExp(
-    /// [in] Handle of the command-buffer kernel command to update.
-    ur_exp_command_buffer_command_handle_t hCommand,
-    /// [in] Struct defining how the kernel command is to be updated.
+    /// [in] Handle of the command-buffer object.
+    ur_exp_command_buffer_handle_t hCommandBuffer,
+    /// [in] Length of pUpdateKernelLaunch.
+    uint32_t numKernelUpdates,
+    /// [in][range(0, numKernelUpdates)]  List of structs defining how a
+    /// kernel commands are to be updated.
     const ur_exp_command_buffer_update_kernel_launch_desc_t
         *pUpdateKernelLaunch);
 
@@ -14820,7 +14832,8 @@ typedef struct ur_command_buffer_enqueue_exp_params_t {
 /// @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_update_kernel_launch_exp_params_t {
-  ur_exp_command_buffer_command_handle_t *phCommand;
+  ur_exp_command_buffer_handle_t *phCommandBuffer;
+  uint32_t *pnumKernelUpdates;
   const ur_exp_command_buffer_update_kernel_launch_desc_t *
       *ppUpdateKernelLaunch;
 } ur_command_buffer_update_kernel_launch_exp_params_t;

unified-runtime/include/ur_api.h

@@ -1682,7 +1682,7 @@ typedef ur_result_t(UR_APICALL *ur_pfnCommandBufferEnqueueExp_t)(
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function-pointer for urCommandBufferUpdateKernelLaunchExp
 typedef ur_result_t(UR_APICALL *ur_pfnCommandBufferUpdateKernelLaunchExp_t)(
-    ur_exp_command_buffer_command_handle_t,
+    ur_exp_command_buffer_handle_t, uint32_t,
     const ur_exp_command_buffer_update_kernel_launch_desc_t *);
 
 ///////////////////////////////////////////////////////////////////////////////

unified-runtime/include/ur_ddi.h

@@ -11481,6 +11481,11 @@ inline std::ostream &operator<<(
 
   ur::details::printStruct(os, (params.pNext));
 
+  os << ", ";
+  os << ".hCommand = ";
+
+  ur::details::printPtr(os, (params.hCommand));
+
   os << ", ";
   os << ".hNewKernel = ";
 
@@ -19432,14 +19437,30 @@ inline std::ostream &
 operator<<(std::ostream &os, [[maybe_unused]] const struct
            ur_command_buffer_update_kernel_launch_exp_params_t *params) {
 
-  os << ".hCommand = ";
+  os << ".hCommandBuffer = ";
 
-  ur::details::printPtr(os, *(params->phCommand));
+  ur::details::printPtr(os, *(params->phCommandBuffer));
+
+  os << ", ";
+  os << ".numKernelUpdates = ";
+
+  os << *(params->pnumKernelUpdates);
 
   os << ", ";
   os << ".pUpdateKernelLaunch = ";
+  ur::details::printPtr(
+      os, reinterpret_cast<const void *>(*(params->ppUpdateKernelLaunch)));
+  if (*(params->ppUpdateKernelLaunch) != NULL) {
+    os << " {";
+    for (size_t i = 0; i < *params->pnumKernelUpdates; ++i) {
+      if (i != 0) {
+        os << ", ";
+      }
 
-  ur::details::printPtr(os, *(params->ppUpdateKernelLaunch));
+      os << (*(params->ppUpdateKernelLaunch))[i];
+    }
+    os << "}";
+  }
 
   return os;
 }

unified-runtime/include/ur_print.hpp

@@ -309,7 +309,8 @@ ${x}CommandBufferUpdateKernelLaunchExp.
     ${x}_exp_command_buffer_update_kernel_launch_desc_t update {
         UR_STRUCTURE_TYPE_EXP_COMMAND_BUFFER_UPDATE_KERNEL_LAUNCH_DESC, // stype
         nullptr, // pNext
-        hNewKernel  // hNewKernel
+        hCommand, // hCommand
+        hNewKernel,  // hNewKernel
         2, // numNewMemobjArgs
         0, // numNewPointerArgs
         0, // numNewValueArgs
@@ -325,7 +326,7 @@ ${x}CommandBufferUpdateKernelLaunchExp.
     };
 
     // Perform the update
-    ${x}CommandBufferUpdateKernelLaunchExp(hCommand, &update);
+    ${x}CommandBufferUpdateKernelLaunchExp(hCommandBuffer, 1, &update);
 
 Command Event Update
 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
@@ -514,6 +515,8 @@ Changelog
 +-----------+-------------------------------------------------------+
 | 1.7       | Remove command handle reference counting and querying |
 +-----------+-------------------------------------------------------+
+| 1.8       | Change Kernel command update API to take a list       |
++-----------+-------------------------------------------------------+
 
 Contributors
 --------------------------------------------------------------------------------