[SYCL][Graph] Implement graph enqueue for CUDA backend (#305)

* [SYCL][Graph] Implement graph enqueue for CUDA backend

Implements urCommandBufferAppendKernelLaunchExp function to append kernel to a command buffer.
Implements urCommandBufferAppendKernelLaunchExp function to enqueue a command buffer to a cuda stream.
Enable all tests supported by the current state of the Cuda backend.

* [SYCL][Graph] Refactors the code using helper to set Kernel params

* [SYCL][Graph] Adds exception catching

* [SYCL][Graph] Reformats paramaters of internal functions in PascalCase

* [SYCL][Graph] Fixes command buffer enqueuing issue in multithreaded environment.

Sets a ScopedContext before enqueuing command buffer.

* [SYCL][Graph] Adds missing exception catch

sycl/plugins/unified_runtime/ur/adapters/cuda/command_buffer.cpp

@@ -8,13 +8,16 @@
 
 #include "command_buffer.hpp"
 #include "common.hpp"
-
-/// Stub implementations of UR experimental feature command-buffers
+#include "enqueue.hpp"
+#include "event.hpp"
+#include "kernel.hpp"
+#include "memory.hpp"
+#include "queue.hpp"
 
 ur_exp_command_buffer_handle_t_::ur_exp_command_buffer_handle_t_(
  ur_context_handle_t hContext, ur_device_handle_t hDevice)
  : Context(hContext),
- Device(hDevice), cudaGraph{nullptr}, cudaGraphExec{nullptr}, RefCount{1} {
+ Device(hDevice), CudaGraph{nullptr}, CudaGraphExec{nullptr}, RefCount{1} {
  urContextRetain(hContext);
  urDeviceRetain(hDevice);
 }
@@ -29,10 +32,43 @@ ur_exp_command_buffer_handle_t_::~ur_exp_command_buffer_handle_t_() {
  urDeviceRelease(Device);
 
  // Release the memory allocated to the CudaGraph
- cuGraphDestroy(cudaGraph);
+ cuGraphDestroy(CudaGraph);
 
  // Release the memory allocated to the CudaGraphExec
- cuGraphExecDestroy(cudaGraphExec);
+ cuGraphExecDestroy(CudaGraphExec);
+}
+
+/// Helper function for finding the Cuda Nodes associated with the
+/// commands in a command-buffer, each event is pointed to by a sync-point in
+/// the wait list.
+///
+/// @param[in] CommandBuffer to lookup the events from.
+/// @param[in] NumSyncPointsInWaitList Length of \p SyncPointWaitList.
+/// @param[in] SyncPointWaitList List of sync points in \p CommandBuffer
+/// to find the events for.
+/// @param[out] CuNodesList Return parameter for the Cuda Nodes associated with
+/// each sync-point in \p SyncPointWaitList.
+///
+/// @return UR_RESULT_SUCCESS or an error code on failure
+static ur_result_t getNodesFromSyncPoints(
+ const ur_exp_command_buffer_handle_t &CommandBuffer,
+ size_t NumSyncPointsInWaitList,
+ const ur_exp_command_buffer_sync_point_t *SyncPointWaitList,
+ std::vector<CUgraphNode> &CuNodesList) {
+ // Map of ur_exp_command_buffer_sync_point_t to ur_event_handle_t defining
+ // the event associated with each sync-point
+ auto SyncPoints = CommandBuffer->SyncPoints;
+
+ // For each sync-point add associated L0 event to the return list.
+ for (size_t i = 0; i < NumSyncPointsInWaitList; i++) {
+ if (auto NodeHandle = SyncPoints.find(SyncPointWaitList[i]);
+ NodeHandle != SyncPoints.end()) {
+ CuNodesList.push_back(*NodeHandle->second.get());
+ } else {
+ return UR_RESULT_ERROR_INVALID_VALUE;
+ }
+ }
+ return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferCreateExp(
@@ -51,7 +87,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferCreateExp(
 
  auto RetCommandBuffer = *hCommandBuffer;
  try {
- UR_CHECK_ERROR(cuGraphCreate(&RetCommandBuffer->cudaGraph, 0));
+ UR_CHECK_ERROR(cuGraphCreate(&RetCommandBuffer->CudaGraph, 0));
  } catch (...) {
  return UR_RESULT_ERROR_OUT_OF_RESOURCES;
  }
@@ -77,8 +113,8 @@ urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
 UR_APIEXPORT ur_result_t UR_APICALL
 urCommandBufferFinalizeExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
  try {
- UR_CHECK_ERROR(cuGraphInstantiate(&hCommandBuffer->cudaGraphExec,
- hCommandBuffer->cudaGraph, 0));
+ UR_CHECK_ERROR(cuGraphInstantiate(&hCommandBuffer->CudaGraphExec,
+ hCommandBuffer->CudaGraph, 0));
  } catch (...) {
  return UR_RESULT_ERROR_UNKNOWN;
  }
@@ -92,19 +128,82 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferAppendKernelLaunchExp(
  uint32_t numSyncPointsInWaitList,
  const ur_exp_command_buffer_sync_point_t *pSyncPointWaitList,
  ur_exp_command_buffer_sync_point_t *pSyncPoint) {
- (void)hCommandBuffer;
- (void)hKernel;
- (void)workDim;
- (void)pGlobalWorkOffset;
- (void)pGlobalWorkSize;
- (void)pLocalWorkSize;
- (void)numSyncPointsInWaitList;
- (void)pSyncPointWaitList;
- (void)pSyncPoint;
+ // Preconditions
+ UR_ASSERT(hCommandBuffer->Context == 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);
+
+ ur_result_t Result = UR_RESULT_SUCCESS;
+ CUgraphNode GraphNode;
+
+ std::vector<CUgraphNode> DepsList;
+ UR_CALL(getNodesFromSyncPoints(hCommandBuffer, numSyncPointsInWaitList,
+ pSyncPointWaitList, DepsList));
+
+ if (*pGlobalWorkSize == 0) {
+ try {
+ // Create a empty node if the kernel worload size is zero
+ Result = UR_CHECK_ERROR(
+ cuGraphAddEmptyNode(&GraphNode, hCommandBuffer->CudaGraph,
+ DepsList.data(), DepsList.size()));
+
+ // Get sync point and register the event with it.
+ *pSyncPoint = hCommandBuffer->GetNextSyncPoint();
+ hCommandBuffer->RegisterSyncPoint(
+ *pSyncPoint, std::make_shared<CUgraphNode>(GraphNode));
+ } catch (ur_result_t Err) {
+ Result = Err;
+ }
+ return Result;
+ }
 
- detail::ur::die("Experimental Command-buffer feature is not "
- "implemented for CUDA adapter.");
- return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+ // 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();
+ CUfunction CuFunc = hKernel->get();
+
+ if ((Result = setKernelParams(
+ hCommandBuffer->Context, hCommandBuffer->Device, workDim,
+ pGlobalWorkOffset, pGlobalWorkSize, pLocalWorkSize, hKernel, CuFunc,
+ ThreadsPerBlock, BlocksPerGrid)) != UR_RESULT_SUCCESS) {
+ return Result;
+ }
+
+ try {
+ // Set node param structure with the kernel related data
+ auto &ArgIndices = hKernel->getArgIndices();
+ CUDA_KERNEL_NODE_PARAMS nodeParams;
+ nodeParams.func = CuFunc;
+ nodeParams.gridDimX = BlocksPerGrid[0];
+ nodeParams.gridDimY = BlocksPerGrid[1];
+ nodeParams.gridDimZ = BlocksPerGrid[2];
+ nodeParams.blockDimX = ThreadsPerBlock[0];
+ nodeParams.blockDimY = ThreadsPerBlock[1];
+ nodeParams.blockDimZ = ThreadsPerBlock[2];
+ nodeParams.sharedMemBytes = LocalSize;
+ nodeParams.kernelParams = const_cast<void **>(ArgIndices.data());
+ nodeParams.extra = nullptr;
+
+ // Create and add an new kernel node to the Cuda graph
+ Result = UR_CHECK_ERROR(
+ cuGraphAddKernelNode(&GraphNode, hCommandBuffer->CudaGraph,
+ DepsList.data(), DepsList.size(), &nodeParams));
+
+ if (LocalSize != 0)
+ hKernel->clearLocalSize();
+
+ // Get sync point and register the event with it.
+ *pSyncPoint = hCommandBuffer->GetNextSyncPoint();
+ hCommandBuffer->RegisterSyncPoint(*pSyncPoint,
+ std::make_shared<CUgraphNode>(GraphNode));
+ } catch (ur_result_t Err) {
+ Result = Err;
+ }
+ return Result;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferAppendMemcpyUSMExp(
@@ -275,13 +374,38 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferEnqueueExp(
  ur_exp_command_buffer_handle_t hCommandBuffer, ur_queue_handle_t hQueue,
  uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
  ur_event_handle_t *phEvent) {
- (void)hCommandBuffer;
- (void)hQueue;
- (void)numEventsInWaitList;
- (void)phEventWaitList;
- (void)phEvent;
+ ur_result_t Result = UR_RESULT_SUCCESS;
 
- detail::ur::die("Experimental Command-buffer feature is not "
- "implemented for CUDA adapter.");
- return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+ try {
+ std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
+ ScopedContext Active(hQueue->getContext());
+ uint32_t StreamToken;
+ ur_stream_guard_ Guard;
+ CUstream CuStream = hQueue->getNextComputeStream(
+ numEventsInWaitList, phEventWaitList, Guard, &StreamToken);
+
+ if ((Result = enqueueEventsWait(hQueue, CuStream, numEventsInWaitList,
+ phEventWaitList)) != UR_RESULT_SUCCESS) {
+ return Result;
+ }
+
+ if (phEvent) {
+ RetImplEvent =
+ std::unique_ptr<ur_event_handle_t_>(ur_event_handle_t_::makeNative(
+ UR_COMMAND_KERNEL_LAUNCH, hQueue, CuStream, StreamToken));
+ RetImplEvent->start();
+ }
+
+ // Launch graph
+ Result =
+ UR_CHECK_ERROR(cuGraphLaunch(hCommandBuffer->CudaGraphExec, CuStream));
+
+ if (phEvent) {
+ Result = RetImplEvent->record();
+ *phEvent = RetImplEvent.release();
+ }
+ } catch (ur_result_t Err) {
+ Result = Err;
+ }
+ return Result;
 }