[SYCL][CUDA][PI] Improve performance of event synchronization

sycl/plugins/cuda/pi_cuda.cpp

@@ -300,7 +300,11 @@ pi_result enqueueEventsWait(pi_queue command_queue, CUstream stream,
     auto result = forLatestEvents(
         event_wait_list, num_events_in_wait_list,
         [stream](pi_event event) -> pi_result {
-          return PI_CHECK_ERROR(cuStreamWaitEvent(stream, event->get(), 0));
+          if (event->get_stream() == stream) {
+            return PI_SUCCESS;
+          } else {
+            return PI_CHECK_ERROR(cuStreamWaitEvent(stream, event->get(), 0));
+          }
         });
 
     if (result != PI_SUCCESS) {
@@ -367,18 +371,47 @@ pi_result cuda_piEventRetain(pi_event event);
 
 /// \endcond
 
-CUstream _pi_queue::get_next_compute_stream() {
-  if (num_compute_streams_ < compute_streams_.size()) {
-    // the check above is for performance - so as not to lock mutex every time
-    std::lock_guard<std::mutex> guard(compute_stream_mutex_);
-    // The second check is done after mutex is locked so other threads can not
-    // change num_compute_streams_ after that
+CUstream _pi_queue::get_next_compute_stream(pi_uint32 *stream_token) {
+  pi_uint32 stream_i;
+  while (true) {
     if (num_compute_streams_ < compute_streams_.size()) {
-      PI_CHECK_ERROR(
-          cuStreamCreate(&compute_streams_[num_compute_streams_++], flags_));
+      // the check above is for performance - so as not to lock mutex every time
+      std::lock_guard<std::mutex> guard(compute_stream_mutex_);
+      // The second check is done after mutex is locked so other threads can not
+      // change num_compute_streams_ after that
+      if (num_compute_streams_ < compute_streams_.size()) {
+        PI_CHECK_ERROR(
+            cuStreamCreate(&compute_streams_[num_compute_streams_++], flags_));
+      }
+    }
+    stream_i = compute_stream_idx_++;
+    if (delay_compute_[stream_i % compute_streams_.size()]) {
+      delay_compute_[stream_i % compute_streams_.size()] = 0;
+    } else {
+      break;
     }
   }
-  return compute_streams_[compute_stream_idx_++ % compute_streams_.size()];
+  if (stream_token) {
+    *stream_token = stream_i;
+  }
+  return compute_streams_[stream_i % compute_streams_.size()];
+}
+
+CUstream getComputeStream(pi_queue queue, pi_uint32 num_events_in_wait_list,
+                          const pi_event *event_wait_list,
+                          pi_uint32 *stream_token = nullptr) {
+  for (pi_uint32 i = 0; i < num_events_in_wait_list; i++) {
+    pi_uint32 token = event_wait_list[i]->get_stream_token();
+    if (event_wait_list[i]->get_queue() == queue &&
+        queue->is_last_command(token)) {
+      queue->delay_stream(token);
+      if (stream_token) {
+        *stream_token = token;
+      }
+      return event_wait_list[i]->get_stream();
+    }
+  }
+  return queue->get_next_compute_stream(stream_token);
 }
 
 CUstream _pi_queue::get_next_transfer_stream() {
@@ -399,9 +432,10 @@ CUstream _pi_queue::get_next_transfer_stream() {
 }
 
 _pi_event::_pi_event(pi_command_type type, pi_context context, pi_queue queue,
-                     CUstream stream)
+                     CUstream stream, pi_uint32 stream_token)
     : commandType_{type}, refCount_{1}, hasBeenWaitedOn_{false},
-      isRecorded_{false}, isStarted_{false}, evEnd_{nullptr}, evStart_{nullptr},
+      isRecorded_{false}, isStarted_{false},
+      streamToken_{stream_token}, evEnd_{nullptr}, evStart_{nullptr},
       evQueued_{nullptr}, queue_{queue}, stream_{stream}, context_{context} {
 
   bool profilingEnabled = queue_->properties_ & PI_QUEUE_PROFILING_ENABLE;
@@ -2837,7 +2871,9 @@ pi_result cuda_piEnqueueKernelLaunch(
 
     std::unique_ptr<_pi_event> retImplEv{nullptr};
 
-    CUstream cuStream = command_queue->get_next_compute_stream();
+    pi_uint32 stream_token;
+    CUstream cuStream = getComputeStream(command_queue, num_events_in_wait_list,
+                                         event_wait_list, &stream_token);
     CUfunction cuFunc = kernel->get();
 
     retError = enqueueEventsWait(command_queue, cuStream,
@@ -2862,8 +2898,9 @@ pi_result cuda_piEnqueueKernelLaunch(
     auto &argIndices = kernel->get_arg_indices();
 
     if (event) {
-      retImplEv = std::unique_ptr<_pi_event>(_pi_event::make_native(
-          PI_COMMAND_TYPE_NDRANGE_KERNEL, command_queue, cuStream));
+      retImplEv = std::unique_ptr<_pi_event>(
+          _pi_event::make_native(PI_COMMAND_TYPE_NDRANGE_KERNEL, command_queue,
+                                 cuStream, stream_token));
       retImplEv->start();
     }
 
@@ -3707,8 +3744,12 @@ pi_result cuda_piEnqueueEventsWaitWithBarrier(pi_queue command_queue,
     }
 
     if (event) {
+      pi_uint32 stream_token;
+      CUstream cuStream =
+          getComputeStream(command_queue, num_events_in_wait_list,
+                           event_wait_list, &stream_token);
       *event = _pi_event::make_native(PI_COMMAND_TYPE_MARKER, command_queue,
-                                      command_queue->get_next_compute_stream());
+                                      cuStream, stream_token);
       (*event)->start();
       (*event)->record();
     }
@@ -4766,12 +4807,14 @@ pi_result cuda_piextUSMEnqueueMemset(pi_queue queue, void *ptr, pi_int32 value,
 
   try {
     ScopedContext active(queue->get_context());
-    CUstream cuStream = queue->get_next_compute_stream();
+    pi_uint32 stream_token;
+    CUstream cuStream = getComputeStream(queue, num_events_in_waitlist,
+                                         events_waitlist, &stream_token);
     result = enqueueEventsWait(queue, cuStream, num_events_in_waitlist,
                                events_waitlist);
     if (event) {
       event_ptr = std::unique_ptr<_pi_event>(_pi_event::make_native(
-          PI_COMMAND_TYPE_MEM_BUFFER_FILL, queue, cuStream));
+          PI_COMMAND_TYPE_MEM_BUFFER_FILL, queue, cuStream, stream_token));
       event_ptr->start();
     }
     result = PI_CHECK_ERROR(cuMemsetD8Async(

sycl/plugins/cuda/pi_cuda.hpp

@@ -382,6 +382,7 @@ struct _pi_queue {
 
   std::vector<native_type> compute_streams_;
   std::vector<native_type> transfer_streams_;
+  std::vector<char> delay_compute_;
   _pi_context *context_;
   _pi_device *device_;
   pi_queue_properties properties_;
@@ -400,7 +401,8 @@ struct _pi_queue {
             _pi_device *device, pi_queue_properties properties,
             unsigned int flags)
       : compute_streams_{std::move(compute_streams)},
-        transfer_streams_{std::move(transfer_streams)}, context_{context},
+        transfer_streams_{std::move(transfer_streams)},
+        delay_compute_(compute_streams_.size(), 0), context_{context},
         device_{device}, properties_{properties}, refCount_{1}, eventCount_{0},
         compute_stream_idx_{0}, transfer_stream_idx_{0},
         num_compute_streams_{0}, num_transfer_streams_{0}, flags_(flags) {
@@ -415,10 +417,21 @@ struct _pi_queue {
 
   // get_next_compute/transfer_stream() functions return streams from
   // appropriate pools in round-robin fashion
-  native_type get_next_compute_stream();
+  native_type get_next_compute_stream(pi_uint32 *stream_token = nullptr);
   native_type get_next_transfer_stream();
   native_type get() { return get_next_compute_stream(); };
 
+  bool is_last_command(pi_uint32 stream_token) {
+    if (stream_token == std::numeric_limits<pi_uint32>::max()) {
+      return true;
+    }
+    return (compute_stream_idx_ - stream_token) <= compute_streams_.size();
+  }
+
+  void delay_stream(pi_uint32 stream_token) {
+    delay_compute_[stream_token % delay_compute_.size()] = 1;
+  }
+
   template <typename T> void for_each_stream(T &&f) {
     {
       std::lock_guard<std::mutex> compute_guard(compute_stream_mutex_);
@@ -473,6 +486,8 @@ struct _pi_event {
 
   CUstream get_stream() const noexcept { return stream_; }
 
+  pi_uint32 get_stream_token() const noexcept { return streamToken_; }
+
   pi_command_type get_command_type() const noexcept { return commandType_; }
 
   pi_uint32 get_reference_count() const noexcept { return refCount_; }
@@ -516,9 +531,11 @@ struct _pi_event {
   pi_uint64 get_end_time() const;
 
   // construct a native CUDA. This maps closely to the underlying CUDA event.
-  static pi_event make_native(pi_command_type type, pi_queue queue,
-                              CUstream stream) {
-    return new _pi_event(type, queue->get_context(), queue, stream);
+  static pi_event
+  make_native(pi_command_type type, pi_queue queue, CUstream stream,
+              pi_uint32 stream_token = std::numeric_limits<pi_uint32>::max()) {
+    return new _pi_event(type, queue->get_context(), queue, stream,
+                         stream_token);
   }
 
   pi_result release();
@@ -529,7 +546,7 @@ struct _pi_event {
   // This constructor is private to force programmers to use the make_native /
   // make_user static members in order to create a pi_event for CUDA.
   _pi_event(pi_command_type type, pi_context context, pi_queue queue,
-            CUstream stream);
+            CUstream stream, pi_uint32 stream_token);
 
   pi_command_type commandType_; // The type of command associated with event.
 
@@ -545,6 +562,7 @@ struct _pi_event {
                     // PI event has started or not
                     //
 
+  pi_uint32 streamToken_;
   pi_uint32 eventId_; // Queue identifier of the event.
 
   native_type evEnd_; // CUDA event handle. If this _pi_event represents a user