[cudax][cuco] Migrate host rehash APIs for fixed_capacity_set

cudax/include/cuda/experimental/__cuco/fixed_capacity_set.cuh

@@ -200,6 +200,64 @@ public:
     __impl_->clear_async(__stream);
   }
 
+  //! @brief Reinserts every stored key into new storage of the same capacity.
+  //!
+  //! @note Synchronizes the supplied stream. For asynchronous execution use `rehash_async`.
+  //! @note Requires exclusive access. Refs and storage pointers obtained earlier are invalidated.
+  //! @throws cuda_error if allocation, the operation, or stream synchronization fails
+  //! @param[in] __stream Stream used for rehashing
+  _CCCL_HOST_API void rehash(::cuda::stream_ref __stream)
+  {
+    rehash_async(__stream);
+    __stream.sync();
+  }
+
+  //! @brief Reinserts every stored key into new storage of at least `__capacity` slots.
+  //!
+  //! The new capacity is the smallest valid slot count not less than `__capacity`.
+  //! @pre `__capacity` is large enough to hold every stored key.
+  //! @note Synchronizes the supplied stream. For asynchronous execution use `rehash_async`.
+  //! @note Requires exclusive access. Refs and storage pointers obtained earlier are invalidated.
+  //! @note Only available for sets with dynamic capacity.
+  //! @throws cuda_error if allocation, the operation, or stream synchronization fails
+  //! @throws std::logic_error if a valid capacity cannot be represented
+  //! @param[in] __stream Stream used for rehashing
+  //! @param[in] __capacity Requested minimum number of slots
+  _CCCL_TEMPLATE(::cuda::std::size_t _C = _Capacity)
+  _CCCL_REQUIRES((_C == _Capacity) _CCCL_AND(_C == ::cuda::std::dynamic_extent))
+  _CCCL_HOST_API void rehash(::cuda::stream_ref __stream, size_type __capacity)
+  {
+    rehash_async(__stream, __capacity);
+    __stream.sync();
+  }
+
+  //! @brief Asynchronously reinserts every stored key into new storage of the same capacity.
+  //!
+  //! @note Requires exclusive access. Refs and storage pointers obtained earlier are invalidated.
+  //! @throws cuda_error if allocation or the operation fails to launch
+  //! @param[in] __stream Stream used for rehashing
+  _CCCL_HOST_API void rehash_async(::cuda::stream_ref __stream)
+  {
+    __impl_->rehash_async(__stream, *this);
+  }
+
+  //! @brief Asynchronously reinserts every stored key into new storage of at least `__capacity` slots.
+  //!
+  //! The new capacity is the smallest valid slot count not less than `__capacity`.
+  //! @pre `__capacity` is large enough to hold every stored key.
+  //! @note Requires exclusive access. Refs and storage pointers obtained earlier are invalidated.
+  //! @note Only available for sets with dynamic capacity.
+  //! @throws cuda_error if allocation or the operation fails to launch
+  //! @throws std::logic_error if a valid capacity cannot be represented
+  //! @param[in] __stream Stream used for rehashing
+  //! @param[in] __capacity Requested minimum number of slots
+  _CCCL_TEMPLATE(::cuda::std::size_t _C = _Capacity)
+  _CCCL_REQUIRES((_C == _Capacity) _CCCL_AND(_C == ::cuda::std::dynamic_extent))
+  _CCCL_HOST_API void rehash_async(::cuda::stream_ref __stream, size_type __capacity)
+  {
+    __impl_->rehash_async(__stream, __capacity, *this);
+  }
+
   //! @brief Inserts keys and returns the number of new keys inserted.
   //!
   //! Duplicate keys and keys that cannot fit in the set do not contribute to the count.

cudax/test/cuco/fixed_capacity_set/test_capacity.cu

@@ -15,6 +15,7 @@
 #include <cuda/std/limits>
 #include <cuda/std/span>
 #include <cuda/std/type_traits>
+#include <cuda/std/utility>
 #include <cuda/stream>
 #include <cuda/type_traits>
 
@@ -28,6 +29,17 @@
 
 #include "test_common.cuh"
 
+template <class Set, class = void>
+inline constexpr bool has_capacity_rehash = false;
+
+template <class Set>
+inline constexpr bool has_capacity_rehash<
+  Set,
+  cuda::std::void_t<
+    decltype(cuda::std::declval<Set&>().rehash(cuda::std::declval<cuda::stream_ref>(), cuda::std::size_t{})),
+    decltype(cuda::std::declval<Set&>().rehash_async(cuda::std::declval<cuda::stream_ref>(), cuda::std::size_t{}))>> =
+  true;
+
 using probing_types =
   c2h::type_list<cudax::cuco::linear_probing<4, cuda::hash<int>>, cudax::cuco::double_hashing<2, cuda::hash<int>>>;
 
@@ -74,6 +86,8 @@ C2H_TEST("fixed_capacity_set capacity constructors", "[container][capacity]", pr
   static_assert(!cuda::std::is_copy_assignable_v<dynamic_set>);
   static_assert(cuda::std::is_nothrow_move_constructible_v<dynamic_set>);
   static_assert(cuda::std::is_nothrow_move_assignable_v<dynamic_set>);
+  static_assert(has_capacity_rehash<dynamic_set>);
+  static_assert(!has_capacity_rehash<static_set>);
 
   test_context context;
   const dynamic_set dynamic{context.stream, context.mr, requested, empty_key{-1}};
@@ -91,6 +105,10 @@ C2H_TEST("fixed_capacity_set capacity constructors", "[container][capacity]", pr
   REQUIRE(fixed.insert(context.stream, first, first + num_keys) == num_keys);
   fixed.contains(context.stream, first, first + num_keys, results.begin());
   REQUIRE(context.all_equal(results.data(), num_keys, 1));
+  fixed.rehash(context.stream);
+  REQUIRE(fixed.capacity() == valid);
+  fixed.contains(context.stream, first, first + num_keys, results.begin());
+  REQUIRE(context.all_equal(results.data(), num_keys, 1));
 
   const dynamic_set zero{context.stream, context.mr, cuda::std::size_t{0}, empty_key{-1}};
   REQUIRE(zero.capacity() == cudax::cuco::make_valid_capacity<probing_type, 2>(cuda::std::size_t{0}));

cudax/test/cuco/fixed_capacity_set/test_insert_and_contains.cu

@@ -159,6 +159,20 @@ C2H_TEST("fixed_capacity_set insert and contains", "[container]", key_types, cg_
   set.contains(context.stream, first, first + num_queries, found.begin());
   REQUIRE(context.matches(found.data(), num_queries, num_keys));
 
+  const auto initial_capacity = set.capacity();
+  set.rehash(context.stream);
+  set.rehash_async(context.stream);
+  REQUIRE(set.capacity() == initial_capacity);
+  set.contains(context.stream, first, first + num_queries, found.begin());
+  REQUIRE(context.matches(found.data(), num_queries, num_keys));
+  set.rehash(context.stream, 2 * initial_capacity);
+  REQUIRE(set.capacity() >= 2 * initial_capacity);
+  set.rehash_async(context.stream, initial_capacity);
+  REQUIRE(set.capacity() == initial_capacity);
+  set.contains(context.stream, first, first + num_queries, found.begin());
+  REQUIRE(context.matches(found.data(), num_queries, num_keys));
+  REQUIRE(set.insert(context.stream, first, first + num_keys) == 0);
+
   // The stencil is negative only for the first half while every key is nonnegative, so the stencil, not the key,
   // selects the first half.
   const auto stencil = cuda::counting_iterator<int>{-num_keys / 2};