[SYCL][Graph] Modified the adapters such that it is valid to call release on...

sycl/source/detail/graph_impl.cpp

@@ -977,11 +977,6 @@ exec_graph_impl::~exec_graph_impl() {
     const sycl::detail::AdapterPtr &Adapter =
         sycl::detail::getSyclObjImpl(MContext)->getAdapter();
     MSchedule.clear();
-    // We need to wait on all command buffer executions before we can release
-    // them.
-    for (auto &Event : MExecutionEvents) {
-      Event->wait(Event);
-    }
 
     // Clean up any graph-owned allocations that were allocated
     MGraphImpl->getMemPool().deallocateAndUnmapAll();

unified-runtime/scripts/core/EXP-COMMAND-BUFFER.rst

@@ -563,3 +563,4 @@ Contributors
 * Maxime France-Pillois `[email protected] <[email protected]>`_
 * Aaron Greig `[email protected] <[email protected]>`_
 * Fábio Mestre `[email protected] <[email protected]>`_
+* Konrad Kusiak `[email protected] <[email protected]>`_

unified-runtime/scripts/core/exp-command-buffer.yml

@@ -310,7 +310,7 @@ returns:
     - $X_RESULT_ERROR_OUT_OF_HOST_MEMORY
 --- #--------------------------------------------------------------------------
 type: function
-desc: "Decrement the command-buffer object's reference count and delete the command-buffer object if the reference count becomes zero."
+desc: "Decrement the command-buffer object's reference count and delete the command-buffer object if the reference count becomes zero. It is legal to call the entry-point while `hCommandBuffer` is still executing, which will block on completion if the reference count of `hCommandBuffer` becomes zero."
 class: $xCommandBuffer
 name: ReleaseExp
 params:

unified-runtime/source/adapters/cuda/command_buffer.cpp

@@ -387,6 +387,9 @@ urCommandBufferRetainExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
 UR_APIEXPORT ur_result_t UR_APICALL
 urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
   if (hCommandBuffer->decrementReferenceCount() == 0) {
+    if (hCommandBuffer->CurrentExecution) {
+      UR_CHECK_ERROR(hCommandBuffer->CurrentExecution->wait());
+    }
     // Ref count has reached zero, release of created commands
     for (auto &Command : hCommandBuffer->CommandHandles) {
       commandHandleDestroy(Command);
@@ -1160,18 +1163,17 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueCommandBufferExp(
   UR_CHECK_ERROR(enqueueEventsWait(hQueue, CuStream, numEventsInWaitList,
                                    phEventWaitList));
 
-  if (phEvent) {
-    RetImplEvent = std::make_unique<ur_event_handle_t_>(
-        UR_COMMAND_ENQUEUE_COMMAND_BUFFER_EXP, hQueue, CuStream, StreamToken);
-    UR_CHECK_ERROR(RetImplEvent->start());
-  }
+  RetImplEvent = std::make_unique<ur_event_handle_t_>(
+      UR_COMMAND_ENQUEUE_COMMAND_BUFFER_EXP, hQueue, CuStream, StreamToken);
+  UR_CHECK_ERROR(RetImplEvent->start());
 
   // Launch graph
   UR_CHECK_ERROR(cuGraphLaunch(hCommandBuffer->CudaGraphExec, CuStream));
 
+  UR_CHECK_ERROR(RetImplEvent->record());
+  hCommandBuffer->CurrentExecution = RetImplEvent.release();
   if (phEvent) {
-    UR_CHECK_ERROR(RetImplEvent->record());
-    *phEvent = RetImplEvent.release();
+    *phEvent = hCommandBuffer->CurrentExecution;
   }
   return UR_RESULT_SUCCESS;
 } catch (ur_result_t Err) {

unified-runtime/source/adapters/cuda/command_buffer.hpp

@@ -192,6 +192,8 @@ struct ur_exp_command_buffer_handle_t_ : ur::cuda::handle_base {
   // Atomic variable counting the number of reference to this command_buffer
   // using std::atomic prevents data race when incrementing/decrementing.
   std::atomic_uint32_t RefCount;
+  // Track the event of the current graph execution.
+  ur_event_handle_t CurrentExecution = nullptr;
 
   // Ordered map of sync_points to ur_events, so that we can find the last
   // node added to an in-order command-buffer.

unified-runtime/source/adapters/hip/command_buffer.cpp

@@ -273,6 +273,9 @@ urCommandBufferRetainExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
 UR_APIEXPORT ur_result_t UR_APICALL
 urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
   if (hCommandBuffer->decrementReferenceCount() == 0) {
+    if (hCommandBuffer->CurrentExecution) {
+      UR_CHECK_ERROR(hCommandBuffer->CurrentExecution->wait());
+    }
     delete hCommandBuffer;
   }
   return UR_RESULT_SUCCESS;
@@ -798,19 +801,17 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueCommandBufferExp(
     UR_CHECK_ERROR(enqueueEventsWait(hQueue, HIPStream, numEventsInWaitList,
                                      phEventWaitList));
 
-    if (phEvent) {
-      RetImplEvent = std::make_unique<ur_event_handle_t_>(
-          UR_COMMAND_ENQUEUE_COMMAND_BUFFER_EXP, hQueue, HIPStream,
-          StreamToken);
-      UR_CHECK_ERROR(RetImplEvent->start());
-    }
+    RetImplEvent = std::make_unique<ur_event_handle_t_>(
+        UR_COMMAND_ENQUEUE_COMMAND_BUFFER_EXP, hQueue, HIPStream, StreamToken);
+    UR_CHECK_ERROR(RetImplEvent->start());
 
     // Launch graph
     UR_CHECK_ERROR(hipGraphLaunch(hCommandBuffer->HIPGraphExec, HIPStream));
 
+    UR_CHECK_ERROR(RetImplEvent->record());
+    hCommandBuffer->CurrentExecution = RetImplEvent.release();
     if (phEvent) {
-      UR_CHECK_ERROR(RetImplEvent->record());
-      *phEvent = RetImplEvent.release();
+      *phEvent = hCommandBuffer->CurrentExecution;
     }
   } catch (ur_result_t Err) {
     return Err;

unified-runtime/source/adapters/hip/command_buffer.hpp

@@ -128,6 +128,8 @@ struct ur_exp_command_buffer_handle_t_ : ur::hip::handle_base {
   // Atomic variable counting the number of reference to this command_buffer
   // using std::atomic prevents data race when incrementing/decrementing.
   std::atomic_uint32_t RefCount;
+  // Track the event of the current graph execution.
+  ur_event_handle_t CurrentExecution = nullptr;
 
   // Ordered map of sync_points to ur_events
   std::map<ur_exp_command_buffer_sync_point_t, hipGraphNode_t> SyncPoints;

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

@@ -671,6 +671,25 @@ ur_result_t createMainCommandList(ur_context_handle_t Context,
   return UR_RESULT_SUCCESS;
 }
 
+/**
+ * Waits for any ongoing executions of the command-buffer to finish.
+ * @param CommandBuffer The command-buffer to wait for.
+ * @return UR_RESULT_SUCCESS or an error code on failure
+ */
+ur_result_t
+waitForOngoingExecution(ur_exp_command_buffer_handle_t CommandBuffer) {
+
+  if (ur_event_handle_t &CurrentSubmissionEvent =
+          CommandBuffer->CurrentSubmissionEvent) {
+    ZE2UR_CALL(zeEventHostSynchronize,
+               (CurrentSubmissionEvent->ZeEvent, UINT64_MAX));
+    UR_CALL(urEventReleaseInternal(CurrentSubmissionEvent));
+    CurrentSubmissionEvent = nullptr;
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
 /**
  * Checks whether the command-buffer can be constructed using in order
  * command-lists.
@@ -830,6 +849,7 @@ urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t CommandBuffer) {
   if (!CommandBuffer->RefCount.decrementAndTest())
     return UR_RESULT_SUCCESS;
 
+  UR_CALL(waitForOngoingExecution(CommandBuffer));
   CommandBuffer->cleanupCommandBufferResources();
   delete CommandBuffer;
   return UR_RESULT_SUCCESS;
@@ -1442,25 +1462,6 @@ ur_result_t getZeCommandQueue(ur_queue_handle_t Queue, bool UseCopyEngine,
   return UR_RESULT_SUCCESS;
 }
 
-/**
- * Waits for any ongoing executions of the command-buffer to finish.
- * @param CommandBuffer The command-buffer to wait for.
- * @return UR_RESULT_SUCCESS or an error code on failure
- */
-ur_result_t
-waitForOngoingExecution(ur_exp_command_buffer_handle_t CommandBuffer) {
-
-  if (ur_event_handle_t &CurrentSubmissionEvent =
-          CommandBuffer->CurrentSubmissionEvent) {
-    ZE2UR_CALL(zeEventHostSynchronize,
-               (CurrentSubmissionEvent->ZeEvent, UINT64_MAX));
-    UR_CALL(urEventReleaseInternal(CurrentSubmissionEvent));
-    CurrentSubmissionEvent = nullptr;
-  }
-
-  return UR_RESULT_SUCCESS;
-}
-
 /**
  * Waits for the all the dependencies of the command-buffer
  * @param[in] CommandBuffer The command-buffer.

unified-runtime/source/adapters/opencl/command_buffer.cpp

@@ -115,6 +115,10 @@ urCommandBufferRetainExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
 UR_APIEXPORT ur_result_t UR_APICALL
 urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
   if (hCommandBuffer->decrementReferenceCount() == 0) {
+    if (hCommandBuffer->LastSubmission) {
+      cl_int RetErr = clWaitForEvents(1, &(hCommandBuffer->LastSubmission));
+      CL_RETURN_ON_FAILURE(RetErr);
+    }
     delete hCommandBuffer;
   }
 

unified-runtime/test/conformance/exp_command_buffer/enqueue.cpp

@@ -135,3 +135,16 @@ TEST_P(urEnqueueCommandBufferExpTest, SerializeOutofOrderQueue) {
     ASSERT_EQ(reference, Output[i]);
   }
 }
+
+// Tests releasing command-buffer while it is still executing relying
+// on synchronization during urCommandBufferReleaseExp call.
+TEST_P(urEnqueueCommandBufferExpTest, EnqueueAndRelease) {
+  ASSERT_SUCCESS(urEnqueueCommandBufferExp(out_of_order_queue, cmd_buf_handle,
+                                           0, nullptr, nullptr));
+
+  // Release the command buffer without explicitly waiting beforehand
+  EXPECT_SUCCESS(urCommandBufferReleaseExp(cmd_buf_handle));
+
+  // Wait before exiting
+  ASSERT_SUCCESS(urQueueFinish(out_of_order_queue));
+}