Get rid of custom cuda_stream_ref (#526)

This PR fixes merge conflicts between #517 and #525.
NVIDIA · Jul 3, 2024 · 2ad911f · 2ad911f
1 parent dff0112
commit 2ad911f
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/include/cuco/detail/open_addressing/open_addressing_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_impl.cuh
@@ -415,15 +415,15 @@ class open_addressing_impl {
                              FoundIt found_begin,
                              InsertedIt inserted_begin,
                              Ref container_ref,
-                             cuda_stream_ref stream) noexcept
+                             cuda::stream_ref stream) noexcept
   {
     auto const num_keys = cuco::detail::distance(first, last);
     if (num_keys == 0) { return; }
 
     auto const grid_size = cuco::detail::grid_size(num_keys, cg_size);
 
     detail::insert_and_find<cg_size, cuco::detail::default_block_size()>
-      <<<grid_size, cuco::detail::default_block_size(), 0, stream>>>(
+      <<<grid_size, cuco::detail::default_block_size(), 0, stream.get()>>>(
         first, num_keys, found_begin, inserted_begin, container_ref);
   }
 

diff --git a/include/cuco/detail/static_set/static_set.inl b/include/cuco/detail/static_set/static_set.inl
@@ -187,10 +187,10 @@ void static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>
   InputIt last,
   FoundIt found_begin,
   InsertedIt inserted_begin,
-  cuda_stream_ref stream)
+  cuda::stream_ref stream)
 {
   insert_and_find_async(first, last, found_begin, inserted_begin, stream);
-  stream.synchronize();
+  stream.wait();
 }
 
 template <class Key,
@@ -206,7 +206,7 @@ void static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>
                         InputIt last,
                         FoundIt found_begin,
                         InsertedIt inserted_begin,
-                        cuda_stream_ref stream) noexcept
+                        cuda::stream_ref stream) noexcept
 {
   impl_->insert_and_find_async(
     first, last, found_begin, inserted_begin, ref(op::insert_and_find), stream);

diff --git a/include/cuco/static_set.cuh b/include/cuco/static_set.cuh
@@ -365,7 +365,7 @@ class static_set {
                              InputIt last,
                              FoundIt found_begin,
                              InsertedIt inserted_begin,
-                             cuda_stream_ref stream = {}) noexcept;
+                             cuda::stream_ref stream = {}) noexcept;
 
   /**
    * @brief Inserts all elements in the range `[first, last)`.
@@ -393,7 +393,7 @@ class static_set {
                        InputIt last,
                        FoundIt found_begin,
                        InsertedIt inserted_begin,
-                       cuda_stream_ref stream = {});
+                       cuda::stream_ref stream = {});
 
   /**
    * @brief Erases keys in the range `[first, last)`.