Fix CTAD for static_set

NVIDIA · Dec 14, 2023 · aa16b70 · aa16b70
1 parent 6eb8cc3
commit aa16b70
Show file tree

Hide file tree

Showing 12 changed files with 84 additions and 33 deletions.
diff --git a/benchmarks/hash_table/static_set/insert_bench.cu b/benchmarks/hash_table/static_set/insert_bench.cu
@@ -48,7 +48,7 @@ void static_set_insert(nvbench::state& state, nvbench::type_list<Key, Dist>)
   state.exec(nvbench::exec_tag::sync | nvbench::exec_tag::timer,
              [&](nvbench::launch& launch, auto& timer) {
                cuco::experimental::static_set<Key> set{
-                 size, cuco::empty_key<Key>{-1}, {}, {}, {}, {launch.get_stream()}};
+                 size, cuco::empty_key<Key>{-1}, {}, {}, {}, {}, {}, {launch.get_stream()}};
 
                timer.start();
                set.insert(keys.begin(), keys.end(), {launch.get_stream()});

diff --git a/examples/static_set/device_ref_example.cu b/examples/static_set/device_ref_example.cu
@@ -83,10 +83,8 @@ int main(void)
   auto constexpr load_factor = 0.5;
   std::size_t const capacity = std::ceil(num_keys / load_factor);
 
-  using set_type = cuco::experimental::static_set<Key>;
-
   // Constructs a hash set with at least "capacity" slots using -1 as the empty key sentinel.
-  set_type set{capacity, cuco::empty_key{empty_key_sentinel}};
+  cuco::experimental::static_set<Key> set{capacity, cuco::empty_key{empty_key_sentinel}};
 
   // Create a sequence of keys {0, 1, 2, .., i}
   thrust::device_vector<Key> keys(num_keys);

diff --git a/examples/static_set/device_subsets_example.cu b/examples/static_set/device_subsets_example.cu
@@ -169,7 +169,7 @@ int main()
   for (std::size_t i = 0; i < num; ++i) {
     storage_ref_type storage_ref{valid_sizes[i], set_storage.data() + offsets[i]};
     set_refs.emplace_back(
-      ref_type{cuco::empty_key<key_type>{empty_key_sentinel}, {}, {}, storage_ref});
+      ref_type{cuco::empty_key<key_type>{empty_key_sentinel}, {}, {}, {}, storage_ref});
   }
 
   thrust::device_vector<ref_type> d_set_refs(set_refs);

diff --git a/include/cuco/cuda_thread_scope.cuh b/include/cuco/cuda_thread_scope.cuh
@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuda/std/atomic>  // cuda::thread_scope
+
+namespace cuco {
+namespace experimental {
+
+/**
+ * @brief Strongly-typed wrapper for `cuda::thread_scope`.
+ *
+ * @tparam Scope `cuda::thread_scope` to be wrapped
+ */
+template <cuda::thread_scope Scope>
+struct cuda_thread_scope {
+  /**
+   * @brief Implicit conversion to `cuda::thread_scope`.
+   *
+   * @return The wrapped `cuda::thread_scope`
+   */
+  __host__ __device__ constexpr operator cuda::thread_scope() const noexcept { return Scope; }
+};
+
+// alias definitions
+inline constexpr auto thread_scope_system = cuda_thread_scope<cuda::thread_scope_system>{};
+inline constexpr auto thread_scope_device = cuda_thread_scope<cuda::thread_scope_device>{};
+inline constexpr auto thread_scope_block  = cuda_thread_scope<cuda::thread_scope_block>{};
+inline constexpr auto thread_scope_thread = cuda_thread_scope<cuda::thread_scope_thread>{};
+
+}  // namespace experimental
+}  // namespace cuco
diff --git a/include/cuco/detail/static_set/static_set.inl b/include/cuco/detail/static_set/static_set.inl
@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-#include <cuco/cuda_stream_ref.hpp>
 #include <cuco/detail/static_set/functors.cuh>
 #include <cuco/detail/static_set/kernels.cuh>
 #include <cuco/detail/utility/cuda.hpp>
@@ -39,6 +38,8 @@ constexpr static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Sto
   empty_key<Key> empty_key_sentinel,
   KeyEqual const& pred,
   ProbingScheme const& probing_scheme,
+  cuda_thread_scope<Scope>,
+  Storage,
   Allocator const& alloc,
   cuda_stream_ref stream)
   : impl_{std::make_unique<impl_type>(
@@ -59,6 +60,8 @@ constexpr static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Sto
   empty_key<Key> empty_key_sentinel,
   KeyEqual const& pred,
   ProbingScheme const& probing_scheme,
+  cuda_thread_scope<Scope>,
+  Storage,
   Allocator const& alloc,
   cuda_stream_ref stream)
   : impl_{std::make_unique<impl_type>(n,
@@ -85,6 +88,8 @@ constexpr static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Sto
   erased_key<Key> erased_key_sentinel,
   KeyEqual const& pred,
   ProbingScheme const& probing_scheme,
+  cuda_thread_scope<Scope>,
+  Storage,
   Allocator const& alloc,
   cuda_stream_ref stream)
   : impl_{std::make_unique<impl_type>(capacity,
@@ -471,11 +476,13 @@ auto static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>
            ? ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
                                     impl_->key_eq(),
                                     impl_->probing_scheme(),
+                                    cuda_thread_scope<Scope>{},
                                     impl_->storage_ref()}
            : ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
                                     cuco::erased_key<key_type>(this->erased_key_sentinel()),
                                     impl_->key_eq(),
                                     impl_->probing_scheme(),
+                                    cuda_thread_scope<Scope>{},
                                     impl_->storage_ref()};
 }
 }  // namespace experimental

diff --git a/include/cuco/detail/static_set/static_set_ref.inl b/include/cuco/detail/static_set/static_set_ref.inl
@@ -40,6 +40,7 @@ __host__ __device__ constexpr static_set_ref<
   Operators...>::static_set_ref(cuco::empty_key<Key> empty_key_sentinel,
                                 KeyEqual const& predicate,
                                 ProbingScheme const& probing_scheme,
+                                cuda_thread_scope<Scope>,
                                 StorageRef storage_ref) noexcept
   : impl_{empty_key_sentinel, predicate, probing_scheme, storage_ref}
 {
@@ -61,6 +62,7 @@ __host__ __device__ constexpr static_set_ref<
                                 cuco::erased_key<Key> erased_key_sentinel,
                                 KeyEqual const& predicate,
                                 ProbingScheme const& probing_scheme,
+                                cuda_thread_scope<Scope>,
                                 StorageRef storage_ref) noexcept
   : impl_{empty_key_sentinel, erased_key_sentinel, predicate, probing_scheme, storage_ref}
 {

diff --git a/include/cuco/static_set.cuh b/include/cuco/static_set.cuh
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cuco/cuda_stream_ref.hpp>
+#include <cuco/cuda_thread_scope.cuh>
 #include <cuco/detail/open_addressing/open_addressing_impl.cuh>
 #include <cuco/extent.cuh>
 #include <cuco/hash_functions.cuh>
@@ -145,13 +146,17 @@ class static_set {
    * @param empty_key_sentinel The reserved key value for empty slots
    * @param pred Key equality binary predicate
    * @param probing_scheme Probing scheme
+   * @param scope The scope in which operations will be performed
+   * @param storage Kind of storage to use
    * @param alloc Allocator used for allocating device storage
    * @param stream CUDA stream used to initialize the set
    */
   constexpr static_set(Extent capacity,
                        empty_key<Key> empty_key_sentinel,
                        KeyEqual const& pred                = {},
                        ProbingScheme const& probing_scheme = {},
+                       cuda_thread_scope<Scope> scope      = {},
+                       Storage storage                     = {},
                        Allocator const& alloc              = {},
                        cuda_stream_ref stream              = {});
 
@@ -182,6 +187,8 @@ class static_set {
    * @param empty_key_sentinel The reserved key value for empty slots
    * @param pred Key equality binary predicate
    * @param probing_scheme Probing scheme
+   * @param scope The scope in which operations will be performed
+   * @param storage Kind of storage to use
    * @param alloc Allocator used for allocating device storage
    * @param stream CUDA stream used to initialize the set
    */
@@ -190,6 +197,8 @@ class static_set {
                        empty_key<Key> empty_key_sentinel,
                        KeyEqual const& pred                = {},
                        ProbingScheme const& probing_scheme = {},
+                       cuda_thread_scope<Scope> scope      = {},
+                       Storage storage                     = {},
                        Allocator const& alloc              = {},
                        cuda_stream_ref stream              = {});
 
@@ -212,6 +221,8 @@ class static_set {
    * @param erased_key_sentinel The reserved key to denote erased slots
    * @param pred Key equality binary predicate
    * @param probing_scheme Probing scheme
+   * @param scope The scope in which operations will be performed
+   * @param storage Kind of storage to use
    * @param alloc Allocator used for allocating device storage
    * @param stream CUDA stream used to initialize the set
    */
@@ -220,6 +231,8 @@ class static_set {
                        erased_key<Key> erased_key_sentinel,
                        KeyEqual const& pred                = {},
                        ProbingScheme const& probing_scheme = {},
+                       cuda_thread_scope<Scope> scope      = {},
+                       Storage storage                     = {},
                        Allocator const& alloc              = {},
                        cuda_stream_ref stream              = {});
 

diff --git a/include/cuco/static_set_ref.cuh b/include/cuco/static_set_ref.cuh
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <cuco/cuda_thread_scope.cuh>
 #include <cuco/detail/open_addressing/open_addressing_ref_impl.cuh>
 #include <cuco/hash_functions.cuh>
 #include <cuco/operator.hpp>
@@ -89,11 +90,13 @@ class static_set_ref
    * @param empty_key_sentinel Sentinel indicating empty key
    * @param predicate Key equality binary callable
    * @param probing_scheme Probing scheme
+   * @param scope The scope in which operations will be performed
    * @param storage_ref Non-owning ref of slot storage
    */
   __host__ __device__ explicit constexpr static_set_ref(cuco::empty_key<Key> empty_key_sentinel,
                                                         KeyEqual const& predicate,
                                                         ProbingScheme const& probing_scheme,
+                                                        cuda_thread_scope<Scope> scope,
                                                         StorageRef storage_ref) noexcept;
 
   /**
@@ -103,12 +106,14 @@ class static_set_ref
    * @param erased_key_sentinel Sentinel indicating erased key
    * @param predicate Key equality binary callable
    * @param probing_scheme Probing scheme
+   * @param scope The scope in which operations will be performed
    * @param storage_ref Non-owning ref of slot storage
    */
   __host__ __device__ explicit constexpr static_set_ref(cuco::empty_key<Key> empty_key_sentinel,
                                                         cuco::erased_key<Key> erased_key_sentinel,
                                                         KeyEqual const& predicate,
                                                         ProbingScheme const& probing_scheme,
+                                                        cuda_thread_scope<Scope> scope,
                                                         StorageRef storage_ref) noexcept;
 
   /**

diff --git a/tests/static_set/insert_and_find_test.cu b/tests/static_set/insert_and_find_test.cu
@@ -101,13 +101,8 @@ TEMPLATE_TEST_CASE_SIG(
     cuco::experimental::linear_probing<CGSize, cuco::default_hash_function<Key>>,
     cuco::experimental::double_hashing<CGSize, cuco::default_hash_function<Key>>>;
 
-  auto set = cuco::experimental::static_set<Key,
-                                            cuco::experimental::extent<std::size_t>,
-                                            cuda::thread_scope_device,
-                                            thrust::equal_to<Key>,
-                                            probe,
-                                            cuco::cuda_allocator<std::byte>,
-                                            cuco::experimental::storage<2>>{
-    num_keys, cuco::empty_key<Key>{-1}};
+  auto set = cuco::experimental::static_set{
+    num_keys, cuco::empty_key<Key>{-1}, {}, {}, {}, cuco::experimental::storage<2>{}};
+
   test_insert_and_find(set, num_keys);
 }
diff --git a/tests/static_set/large_input_test.cu b/tests/static_set/large_input_test.cu
@@ -70,9 +70,7 @@ TEMPLATE_TEST_CASE_SIG(
   using probe       = cuco::experimental::double_hashing<CGSize, cuco::default_hash_function<Key>>;
 
   try {
-    auto set = cuco::experimental::
-      static_set<Key, extent_type, cuda::thread_scope_device, thrust::equal_to<Key>, probe>{
-        num_keys * 2, cuco::empty_key<Key>{-1}};
+    auto set = cuco::experimental::static_set{num_keys * 2, cuco::empty_key<Key>{-1}, {}, probe{}};
 
     thrust::device_vector<bool> d_contained(num_keys);
     test_unique_sequence(set, d_contained.data().get(), num_keys);

diff --git a/tests/static_set/retrieve_all_test.cu b/tests/static_set/retrieve_all_test.cu
@@ -80,14 +80,7 @@ TEMPLATE_TEST_CASE_SIG(
     cuco::experimental::linear_probing<CGSize, cuco::default_hash_function<Key>>,
     cuco::experimental::double_hashing<CGSize, cuco::default_hash_function<Key>>>;
 
-  auto set = cuco::experimental::static_set<Key,
-                                            cuco::experimental::extent<std::size_t>,
-                                            cuda::thread_scope_device,
-                                            thrust::equal_to<Key>,
-                                            probe,
-                                            cuco::cuda_allocator<std::byte>,
-                                            cuco::experimental::storage<1>>{
-    num_keys, cuco::empty_key<Key>{-1}};
+  auto set = cuco::experimental::static_set{num_keys, cuco::empty_key<Key>{-1}, {}, probe{}};
 
   REQUIRE(set.capacity() == gold_capacity);
 

diff --git a/tests/static_set/unique_sequence_test.cu b/tests/static_set/unique_sequence_test.cu
@@ -142,14 +142,8 @@ TEMPLATE_TEST_CASE_SIG(
     cuco::experimental::linear_probing<CGSize, cuco::default_hash_function<Key>>,
     cuco::experimental::double_hashing<CGSize, cuco::default_hash_function<Key>>>;
 
-  auto set = cuco::experimental::static_set<Key,
-                                            cuco::experimental::extent<size_type>,
-                                            cuda::thread_scope_device,
-                                            thrust::equal_to<Key>,
-                                            probe,
-                                            cuco::cuda_allocator<std::byte>,
-                                            cuco::experimental::storage<2>>{
-    num_keys, cuco::empty_key<Key>{-1}};
+  auto set = cuco::experimental::static_set{
+    num_keys, cuco::empty_key<Key>{-1}, {}, probe{}, {}, cuco::experimental::storage<2>{}};
 
   REQUIRE(set.capacity() == gold_capacity);