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

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

Author: ldorau

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

@@ -526,6 +526,16 @@ 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;
+
   return UR_RESULT_SUCCESS;
 }
 

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

@@ -162,4 +162,12 @@ 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;
+    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

@@ -1039,7 +1039,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 +1048,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++) {
@@ -1095,3 +1093,126 @@ ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExp(
 
   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 numPropsInLaunchPropList,
+    const ur_kernel_launch_property_t *launchPropList,
+    uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+    ur_event_handle_t phEvent) {
+
+  // TODO: remove memory allocation
+
+  // kernelMemObj contains kernel memory objects that
+  // UR_EXP_KERNEL_ARG_TYPE_MEM_OBJ kernelArgs pointers point to
+  std::vector<void *> kernelMemObj(numArgs, nullptr);
+  std::vector<void *> kernelArgs(numArgs, nullptr);
+  std::scoped_lock<ur_shared_mutex> Lock(hKernel->Mutex);
+
+  for (uint32_t argIndex = 0; argIndex < numArgs; argIndex++) {
+    switch (pArgs[argIndex].type) {
+    case UR_EXP_KERNEL_ARG_TYPE_LOCAL:
+      kernelArgs[argIndex] = (void *)&pArgs[argIndex].size;
+      break;
+    case UR_EXP_KERNEL_ARG_TYPE_VALUE:
+      kernelArgs[argIndex] = (void *)pArgs[argIndex].value.value;
+      break;
+    case UR_EXP_KERNEL_ARG_TYPE_POINTER:
+      kernelArgs[argIndex] = (void *)&pArgs[argIndex].value.pointer;
+      break;
+    case UR_EXP_KERNEL_ARG_TYPE_MEM_OBJ:
+      // prepareForSubmission() will save zePtr in kernelMemObj[argIndex]
+      kernelArgs[argIndex] = &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: {
+      kernelArgs[argIndex] = &pArgs[argIndex].value.sampler->ZeSampler;
+      break;
+    }
+    default:
+      return UR_RESULT_ERROR_INVALID_ENUMERATION;
+    }
+  }
+
+  for (uint32_t propIndex = 0; propIndex < numPropsInLaunchPropList;
+       propIndex++) {
+    if (launchPropList[propIndex].id != UR_KERNEL_LAUNCH_PROPERTY_ID_IGNORE) {
+      // We don't support any other properties.
+      return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    }
+  }
+
+  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());
+
+  ze_group_count_t zeThreadGroupDimensions{1, 1, 1};
+  uint32_t WG[3]{};
+  UR_CALL(calculateKernelWorkDimensions(hZeKernel, hDevice.get(),
+                                        zeThreadGroupDimensions, WG, workDim,
+                                        pGlobalWorkSize, pLocalWorkSize));
+
+  ze_group_size_t groupSize = {WG[0], WG[1], WG[2]};
+
+  auto zeSignalEvent = getSignalEvent(phEvent, UR_COMMAND_KERNEL_LAUNCH);
+  auto waitListView = getWaitListView(phEventWaitList, numEventsInWaitList);
+
+  UR_CALL(hKernel->prepareForSubmission(
+      hContext.get(), hDevice.get(), pGlobalWorkOffset, workDim, WG[0], WG[1],
+      WG[2], getZeCommandList(), waitListView, &kernelArgs));
+
+  {
+    TRACK_SCOPE_LATENCY("ur_command_list_manager::"
+                        "zeCommandListAppendLaunchKernelWithArguments");
+    ZE2UR_CALL(hContext->getPlatform()
+                   ->ZeCommandListAppendLaunchKernelWithArgumentsExt
+                   .zeCommandListAppendLaunchKernelWithArguments,
+               (getZeCommandList(), hZeKernel, zeThreadGroupDimensions,
+                groupSize, kernelArgs.data(), nullptr, zeSignalEvent,
+                waitListView.num, waitListView.handles));
+  }
+
+  recordSubmittedKernel(hKernel);
+
+  postSubmit(hZeKernel, pGlobalWorkOffset);
+
+  return UR_RESULT_SUCCESS;
+}
+
+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");
+
+  UR_ASSERT(hKernel, UR_RESULT_ERROR_INVALID_NULL_HANDLE);
+
+  if (hContext->getPlatform()
+          ->ZeCommandListAppendLaunchKernelWithArgumentsExt.Supported) {
+    return appendKernelLaunchWithArgsExpNew(
+        hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize, pLocalWorkSize,
+        numArgs, pArgs, numPropsInLaunchPropList, launchPropList,
+        numEventsInWaitList, phEventWaitList, phEvent);
+  } else {
+    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,26 @@ 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 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 appendGenericCommandListsExp(
       uint32_t numCommandLists, ze_command_list_handle_t *phCommandLists,
       ur_event_handle_t phEvent, uint32_t numEventsInWaitList,

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

@@ -271,7 +271,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 *> *kernelArgs) {
   auto &deviceKernelOpt = deviceKernels[deviceIndex(hDevice)];
   if (!deviceKernelOpt.has_value())
     return UR_RESULT_ERROR_INVALID_KERNEL;
@@ -299,8 +300,25 @@ 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));
+
+    // kernelArgs must be non-nullptr in the path of
+    // zeCommandListAppendLaunchKernelWithArguments()
+    if (kernelArgs) {
+      // 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 kernelArgs
+      // for this future call.
+      if (pending.argIndex > kernelArgs->size() - 1) {
+        return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
+      }
+      UR_ASSERT((*kernelArgs)[pending.argIndex] != nullptr,
+                UR_RESULT_ERROR_INVALID_NULL_POINTER);
+      *((void **)(*kernelArgs)[pending.argIndex]) = zePtr;
+    } else {
+      // Set the argument only on this device's kernel.
+      UR_CALL(deviceKernel.setArgPointer(pending.argIndex, zePtr));
+    }
   }
   pending_allocations.clear();
 

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

@@ -96,13 +96,16 @@ struct ur_kernel_handle_t_ : ur_object {
 
   // Set all required values for the kernel before submission (including pending
   // memory allocations).
+  // The kernelArgs argument must be non-nullptr
+  // in the path of zeCommandListAppendLaunchKernelWithArguments()
   ur_result_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 cmdList,
-                                   wait_list_view &waitListView);
+                                   wait_list_view &waitListView,
+                                   std::vector<void *> *kernelArgs = nullptr);
 
   // Get context of the kernel.
   ur_context_handle_t getContext() const { return hProgram->Context; }