[UR][SYCL] Add support for zeCommandListAppendLaunchKernelWithArguments()

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

unified-runtime/cmake/FetchLevelZero.cmake

@@ -125,7 +125,7 @@ if(L0_COMPUTE_RUNTIME_HEADERS)
     set(COMPUTE_RUNTIME_REPO_PATH "${L0_COMPUTE_RUNTIME_HEADERS}")
 else()
     set(UR_COMPUTE_RUNTIME_REPO "https://github.com/intel/compute-runtime.git")
-    set(UR_COMPUTE_RUNTIME_TAG 25.31.34666.3)
+    set(UR_COMPUTE_RUNTIME_TAG 25.35.35096.9)
 
     include(FetchContent)
     # Sparse fetch only the dir with level zero headers for experimental features to avoid pulling in the entire compute-runtime.

unified-runtime/source/adapters/level_zero/common.cpp

@@ -324,6 +324,11 @@ template <>
 ze_structure_type_t getZeStructureType<ze_device_cache_line_size_ext_t>() {
   return ZE_STRUCTURE_TYPE_DEVICE_CACHELINE_SIZE_EXT;
 }
+template <>
+ze_structure_type_t getZeStructureType<
+    ze_command_list_append_launch_kernel_param_cooperative_desc_t>() {
+  return ZE_STRUCTURE_TYPE_COMMAND_LIST_APPEND_PARAM_COOPERATIVE_DESC;
+}
 
 #ifdef ZE_INTEL_DEVICE_BLOCK_ARRAY_EXP_NAME
 template <>

unified-runtime/source/adapters/level_zero/platform.cpp

@@ -526,6 +526,22 @@ ur_result_t ur_platform_handle_t_::initialize() {
   ZeMemGetPitchFor2dImageExt.Supported =
       ZeMemGetPitchFor2dImageExt.zeMemGetPitchFor2dImage != nullptr;
 
+  ZE_CALL_NOCHECK(zeDriverGetExtensionFunctionAddress,
+                  (ZeDriver, "zeCommandListAppendLaunchKernelWithArguments",
+                   reinterpret_cast<void **>(
+                       &ZeCommandListAppendLaunchKernelWithArgumentsExt
+                            .zeCommandListAppendLaunchKernelWithArguments)));
+
+  ZeCommandListAppendLaunchKernelWithArgumentsExt.Supported =
+      ZeCommandListAppendLaunchKernelWithArgumentsExt
+          .zeCommandListAppendLaunchKernelWithArguments != nullptr;
+
+  // Check if the driver supports zeCommandListAppendLaunchKernelWithArguments()
+  // with cooperative mode (version >= 1.6.35005)
+  ZeCommandListAppendLaunchKernelWithArgumentsExt
+      .DriverSupportsCooperativeKernelLaunchWithArgs =
+      this->isDriverVersionNewerOrSimilar(1, 6, 35005);
+
   return UR_RESULT_SUCCESS;
 }
 

unified-runtime/source/adapters/level_zero/platform.hpp

@@ -162,4 +162,13 @@ struct ur_platform_handle_t_ : ur::handle_base<ur::level_zero::ddi_getter>,
                                            ze_device_handle_t, size_t, size_t,
                                            unsigned int, size_t *);
   } ZeMemGetPitchFor2dImageExt;
+
+  struct ZeCommandListAppendLaunchKernelWithArgumentsExtension {
+    bool Supported = false;
+    bool DriverSupportsCooperativeKernelLaunchWithArgs = false;
+    ze_result_t (*zeCommandListAppendLaunchKernelWithArguments)(
+        ze_command_list_handle_t, ze_kernel_handle_t, const ze_group_count_t,
+        const ze_group_size_t, void **, void *, ze_event_handle_t, uint32_t,
+        ze_event_handle_t *);
+  } ZeCommandListAppendLaunchKernelWithArgumentsExt;
 };

unified-runtime/source/adapters/level_zero/v2/command_list_manager.cpp

@@ -14,6 +14,7 @@
 #include "../sampler.hpp"
 #include "../ur_interface_loader.hpp"
 #include "command_buffer.hpp"
+#include "common.hpp"
 #include "context.hpp"
 #include "kernel.hpp"
 #include "memory.hpp"
@@ -149,21 +150,13 @@ ur_command_list_manager::getSignalEvent(ur_event_handle_t hUserEvent,
   }
 }
 
-ur_result_t ur_command_list_manager::appendKernelLaunchUnlocked(
-    ur_kernel_handle_t hKernel, uint32_t workDim,
+// must be called with hKernel->Mutex held
+ur_result_t ur_command_list_manager::appendKernelLaunchLocked(
+    ur_kernel_handle_t hKernel, ze_kernel_handle_t hZeKernel, uint32_t workDim,
     const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
     const size_t *pLocalWorkSize, uint32_t numEventsInWaitList,
     const ur_event_handle_t *phEventWaitList, ur_event_handle_t phEvent,
-    bool cooperative) {
-  UR_ASSERT(hKernel, UR_RESULT_ERROR_INVALID_NULL_HANDLE);
-  UR_ASSERT(hKernel->getProgramHandle(), UR_RESULT_ERROR_INVALID_NULL_POINTER);
-
-  UR_ASSERT(workDim > 0, UR_RESULT_ERROR_INVALID_WORK_DIMENSION);
-  UR_ASSERT(workDim < 4, UR_RESULT_ERROR_INVALID_WORK_DIMENSION);
-
-  ze_kernel_handle_t hZeKernel = hKernel->getZeHandle(hDevice.get());
-
-  std::scoped_lock<ur_shared_mutex> Lock(hKernel->Mutex);
+    bool cooperative, std::vector<void *> &kMemObj, void *pNext) {
 
   ze_group_count_t zeThreadGroupDimensions{1, 1, 1};
   uint32_t WG[3]{};
@@ -176,15 +169,28 @@ ur_result_t ur_command_list_manager::appendKernelLaunchUnlocked(
 
   UR_CALL(hKernel->prepareForSubmission(
       hContext.get(), hDevice.get(), pGlobalWorkOffset, workDim, WG[0], WG[1],
-      WG[2], getZeCommandList(), waitListView));
+      WG[2], getZeCommandList(), waitListView, kMemObj));
 
-  if (cooperative) {
+  if (!kMemObj.empty()) {
+    // zeCommandListAppendLaunchKernelWithArguments
+    TRACK_SCOPE_LATENCY("ur_command_list_manager::"
+                        "zeCommandListAppendLaunchKernelWithArguments");
+    ze_group_size_t groupSize = {WG[0], WG[1], WG[2]};
+    ZE2UR_CALL(hContext->getPlatform()
+                   ->ZeCommandListAppendLaunchKernelWithArgumentsExt
+                   .zeCommandListAppendLaunchKernelWithArguments,
+               (getZeCommandList(), hZeKernel, zeThreadGroupDimensions,
+                groupSize, hKernel->kernelArgs.data(), pNext, zeSignalEvent,
+                waitListView.num, waitListView.handles));
+  } else if (cooperative) {
+    // zeCommandListAppendLaunchCooperativeKernel
     TRACK_SCOPE_LATENCY("ur_command_list_manager::"
                         "zeCommandListAppendLaunchCooperativeKernel");
     ZE2UR_CALL(zeCommandListAppendLaunchCooperativeKernel,
                (getZeCommandList(), hZeKernel, &zeThreadGroupDimensions,
                 zeSignalEvent, waitListView.num, waitListView.handles));
   } else {
+    // zeCommandListAppendLaunchKernel
     TRACK_SCOPE_LATENCY("ur_command_list_manager::"
                         "zeCommandListAppendLaunchKernel");
     ZE2UR_CALL(zeCommandListAppendLaunchKernel,
@@ -199,6 +205,39 @@ ur_result_t ur_command_list_manager::appendKernelLaunchUnlocked(
   return UR_RESULT_SUCCESS;
 }
 
+static ur_result_t kernelLaunchChecks(ur_kernel_handle_t hKernel,
+                                      uint32_t workDim) {
+  UR_ASSERT(hKernel, UR_RESULT_ERROR_INVALID_NULL_HANDLE);
+  UR_ASSERT(hKernel->getProgramHandle(), UR_RESULT_ERROR_INVALID_NULL_POINTER);
+  UR_ASSERT(workDim > 0, UR_RESULT_ERROR_INVALID_WORK_DIMENSION);
+  UR_ASSERT(workDim < 4, UR_RESULT_ERROR_INVALID_WORK_DIMENSION);
+
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t ur_command_list_manager::appendKernelLaunchUnlocked(
+    ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+    const size_t *pLocalWorkSize, uint32_t numEventsInWaitList,
+    const ur_event_handle_t *phEventWaitList, ur_event_handle_t phEvent,
+    bool cooperative) {
+
+  ur_result_t checkResult = kernelLaunchChecks(hKernel, workDim);
+  if (checkResult != UR_RESULT_SUCCESS) {
+    return checkResult;
+  }
+
+  ze_kernel_handle_t hZeKernel = hKernel->getZeHandle(hDevice.get());
+  std::vector<void *> emptyKMemObj;
+
+  std::scoped_lock<ur_shared_mutex> Lock(hKernel->Mutex);
+
+  return appendKernelLaunchLocked(
+      hKernel, hZeKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
+      pLocalWorkSize, numEventsInWaitList, phEventWaitList, phEvent,
+      cooperative, emptyKMemObj, nullptr /* pNext */);
+}
+
 ur_result_t ur_command_list_manager::appendKernelLaunch(
     ur_kernel_handle_t hKernel, uint32_t workDim,
     const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
@@ -1039,7 +1078,7 @@ ur_result_t ur_command_list_manager::releaseSubmittedKernels() {
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExp(
+ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExpOld(
     ur_kernel_handle_t hKernel, uint32_t workDim,
     const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
     const size_t *pLocalWorkSize, uint32_t numArgs,
@@ -1048,8 +1087,6 @@ ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExp(
     const ur_kernel_launch_property_t *launchPropList,
     uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
     ur_event_handle_t phEvent) {
-  TRACK_SCOPE_LATENCY(
-      "ur_queue_immediate_in_order_t::enqueueKernelLaunchWithArgsExp");
   {
     std::scoped_lock<ur_shared_mutex> guard(hKernel->Mutex);
     for (uint32_t argIndex = 0; argIndex < numArgs; argIndex++) {
@@ -1091,7 +1128,129 @@ ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExp(
                              numPropsInLaunchPropList, launchPropList,
                              numEventsInWaitList, phEventWaitList, phEvent));
 
-  recordSubmittedKernel(hKernel);
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExpNew(
+    ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+    const size_t *pLocalWorkSize, uint32_t numArgs,
+    const ur_exp_kernel_arg_properties_t *pArgs, uint32_t numEventsInWaitList,
+    const ur_event_handle_t *phEventWaitList, ur_event_handle_t phEvent,
+    bool cooperativeKernelLaunchRequested) {
+
+  ur_result_t checkResult = kernelLaunchChecks(hKernel, workDim);
+  if (checkResult != UR_RESULT_SUCCESS) {
+    return checkResult;
+  }
+
+  // It is needed in case of UR_KERNEL_LAUNCH_PROPERTY_ID_COOPERATIVE
+  // to launch the cooperative kernel.
+  ZeStruct<ze_command_list_append_launch_kernel_param_cooperative_desc_t>
+      cooperativeDesc;
+  cooperativeDesc.isCooperative = static_cast<ze_bool_t>(true);
+
+  void *pNext = nullptr;
+  if (cooperativeKernelLaunchRequested) {
+    pNext = &cooperativeDesc;
+  }
+
+  ze_kernel_handle_t hZeKernel = hKernel->getZeHandle(hDevice.get());
+
+  std::scoped_lock<ur_shared_mutex> Lock(hKernel->Mutex);
+
+  // kernelMemObj contains kernel memory objects that
+  // UR_EXP_KERNEL_ARG_TYPE_MEM_OBJ kernelArgs pointers point to
+  hKernel->kernelMemObj.resize(numArgs, 0);
+  hKernel->kernelArgs.resize(numArgs, 0);
+
+  for (uint32_t argIndex = 0; argIndex < numArgs; argIndex++) {
+    switch (pArgs[argIndex].type) {
+    case UR_EXP_KERNEL_ARG_TYPE_LOCAL:
+      hKernel->kernelArgs[argIndex] = (void *)&pArgs[argIndex].size;
+      break;
+    case UR_EXP_KERNEL_ARG_TYPE_VALUE:
+      hKernel->kernelArgs[argIndex] = (void *)pArgs[argIndex].value.value;
+      break;
+    case UR_EXP_KERNEL_ARG_TYPE_POINTER:
+      hKernel->kernelArgs[argIndex] = (void *)&pArgs[argIndex].value.pointer;
+      break;
+    case UR_EXP_KERNEL_ARG_TYPE_MEM_OBJ:
+      // prepareForSubmission() will save zePtr in kernelMemObj[argIndex]
+      hKernel->kernelArgs[argIndex] = &hKernel->kernelMemObj[argIndex];
+      UR_CALL(hKernel->addPendingMemoryAllocation(
+          {pArgs[argIndex].value.memObjTuple.hMem,
+           ur_mem_buffer_t::device_access_mode_t::read_write,
+           pArgs[argIndex].index}));
+      break;
+    case UR_EXP_KERNEL_ARG_TYPE_SAMPLER:
+      hKernel->kernelArgs[argIndex] = &pArgs[argIndex].value.sampler->ZeSampler;
+      break;
+    default:
+      return UR_RESULT_ERROR_INVALID_ENUMERATION;
+    }
+  }
+
+  return appendKernelLaunchLocked(
+      hKernel, hZeKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
+      pLocalWorkSize, numEventsInWaitList, phEventWaitList, phEvent,
+      cooperativeKernelLaunchRequested, hKernel->kernelMemObj, pNext);
+}
+
+ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExp(
+    ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+    const size_t *pLocalWorkSize, uint32_t numArgs,
+    const ur_exp_kernel_arg_properties_t *pArgs,
+    uint32_t numPropsInLaunchPropList,
+    const ur_kernel_launch_property_t *launchPropList,
+    uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+    ur_event_handle_t phEvent) {
+  TRACK_SCOPE_LATENCY(
+      "ur_queue_immediate_in_order_t::enqueueKernelLaunchWithArgsExp");
+
+  bool cooperativeKernelLaunchRequested = false;
+
+  for (uint32_t propIndex = 0; propIndex < numPropsInLaunchPropList;
+       propIndex++) {
+    switch (launchPropList[propIndex].id) {
+    case UR_KERNEL_LAUNCH_PROPERTY_ID_IGNORE:
+      break;
+    case UR_KERNEL_LAUNCH_PROPERTY_ID_COOPERATIVE:
+      if (launchPropList[propIndex].value.cooperative) {
+        cooperativeKernelLaunchRequested = true;
+      }
+      break;
+    default:
+      // We don't support any other properties.
+      return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    }
+  }
+
+  ur_platform_handle_t hPlatform = hContext->getPlatform();
+  bool KernelWithArgsSupported =
+      hPlatform->ZeCommandListAppendLaunchKernelWithArgumentsExt.Supported;
+  bool CooperativeCompatible =
+      hPlatform->ZeCommandListAppendLaunchKernelWithArgumentsExt
+          .DriverSupportsCooperativeKernelLaunchWithArgs;
+  bool RunNewPath =
+      KernelWithArgsSupported &&
+      (!cooperativeKernelLaunchRequested ||
+       (cooperativeKernelLaunchRequested && CooperativeCompatible));
+  if (RunNewPath) {
+    return appendKernelLaunchWithArgsExpNew(
+        hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize, pLocalWorkSize,
+        numArgs, pArgs, numEventsInWaitList, phEventWaitList, phEvent,
+        cooperativeKernelLaunchRequested);
+  } else {
+    // We cannot pass cooperativeKernelLaunchRequested to
+    // appendKernelLaunchWithArgsExpOld() because appendKernelLaunch() must
+    // check it on its own since it is called also from enqueueKernelLaunch().
+    return appendKernelLaunchWithArgsExpOld(
+        hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize, pLocalWorkSize,
+        numArgs, pArgs, numPropsInLaunchPropList, launchPropList,
+        numEventsInWaitList, phEventWaitList, phEvent);
+  }
 
   return UR_RESULT_SUCCESS;
 }

unified-runtime/source/adapters/level_zero/v2/command_list_manager.hpp

@@ -244,6 +244,24 @@ struct ur_command_list_manager {
       ur_event_handle_t phEvent);
 
 private:
+  ur_result_t appendKernelLaunchWithArgsExpOld(
+      ur_kernel_handle_t hKernel, uint32_t workDim,
+      const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+      const size_t *pLocalWorkSize, uint32_t numArgs,
+      const ur_exp_kernel_arg_properties_t *pArgs,
+      uint32_t numPropsInLaunchPropList,
+      const ur_kernel_launch_property_t *launchPropList,
+      uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+      ur_event_handle_t phEvent);
+
+  ur_result_t appendKernelLaunchWithArgsExpNew(
+      ur_kernel_handle_t hKernel, uint32_t workDim,
+      const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+      const size_t *pLocalWorkSize, uint32_t numArgs,
+      const ur_exp_kernel_arg_properties_t *pArgs, uint32_t numEventsInWaitList,
+      const ur_event_handle_t *phEventWaitList, ur_event_handle_t phEvent,
+      bool cooperativeKernelLaunchRequested);
+
   ur_result_t appendGenericCommandListsExp(
       uint32_t numCommandLists, ze_command_list_handle_t *phCommandLists,
       ur_event_handle_t phEvent, uint32_t numEventsInWaitList,
@@ -258,6 +276,14 @@ struct ur_command_list_manager {
   ze_event_handle_t getSignalEvent(ur_event_handle_t hUserEvent,
                                    ur_command_t commandType);
 
+  ur_result_t appendKernelLaunchLocked(
+      ur_kernel_handle_t hKernel, ze_kernel_handle_t hZeKernel,
+      uint32_t workDim, const size_t *pGlobalWorkOffset,
+      const size_t *pGlobalWorkSize, const size_t *pLocalWorkSize,
+      uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+      ur_event_handle_t phEvent, bool cooperative, std::vector<void *> &kMemObj,
+      void *pNext);
+
   ur_result_t appendKernelLaunchUnlocked(
       ur_kernel_handle_t hKernel, uint32_t workDim,
       const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,

unified-runtime/source/adapters/level_zero/v2/kernel.cpp

@@ -270,7 +270,8 @@ ur_result_t ur_kernel_handle_t_::prepareForSubmission(
     ur_context_handle_t hContext, ur_device_handle_t hDevice,
     const size_t *pGlobalWorkOffset, uint32_t workDim, uint32_t groupSizeX,
     uint32_t groupSizeY, uint32_t groupSizeZ,
-    ze_command_list_handle_t commandList, wait_list_view &waitListView) {
+    ze_command_list_handle_t commandList, wait_list_view &waitListView,
+    std::vector<void *> &kMemObj) {
   auto &deviceKernelOpt = deviceKernels[deviceIndex(hDevice)];
   if (!deviceKernelOpt.has_value())
     return UR_RESULT_ERROR_INVALID_KERNEL;
@@ -298,8 +299,23 @@ ur_result_t ur_kernel_handle_t_::prepareForSubmission(
         zePtr = reinterpret_cast<void *>(hImage->getZeImage());
       }
     }
-    // Set the argument only on this device's kernel.
-    UR_CALL(deviceKernel.setArgPointer(pending.argIndex, zePtr));
+
+    // kMemObj must be a non-empty vector in the path of
+    // zeCommandListAppendLaunchKernelWithArguments()
+    if (!kMemObj.empty()) {
+      // zeCommandListAppendLaunchKernelWithArguments()
+      // (==CommandListCoreFamily<gfxCoreFamily>::appendLaunchKernelWithArguments())
+      // calls setArgumentValue(i, argSize, argValue) for all arguments on its
+      // own so do not call it here, but save the zePtr pointer in kMemObj
+      // for this future call.
+      if (pending.argIndex > kMemObj.size() - 1) {
+        return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
+      }
+      kMemObj[pending.argIndex] = zePtr;
+    } else {
+      // Set the argument only on this device's kernel.
+      UR_CALL(deviceKernel.setArgPointer(pending.argIndex, zePtr));
+    }
   }
   pending_allocations.clear();