Update interface load with new api.

- Move new function to enqueue.cpp.
- Fix impl/tests.
- Add device extension string.
- Clean up documentation.

Signed-off-by: JackAKirk <[email protected]>

Author: JackAKirk

scripts/core/EXP-LAUNCH-ATTRIBUTES.rst

@@ -35,17 +35,17 @@ dimension over which the shared memory is accessible. Additionally some
 applications require specification of kernel attributes at runtime.
 
 This extension is a future-proof and portable solution that supports these two requirements.
-Instead of using a fixed set of kernel arguments, the approach is to introduce 
-"exp_launch_attribute_t" that enables a more flexible approach.
-Each exp_launch_attr_t corresponds to a specific kernel launch attribute.
-One new function is introduced. "urEnqueueKernelLaunchCustomExp" takes an
-array of launch_attribute_t as an argument, and launches a kernel using these
-attributes. "urEnqueueKernelLaunchCustomExp" corresponds to the CUDA Driver API
-"cuLaunchKernelEx".
-
-Many kernel properties could be supported, such as cooperative kernels. As such,
+Instead of using a fixed set of kernel arguments, the approach is to introduce the
+`exp_launch_attribute_t` type that enables a more flexible approach.
+Each `exp_launch_attribute_t` instance corresponds to a specific kernel launch attribute.
+One new function is introduced; `urEnqueueKernelLaunchCustomExp` takes an
+array of `exp_launch_attribute_t` as an argument, and launches a kernel using these
+attributes. `urEnqueueKernelLaunchCustomExp` corresponds closely to the CUDA Driver API
+`cuLaunchKernelEx`.
+
+Many kernel properties can be supported, such as cooperative kernels. As such,
 eventually this extension should be able to replace the cooperative kernels
-UR extension. 
+UR extension.
 
 API
 --------------------------------------------------------------------------------

source/adapters/cuda/device.cpp

@@ -628,6 +628,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     // Return supported for the UR command-buffer experimental feature
     SupportedExtensions += "ur_exp_command_buffer ";
     SupportedExtensions += "ur_exp_usm_p2p ";
+    SupportedExtensions += "ur_exp_launch_attributes ";
     SupportedExtensions += " ";
 
     int Major = 0;

source/adapters/cuda/enqueue.cpp

@@ -484,6 +484,133 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueCooperativeKernelLaunchExp(
                                numEventsInWaitList, phEventWaitList, phEvent);
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunchCustomExp(
+    ur_queue_handle_t hQueue, ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkSize, const size_t *pLocalWorkSize,
+    uint32_t numAttrsInLaunchAttrList,
+    const ur_exp_launch_attribute_t *launchAttrList,
+    uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+    ur_event_handle_t *phEvent) {
+
+  if (numAttrsInLaunchAttrList == 0) {
+    urEnqueueKernelLaunch(hQueue, hKernel, workDim, nullptr, pGlobalWorkSize,
+                          pLocalWorkSize, numEventsInWaitList, phEventWaitList,
+                          phEvent);
+  }
+
+  // Preconditions
+  UR_ASSERT(hQueue->getContext() == hKernel->getContext(),
+            UR_RESULT_ERROR_INVALID_KERNEL);
+  UR_ASSERT(workDim > 0, UR_RESULT_ERROR_INVALID_WORK_DIMENSION);
+  UR_ASSERT(workDim < 4, UR_RESULT_ERROR_INVALID_WORK_DIMENSION);
+
+  if (launchAttrList == NULL) {
+    return UR_RESULT_ERROR_INVALID_NULL_POINTER;
+  }
+
+  std::vector<CUlaunchAttribute> launch_attribute(numAttrsInLaunchAttrList);
+  for (uint32_t i = 0; i < numAttrsInLaunchAttrList; i++) {
+    switch (launchAttrList[i].id) {
+    case UR_EXP_LAUNCH_ATTRIBUTE_ID_IGNORE: {
+      launch_attribute[i].id = CU_LAUNCH_ATTRIBUTE_IGNORE;
+      break;
+    }
+    case UR_EXP_LAUNCH_ATTRIBUTE_ID_CLUSTER_DIMENSION: {
+
+      launch_attribute[i].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
+      launch_attribute[i].value.clusterDim.x =
+          launchAttrList[i].value.clusterDim[0];
+      launch_attribute[i].value.clusterDim.y =
+          launchAttrList[i].value.clusterDim[1];
+      launch_attribute[i].value.clusterDim.z =
+          launchAttrList[i].value.clusterDim[2];
+      break;
+    }
+    case UR_EXP_LAUNCH_ATTRIBUTE_ID_COOPERATIVE: {
+      launch_attribute[i].id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE;
+      launch_attribute[i].value.cooperative =
+          launchAttrList[i].value.cooperative;
+      break;
+    }
+    default: {
+      return UR_RESULT_ERROR_INVALID_ENUMERATION;
+    }
+    }
+  }
+
+  if (*pGlobalWorkSize == 0) {
+    return urEnqueueEventsWaitWithBarrier(hQueue, numEventsInWaitList,
+                                          phEventWaitList, phEvent);
+  }
+
+  // Set the number of threads per block to the number of threads per warp
+  // by default unless user has provided a better number
+  size_t ThreadsPerBlock[3] = {32u, 1u, 1u};
+  size_t BlocksPerGrid[3] = {1u, 1u, 1u};
+
+  uint32_t LocalSize = hKernel->getLocalSize();
+  ur_result_t Result = UR_RESULT_SUCCESS;
+  CUfunction CuFunc = hKernel->get();
+
+  Result = setKernelParams(hQueue->getContext(), hQueue->Device, workDim,
+                           nullptr, pGlobalWorkSize, pLocalWorkSize, hKernel,
+                           CuFunc, ThreadsPerBlock, BlocksPerGrid);
+  if (Result != UR_RESULT_SUCCESS) {
+    return Result;
+  }
+
+  try {
+    std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
+
+    uint32_t StreamToken;
+    ur_stream_guard_ Guard;
+    CUstream CuStream = hQueue->getNextComputeStream(
+        numEventsInWaitList, phEventWaitList, Guard, &StreamToken);
+
+    Result = enqueueEventsWait(hQueue, CuStream, numEventsInWaitList,
+                               phEventWaitList);
+
+    if (phEvent) {
+      RetImplEvent =
+          std::unique_ptr<ur_event_handle_t_>(ur_event_handle_t_::makeNative(
+              UR_COMMAND_KERNEL_LAUNCH, hQueue, CuStream, StreamToken));
+      UR_CHECK_ERROR(RetImplEvent->start());
+    }
+
+    auto &ArgIndices = hKernel->getArgIndices();
+
+    CUlaunchConfig launch_config;
+    launch_config.gridDimX = BlocksPerGrid[0];
+    launch_config.gridDimY = BlocksPerGrid[1];
+    launch_config.gridDimZ = BlocksPerGrid[2];
+    launch_config.blockDimX = ThreadsPerBlock[0];
+    launch_config.blockDimY = ThreadsPerBlock[1];
+    launch_config.blockDimZ = ThreadsPerBlock[2];
+
+    launch_config.sharedMemBytes = LocalSize;
+    launch_config.hStream = CuStream;
+    launch_config.attrs = &launch_attribute[0];
+    launch_config.numAttrs = numAttrsInLaunchAttrList;
+
+    UR_CHECK_ERROR(cuLaunchKernelEx(&launch_config, CuFunc,
+                                    const_cast<void **>(ArgIndices.data()),
+                                    nullptr));
+
+    if (LocalSize != 0)
+      hKernel->clearLocalSize();
+
+    if (phEvent) {
+      UR_CHECK_ERROR(RetImplEvent->record());
+      *phEvent = RetImplEvent.release();
+    }
+
+  } catch (ur_result_t Err) {
+    Result = Err;
+  }
+  return Result;
+}
+
+
 /// Set parameters for general 3D memory copy.
 /// If the source and/or destination is on the device, SrcPtr and/or DstPtr
 /// must be a pointer to a CUdeviceptr

source/adapters/cuda/launch_attributes.cpp

@@ -408,6 +408,7 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetEnqueueExpProcAddrTable(
   pDdiTable->pfnCooperativeKernelLaunchExp =
       urEnqueueCooperativeKernelLaunchExp;
   pDdiTable->pfnTimestampRecordingExp = urEnqueueTimestampRecordingExp;
+  pDdiTable->pfnKernelLaunchCustomExp = urEnqueueKernelLaunchCustomExp;
 
   return UR_RESULT_SUCCESS;
 }