[SYCL][Graph] Implementation of explicit update with indices (#12840)

- Experimental implementation of explicit update with indices (based on
API in #12486 )
- New scheduler command for updating a command buffer command
- PI equivalents for new UR APIs
- E2E and Unit Tests

---------

Co-authored-by: Ewan Crawford <[email protected]>

llvm/include/llvm/SYCLLowerIR/DeviceConfigFile.td

@@ -71,6 +71,7 @@ def AspectExt_oneapi_is_composite : Aspect<"ext_oneapi_is_composite">;
 def AspectExt_oneapi_is_component : Aspect<"ext_oneapi_is_component">;
 def AspectExt_oneapi_graph : Aspect<"ext_oneapi_graph">;
 def AspectExt_intel_fpga_task_sequence : Aspect<"ext_intel_fpga_task_sequence">;
+def AspectExt_oneapi_limited_graph : Aspect<"ext_oneapi_limited_graph">;
 // Deprecated aspects
 def AspectInt64_base_atomics : Aspect<"int64_base_atomics">;
 def AspectInt64_extended_atomics : Aspect<"int64_extended_atomics">;
@@ -122,7 +123,7 @@ def : TargetInfo<"__TestAspectList",
     AspectExt_oneapi_mipmap, AspectExt_oneapi_mipmap_anisotropy, AspectExt_oneapi_mipmap_level_reference, AspectExt_intel_esimd,
     AspectExt_oneapi_ballot_group, AspectExt_oneapi_fixed_size_group, AspectExt_oneapi_opportunistic_group,
     AspectExt_oneapi_tangle_group, AspectExt_intel_matrix, AspectExt_oneapi_is_composite, AspectExt_oneapi_is_component,
-    AspectExt_oneapi_graph, AspectExt_intel_fpga_task_sequence],
+    AspectExt_oneapi_graph, AspectExt_intel_fpga_task_sequence, AspectExt_oneapi_limited_graph],
     []>;
 // This definition serves the only purpose of testing whether the deprecated aspect list defined in here and in SYCL RT
 // match.

sycl/doc/design/CommandGraph.md

@@ -46,6 +46,7 @@ with the following entry-points:
 | `urCommandBufferAppendMemBufferReadRectExp`  | Append a rectangular memory read command to a command-buffer object. |
 | `urCommandBufferAppendMemBufferFillExp`      | Append a memory fill command to a command-buffer object. |
 | `urCommandBufferEnqueueExp`                  | Submit command-buffer to a command-queue for execution. |
+| `urCommandBufferUpdateKernelLaunchExp`       | Updates the parameters of a previous kernel launch command. |
 
 See the [UR EXP-COMMAND-BUFFER](https://oneapi-src.github.io/unified-runtime/core/EXP-COMMAND-BUFFER.html)
 specification for more details.
@@ -230,6 +231,47 @@ on buffer usage in a graph so that their lifetime semantics are compatible with
 a lazy work execution model. However these changes to storage lifetimes have not
 yet been implemented.
 
+## Graph Update
+
+### Design Challenges
+
+Graph update faces significant design challenges in SYCL:
+
+* Lambda capture order is explicitly undefined in C++, so the user cannot reason
+  about the indices of arguments captured by kernel lambdas.
+* Once arguments have been captured the actual type information is lost in the
+  transition through the integration header and extracting arguments in the SYCL
+  runtime, therefore we cannot automatically match new argument values by
+  querying the captured arguments without significant possibility for
+  collisions. For example, if a kernel captures two USM pointers and the user
+  wishes to update one, we cannot reason about which pointer they actually want
+  to update when we only know that: they are pointer args of a certain size.
+
+The current approach is to limit graph update to the explicit APIs and where the
+user is using `handler::set_arg()` or some equivalent to manually set kernel
+arguments using indices. Therefore when updating we can use indices to avoid
+collisions. In practice there are only a few current scenarios where `set_arg()`
+can be used:
+
+* The proposed ["Free Function Kernel"
+  extension](../extensions/proposed/sycl_ext_oneapi_free_function_kernels.asciidoc)
+* OpenCL interop kernels created from SPIR-V source at runtime.
+
+A possible future workaround lambda capture issues could be "Whole-Graph Update"
+where if we can guarantee that lambda capture order is the same across two
+different recordings we can then match parameter order when updating.
+
+### Scheduler Integration
+
+Graph updates in the runtime are synchronous calls however they can optionally
+be done through the scheduler using a new command,
+`sycl::detail::UpdateCommandBufferCommand`. This is needed when dealing with
+accessor updates. Since a new buffer which the user creates for updating may not
+yet have been lazily initialized on device we schedule a new command which has
+requirements for these new accessors to correctly trigger allocations before
+updating. This is similar to how individual graph commands are enqueued when
+accessors are used in a graph node.
+
 ## Backend Implementation
 
 Implementation of UR command-buffers

sycl/include/sycl/detail/pi.def

@@ -184,6 +184,9 @@ _PI_API(piextCommandBufferFillUSM)
 _PI_API(piextCommandBufferPrefetchUSM)
 _PI_API(piextCommandBufferAdviseUSM)
 _PI_API(piextEnqueueCommandBuffer)
+_PI_API(piextCommandBufferUpdateKernelLaunch)
+_PI_API(piextCommandBufferRetainCommand)
+_PI_API(piextCommandBufferReleaseCommand)
 
 _PI_API(piextUSMPitchedAlloc)
 

sycl/include/sycl/detail/pi.h

@@ -156,9 +156,10 @@
 //       piextEnqueueCooperativeKernelLaunch.
 // 15.46 Add piextGetGlobalVariablePointer
 // 15.47 Added PI_ERROR_FEATURE_UNSUPPORTED.
+// 15.48 Add CommandBuffer update definitions
 
 #define _PI_H_VERSION_MAJOR 15
-#define _PI_H_VERSION_MINOR 47
+#define _PI_H_VERSION_MINOR 48
 
 #define _PI_STRING_HELPER(a) #a
 #define _PI_CONCAT(a, b) _PI_STRING_HELPER(a.b)
@@ -445,6 +446,10 @@ typedef enum {
   // Composite device
   PI_EXT_ONEAPI_DEVICE_INFO_COMPONENT_DEVICES = 0x20111,
   PI_EXT_ONEAPI_DEVICE_INFO_COMPOSITE_DEVICE = 0x20112,
+
+  // Command Buffers
+  PI_EXT_ONEAPI_DEVICE_INFO_COMMAND_BUFFER_SUPPORT = 0x20113,
+  PI_EXT_ONEAPI_DEVICE_INFO_COMMAND_BUFFER_UPDATE_SUPPORT = 0x20114,
 } _pi_device_info;
 
 typedef enum {
@@ -2320,7 +2325,10 @@ __SYCL_EXPORT pi_result piGetDeviceAndHostTimer(pi_device Device,
 /// Command buffer extension
 struct _pi_ext_command_buffer;
 struct _pi_ext_sync_point;
+struct _pi_ext_command_buffer_command;
+
 using pi_ext_command_buffer = _pi_ext_command_buffer *;
+using pi_ext_command_buffer_command = _pi_ext_command_buffer_command *;
 using pi_ext_sync_point = pi_uint32;
 
 typedef enum {
@@ -2330,7 +2338,40 @@ typedef enum {
 struct pi_ext_command_buffer_desc final {
   pi_ext_structure_type stype;
   const void *pNext;
-  pi_queue_properties *properties;
+  pi_bool is_updatable;
+};
+
+// Command Buffer Update types
+struct pi_ext_command_buffer_update_memobj_arg_desc_t final {
+  uint32_t arg_index;
+  const pi_mem_obj_property *properties;
+  pi_mem new_mem_obj;
+};
+
+struct pi_ext_command_buffer_update_pointer_arg_desc_t final {
+  uint32_t arg_index;
+  void *new_ptr;
+};
+
+struct pi_ext_command_buffer_update_value_arg_desc_t final {
+  uint32_t arg_index;
+  uint32_t arg_size;
+  void *new_value;
+};
+
+struct pi_ext_command_buffer_update_kernel_launch_desc final {
+  uint32_t num_mem_obj_args;
+  uint32_t num_ptr_args;
+  uint32_t num_value_args;
+  uint32_t num_work_dim;
+
+  pi_ext_command_buffer_update_memobj_arg_desc_t *mem_obj_arg_list;
+  pi_ext_command_buffer_update_pointer_arg_desc_t *ptr_arg_list;
+  pi_ext_command_buffer_update_value_arg_desc_t *value_arg_list;
+
+  size_t *global_work_offset;
+  size_t *global_work_size;
+  size_t *local_work_size;
 };
 
 /// API to create a command-buffer.
@@ -2374,12 +2415,14 @@ piextCommandBufferFinalize(pi_ext_command_buffer command_buffer);
 /// \param sync_point_wait_list A list of sync points that this command must
 /// wait on.
 /// \param sync_point The sync_point associated with this kernel execution.
+/// \param command Return pointer to the command representing this kernel
+/// execution.
 __SYCL_EXPORT pi_result piextCommandBufferNDRangeKernel(
     pi_ext_command_buffer command_buffer, pi_kernel kernel, pi_uint32 work_dim,
     const size_t *global_work_offset, const size_t *global_work_size,
     const size_t *local_work_size, pi_uint32 num_sync_points_in_wait_list,
     const pi_ext_sync_point *sync_point_wait_list,
-    pi_ext_sync_point *sync_point);
+    pi_ext_sync_point *sync_point, pi_ext_command_buffer_command *command);
 
 /// API to append a USM memcpy command to the command-buffer.
 /// \param command_buffer The command-buffer to append onto.
@@ -2607,6 +2650,25 @@ piextEnqueueCommandBuffer(pi_ext_command_buffer command_buffer, pi_queue queue,
                           pi_uint32 num_events_in_wait_list,
                           const pi_event *event_wait_list, pi_event *event);
 
+/// API to update a kernel launch command inside of a command-buffer.
+/// @param command The command to be updated.
+/// @param desc Descriptor which describes the updated parameters of the kernel
+/// launch.
+__SYCL_EXPORT pi_result piextCommandBufferUpdateKernelLaunch(
+    pi_ext_command_buffer_command command,
+    pi_ext_command_buffer_update_kernel_launch_desc *desc);
+
+/// API to increment the reference count of a command-buffer command.
+/// \param command The command to release.
+__SYCL_EXPORT pi_result
+piextCommandBufferRetainCommand(pi_ext_command_buffer_command command);
+
+/// API to decrement the reference count of a command-buffer command. After the
+/// command reference count becomes zero, the command is deleted.
+/// \param command The command to release.
+__SYCL_EXPORT pi_result
+piextCommandBufferReleaseCommand(pi_ext_command_buffer_command command);
+
 /// API to destroy bindless unsampled image handles.
 ///
 /// \param context is the pi_context

sycl/include/sycl/detail/pi.hpp

@@ -156,6 +156,7 @@ using PiKernelCacheConfig = ::pi_kernel_cache_config;
 using PiExtSyncPoint = ::pi_ext_sync_point;
 using PiExtCommandBuffer = ::pi_ext_command_buffer;
 using PiExtCommandBufferDesc = ::pi_ext_command_buffer_desc;
+using PiExtCommandBufferCommand = ::pi_ext_command_buffer_command;
 using PiPeerAttr = ::pi_peer_attr;
 using PiImageHandle = ::pi_image_handle;
 using PiImageMemHandle = ::pi_image_mem_handle;

sycl/include/sycl/detail/property_helper.hpp

@@ -47,8 +47,9 @@ enum DataLessPropKind {
   GraphAssumeDataOutlivesBuffer = 22,
   GraphAssumeBufferOutlivesGraph = 23,
   GraphDependOnAllLeaves = 24,
+  GraphUpdatable = 25,
   // Indicates the last known dataless property.
-  LastKnownDataLessPropKind = 24,
+  LastKnownDataLessPropKind = 25,
   // Exceeding 32 may cause ABI breaking change on some of OSes.
   DataLessPropKindSize = 32
 };

sycl/include/sycl/device_aspect_macros.hpp

@@ -323,6 +323,11 @@
 #define __SYCL_ALL_DEVICES_HAVE_ext_intel_fpga_task_sequence__ 0
 #endif
 
+#ifndef __SYCL_ALL_DEVICES_HAVE_ext_oneapi_limited_graph__
+// __SYCL_ASPECT(ext_oneapi_limited_graph, 63)
+#define __SYCL_ALL_DEVICES_HAVE_ext_oneapi_limited_graph__ 0
+#endif
+
 #ifndef __SYCL_ANY_DEVICE_HAS_host__
 // __SYCL_ASPECT(host, 0)
 #define __SYCL_ANY_DEVICE_HAS_host__ 0
@@ -637,3 +642,8 @@
 // __SYCL_ASPECT(ext_intel_fpga_task_sequence__, 62)
 #define __SYCL_ANY_DEVICE_HAS_ext_intel_fpga_task_sequence__ 0
 #endif
+
+#ifndef __SYCL_ANY_DEVICE_HAS_ext_oneapi_limited_graph__
+// __SYCL_ASPECT(ext_oneapi_limited_graph, 63)
+#define __SYCL_ANY_DEVICE_HAS_ext_oneapi_limited_graph__ 0
+#endif