[UR][Cuda] Add support for command-buffer kernel updates

Author: Fábio Mestre

sycl/cmake/modules/FetchUnifiedRuntime.cmake

@@ -116,15 +116,15 @@ if(SYCL_UR_USE_FETCH_CONTENT)
       CACHE PATH "Path to external '${name}' adapter source dir" FORCE)
   endfunction()
 
-  set(UNIFIED_RUNTIME_REPO "https://github.com/oneapi-src/unified-runtime.git")
+  set(UNIFIED_RUNTIME_REPO "https://github.com/Bensuo/unified-runtime.git")
   # commit 2ad32681efd2c977f2c1f7f3a30d572d4c15499c
   # Author: Hugh Delaney <[email protected]>
   # Date:   Wed Sep 25 15:55:05 2024 +0100
   #     Associate queue with device in context (#1992)
   #     Making a native queue doesn't require hDevice to be non null, but this
   #     associates the queue with a null device, even if hContext contains valid
   #     devices.
-  set(UNIFIED_RUNTIME_TAG 2ad32681efd2c977f2c1f7f3a30d572d4c15499c)
+  set(UNIFIED_RUNTIME_TAG d944ff3391dfbe69db453406bd0bbcb78716dee0)
 
   set(UMF_BUILD_EXAMPLES OFF CACHE INTERNAL "EXAMPLES")
   # Due to the use of dependentloadflag and no installer for UMF and hwloc we need

sycl/source/detail/device_impl.cpp

@@ -704,17 +704,26 @@ bool device_impl::has(aspect Aspect) const {
     return CallSuccessful && Result != nullptr;
   }
   case aspect::ext_oneapi_graph: {
-    bool SupportsCommandBufferUpdate = false;
+    ur_device_command_buffer_update_capability_flags_t UpdateCapabilities;
     bool CallSuccessful =
         getAdapter()->call_nocheck<UrApiKind::urDeviceGetInfo>(
-            MDevice, UR_DEVICE_INFO_COMMAND_BUFFER_UPDATE_SUPPORT_EXP,
-            sizeof(SupportsCommandBufferUpdate), &SupportsCommandBufferUpdate,
+            MDevice, UR_DEVICE_INFO_COMMAND_BUFFER_UPDATE_CAPABILITIES_EXP,
+            sizeof(UpdateCapabilities), &UpdateCapabilities,
             nullptr) == UR_RESULT_SUCCESS;
     if (!CallSuccessful) {
       return false;
     }
 
-    return has(aspect::ext_oneapi_limited_graph) && SupportsCommandBufferUpdate;
+    /* The kernel handle update capability is not yet required for the
+     * ext_oneapi_graph aspect */
+    ur_device_command_buffer_update_capability_flags_t RequiredCapabilities =
+        UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_KERNEL_ARGUMENTS |
+        UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_LOCAL_WORK_SIZE |
+        UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_GLOBAL_WORK_SIZE |
+        UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_GLOBAL_WORK_OFFSET;
+
+    return has(aspect::ext_oneapi_limited_graph) &&
+           (UpdateCapabilities & RequiredCapabilities) == RequiredCapabilities;
   }
   case aspect::ext_oneapi_limited_graph: {
     bool SupportsCommandBuffers = false;

sycl/source/detail/graph_impl.cpp

@@ -1474,6 +1474,7 @@ void exec_graph_impl::updateImpl(std::shared_ptr<node_impl> Node) {
     }
   }
 
+  UpdateDesc.hNewKernel = UrKernel;
   UpdateDesc.numNewMemObjArgs = MemobjDescs.size();
   UpdateDesc.pNewMemObjArgList = MemobjDescs.data();
   UpdateDesc.numNewPointerArgs = PtrDescs.size();

sycl/source/detail/scheduler/commands.cpp

@@ -2538,7 +2538,7 @@ ur_result_t enqueueImpCommandBufferKernel(
   ur_result_t Res =
       Adapter->call_nocheck<UrApiKind::urCommandBufferAppendKernelLaunchExp>(
           CommandBuffer, UrKernel, NDRDesc.Dims, &NDRDesc.GlobalOffset[0],
-          &NDRDesc.GlobalSize[0], LocalSize, SyncPoints.size(),
+          &NDRDesc.GlobalSize[0], LocalSize, 0, nullptr, SyncPoints.size(),
           SyncPoints.size() ? SyncPoints.data() : nullptr, OutSyncPoint,
           OutCommand);
 

sycl/unittests/helpers/UrMock.hpp

@@ -200,14 +200,27 @@ inline ur_result_t mock_urDeviceGetInfo(void *pParams) {
   case UR_DEVICE_INFO_AVAILABLE:
   case UR_DEVICE_INFO_LINKER_AVAILABLE:
   case UR_DEVICE_INFO_COMPILER_AVAILABLE:
-  case UR_DEVICE_INFO_COMMAND_BUFFER_SUPPORT_EXP:
-  case UR_DEVICE_INFO_COMMAND_BUFFER_UPDATE_SUPPORT_EXP: {
+  case UR_DEVICE_INFO_COMMAND_BUFFER_SUPPORT_EXP: {
     if (*params->ppPropValue)
       *static_cast<ur_bool_t *>(*params->ppPropValue) = true;
     if (*params->ppPropSizeRet)
       **params->ppPropSizeRet = sizeof(true);
     return UR_RESULT_SUCCESS;
   }
+  case UR_DEVICE_INFO_COMMAND_BUFFER_UPDATE_CAPABILITIES_EXP: {
+    if (*params->ppPropValue)
+      *static_cast<ur_device_command_buffer_update_capability_flags_t *>(
+          *params->ppPropValue) =
+          UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_KERNEL_ARGUMENTS |
+          UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_GLOBAL_WORK_SIZE |
+          UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_LOCAL_WORK_SIZE |
+          UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_GLOBAL_WORK_OFFSET |
+          UR_DEVICE_COMMAND_BUFFER_UPDATE_CAPABILITY_FLAG_KERNEL_HANDLE;
+    if (*params->ppPropSizeRet)
+      **params->ppPropSizeRet =
+          sizeof(ur_device_command_buffer_update_capability_flags_t);
+    return UR_RESULT_SUCCESS;
+  }
   // This mock GPU device has no sub-devices
   case UR_DEVICE_INFO_SUPPORTED_PARTITIONS: {
     if (*params->ppPropSizeRet) {