Restore aligned bucket loads

include/cuco/bucket_storage.cuh

@@ -9,9 +9,10 @@
 #include <cuco/extent.cuh>
 #include <cuco/utility/allocator.hpp>
 
+#include <cuda/std/algorithm>
 #include <cuda/std/array>
-#include <cuda/std/bit>
 #include <cuda/std/functional>
+#include <cuda/std/numeric>
 #include <cuda/stream_ref>
 
 #include <cstddef>
@@ -29,15 +30,18 @@ namespace cuco {
  */
 template <typename T, int32_t BucketSize, typename Extent = cuco::extent<std::size_t>>
 class bucket_storage_ref {
+  static_assert(BucketSize > 0, "Bucket size must be positive");
+
  public:
   static constexpr int32_t bucket_size = BucketSize;        ///< Number of elements per bucket
-  static constexpr std::size_t max_vector_load_bytes = 16;  ///< Maximum vector load width in bytes
+  static constexpr std::size_t max_vector_load_bytes = 32;  ///< Maximum vector load width in bytes
 
   using bucket_type = cuda::std::array<T, BucketSize>;  ///< Slot bucket type
 
   static constexpr std::size_t alignment =
-    cuda::std::min(cuda::std::bit_ceil(sizeof(bucket_type)),
-                   max_vector_load_bytes);  ///< Required alignment in bytes
+    cuda::std::max(alignof(T),
+                   cuda::std::gcd(sizeof(T) * BucketSize,
+                                  max_vector_load_bytes));  ///< Required alignment in bytes
 
   using extent_type = Extent;                            ///< Storage extent type
   using size_type   = typename extent_type::value_type;  ///< Storage size type
@@ -46,6 +50,9 @@ class bucket_storage_ref {
   /**
    * @brief Constructor of slot storage ref.
    *
+   * @note `slots` must be aligned to `alignment` bytes. This alignment is
+   * preserved at every bucket boundary, including for non-power-of-two buckets.
+   *
    * @param size Number of slots
    * @param slots Pointer to the slots array
    */
@@ -92,11 +99,28 @@ class bucket_storage_ref {
   /**
    * @brief Returns an array of slots (or a bucket) for a given index.
    *
+   * @pre The complete range `[index, index + bucket_size)` is within the storage.
+   *
+   * @note `index` need not be a multiple of `bucket_size`.
+   *
    * @param index Index of the slot
    * @return An array of slots
    */
   [[nodiscard]] __device__ constexpr bucket_type operator[](size_type index) const noexcept;
 
+  /**
+   * @brief Loads a bucket starting at a bucket-aligned slot index.
+   *
+   * Unlike `operator[]`, this access exposes the guaranteed bucket alignment to the compiler.
+   *
+   * @pre `index` is a multiple of `bucket_size`.
+   * @pre The complete range `[index, index + bucket_size)` is within the storage.
+   *
+   * @param index Index of the first slot in the bucket
+   * @return An array containing the bucket's slots
+   */
+  [[nodiscard]] __device__ constexpr bucket_type load_bucket(size_type index) const noexcept;
+
   /**
    * @brief Gets the total number of slot buckets in the current storage.
    *

include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh

@@ -5,6 +5,7 @@
 
 #pragma once
 
+#include <cuco/bucket_storage.cuh>
 #include <cuco/detail/equal_wrapper.cuh>
 #include <cuco/detail/open_addressing/constraints.cuh>
 #include <cuco/detail/probing_scheme/probing_scheme_base.cuh>
@@ -118,6 +119,24 @@ class open_addressing_ref_impl
     storage_ref_type::bucket_size;             ///< Number of elements handled per bucket
   static constexpr auto thread_scope = Scope;  ///< CUDA thread scope
 
+  /**
+   * @brief Loads the complete bucket at the probing iterator's slot index.
+   *
+   * Probing schemes must produce bucket-aligned slot indices.
+   *
+   * @param index Slot index produced by the probing iterator
+   * @return The bucket at `index`
+   */
+  [[nodiscard]] __device__ bucket_type load_bucket(size_type index) const noexcept
+  {
+    using native_storage_ref = bucket_storage_ref<value_type, bucket_size, extent_type>;
+    if constexpr (cuda::std::is_same_v<storage_ref_type, native_storage_ref>) {
+      return storage_ref_.load_bucket(index);
+    } else {
+      return storage_ref_[index];
+    }
+  }
+
   /**
    * @brief Constructs open_addressing_ref_impl.
    *
@@ -375,7 +394,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       for (auto& slot_content : bucket_slots) {
         auto const eq_res = this->predicate_.template operator()<is_insert::YES>(
@@ -428,7 +447,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -524,7 +543,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       for (auto i = 0; i < bucket_size; ++i) {
         auto const eq_res = this->predicate_.template operator()<is_insert::YES>(
@@ -590,7 +609,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -664,7 +683,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       for (auto& slot_content : bucket_slots) {
         auto const eq_res =
@@ -709,7 +728,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -770,7 +789,7 @@ class open_addressing_ref_impl
 
     while (true) {
       // TODO atomic_ref::load if insert operator is present
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       for (auto i = 0; i < bucket_size; ++i) {
         switch (this->predicate_.template operator()<is_insert::NO>(
@@ -808,7 +827,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       auto const state = [&]() {
         auto res = detail::equal_result::UNEQUAL;
@@ -850,7 +869,7 @@ class open_addressing_ref_impl
 
     while (true) {
       // TODO atomic_ref::load if insert operator is present
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       for (auto i = 0; i < bucket_size; ++i) {
         switch (this->predicate_.template operator()<is_insert::NO>(
@@ -892,7 +911,7 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -945,7 +964,7 @@ class open_addressing_ref_impl
       size_type count     = 0;
 
       while (true) {
-        auto const bucket_slots                = storage_ref_[*probing_iter];
+        auto const bucket_slots                = this->load_bucket(*probing_iter);
         cuda::std::int32_t equals[bucket_size] = {0};
         bool empty_found                       = false;
 
@@ -987,7 +1006,7 @@ class open_addressing_ref_impl
     size_type count     = 0;
 
     while (true) {
-      auto const bucket_slots                = storage_ref_[*probing_iter];
+      auto const bucket_slots                = this->load_bucket(*probing_iter);
       cuda::std::int32_t equals[bucket_size] = {0};
       bool empty_found                       = false;
 
@@ -1298,7 +1317,7 @@ class open_addressing_ref_impl
         while (active_flushing_tile.any(running)) {
           if (running) {
             // TODO atomic_ref::load if insert operator is present
-            auto const bucket_slots = this->storage_ref_[*probing_iter];
+            auto const bucket_slots = this->load_bucket(*probing_iter);
 
             cuda::static_for<bucket_size>([&] __device__(auto i) {
               equals[i()] = false;
@@ -1419,7 +1438,7 @@ class open_addressing_ref_impl
 
     while (true) {
       // TODO atomic_ref::load if insert operator is present
-      auto const bucket_slots = this->storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       bool should_return = false;
       cuda::static_for<bucket_size>([&] __device__(auto i) {
@@ -1476,7 +1495,7 @@ class open_addressing_ref_impl
 
     while (true) {
       // TODO atomic_ref::load if insert operator is present
-      auto const bucket_slots = this->storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) {
         switch (this->predicate_.template operator()<is_insert::NO>(
@@ -1542,7 +1561,7 @@ class open_addressing_ref_impl
 
     while (true) {
       // TODO atomic_ref::load if insert operator is present
-      auto const bucket_slots = this->storage_ref_[*probing_iter];
+      auto const bucket_slots = this->load_bucket(*probing_iter);
 
       for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) {
         switch (this->predicate_.template operator()<is_insert::NO>(

include/cuco/detail/probing_scheme/probing_scheme_base.cuh

@@ -15,6 +15,9 @@ namespace detail {
  *
  * This class should not be used directly.
  *
+ * @note Derived probing schemes must produce bucket-aligned slot indices. For a bucket size
+ * `B`, every index must be a multiple of `B` and the complete bucket must fit within the storage.
+ *
  * @tparam CGSize Size of CUDA Cooperative Groups
  */
 template <int32_t CGSize>

include/cuco/detail/static_map/static_map_ref.inl

@@ -513,7 +513,7 @@ class operator_impl<
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref[*probing_iter];
+      auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter);
 
       for (auto& slot_content : bucket_slots) {
         auto const eq_res =
@@ -563,7 +563,7 @@ class operator_impl<
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref[*probing_iter];
+      auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         detail::bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -888,7 +888,7 @@ class operator_impl<
     auto constexpr wait_for_payload = (not UseDirectApply) and (sizeof(value_type) > 8);
 
     while (true) {
-      auto const bucket_slots = storage_ref[*probing_iter];
+      auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter);
 
       for (auto& slot_content : bucket_slots) {
         auto const eq_res =
@@ -966,7 +966,7 @@ class operator_impl<
     auto constexpr wait_for_payload = (not UseDirectApply) and (sizeof(value_type) > 8);
 
     while (true) {
-      auto const bucket_slots = storage_ref[*probing_iter];
+      auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         detail::bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -1592,4 +1592,4 @@ class operator_impl<
 };
 
 }  // namespace detail
-}  // namespace cuco
+}  // namespace cuco

include/cuco/detail/storage/bucket_storage.inl

@@ -65,6 +65,15 @@ bucket_storage_ref<T, BucketSize, Extent>::operator[](size_type index) const noe
   return *reinterpret_cast<bucket_type*>(this->data() + index);
 }
 
+template <typename T, int BucketSize, typename Extent>
+__device__ constexpr bucket_storage_ref<T, BucketSize, Extent>::bucket_type
+bucket_storage_ref<T, BucketSize, Extent>::load_bucket(size_type index) const noexcept
+{
+  assert(index % bucket_size == 0);
+  assert(index <= capacity() && bucket_size <= capacity() - index);
+  return *reinterpret_cast<bucket_type*>(__builtin_assume_aligned(this->data() + index, alignment));
+}
+
 template <typename T, int BucketSize, typename Extent>
 __host__ __device__ constexpr typename bucket_storage_ref<T, BucketSize, Extent>::size_type
 bucket_storage_ref<T, BucketSize, Extent>::num_buckets() const noexcept

include/cuco/static_map_ref.cuh

@@ -269,7 +269,7 @@ class static_map_ref
    *
    * @param tile The cooperative thread group used to copy the data structure
    * @param memory_to_use Array large enough to support `capacity` elements. Object does not take
-   * the ownership of the memory
+   * the ownership of the memory. Must satisfy the storage reference's alignment requirements.
    * @param scope The thread scope of the newly created device ref
    *
    * @return Copy of the current device ref

include/cuco/static_multimap_ref.cuh

@@ -271,7 +271,7 @@ class static_multimap_ref
    *
    * @param tile The cooperative thread group used to copy the data structure
    * @param memory_to_use Array large enough to support `capacity` elements. Object does not take
-   * the ownership of the memory
+   * the ownership of the memory. Must satisfy the storage reference's alignment requirements.
    * @param scope The thread scope of the newly created device ref
    *
    * @return Copy of the current device ref

include/cuco/static_multiset_ref.cuh

@@ -255,7 +255,7 @@ class static_multiset_ref
    *
    * @param tile The cooperative thread group used to copy the data structure
    * @param memory_to_use Array large enough to support `capacity` elements. Object does not take
-   * the ownership of the memory
+   * the ownership of the memory. Must satisfy the storage reference's alignment requirements.
    * @param scope The thread scope of the newly created device ref
    *
    * @return Copy of the current device ref

include/cuco/static_set_ref.cuh

@@ -253,7 +253,7 @@ class static_set_ref
    *
    * @param tile The cooperative thread group used to copy the data structure
    * @param memory_to_use Array large enough to support `capacity` elements. Object does not take
-   * the ownership of the memory
+   * the ownership of the memory. Must satisfy the storage reference's alignment requirements.
    * @param scope The thread scope of the newly created device ref
    *
    * @return Copy of the current device ref

tests/CMakeLists.txt

@@ -40,6 +40,7 @@ endfunction(ConfigureTest)
 ###################################################################################################
 # - utility tests ---------------------------------------------------------------------------------
 ConfigureTest(UTILITY_TEST
+    utility/aligned_storage_test.cu
     utility/extent_test.cu
     utility/next_prime_test.cu
     utility/storage_test.cu

tests/static_map/shared_memory_test.cu

@@ -32,7 +32,8 @@ __global__ void shared_memory_test_kernel(Ref* maps,
   size_t const map_id = blockIdx.x;
   size_t const offset = map_id * number_of_elements;
 
-  __shared__ typename Ref::value_type sm_buffer[ValidSize];
+  using storage_ref_type = typename Ref::storage_ref_type;
+  alignas(storage_ref_type::alignment) __shared__ typename Ref::value_type sm_buffer[ValidSize];
 
   auto g          = cuco::test::cg::this_thread_block();
   auto insert_ref = maps[map_id].make_copy(g, sm_buffer, cuco::thread_scope_block);

tests/static_set/shared_memory_test.cu

@@ -31,7 +31,8 @@ __global__ void shared_memory_test_kernel(Ref* sets,
   size_t const set_id = blockIdx.x;
   size_t const offset = set_id * number_of_elements;
 
-  __shared__ typename Ref::value_type sm_buffer[ValidSize];
+  using storage_ref_type = typename Ref::storage_ref_type;
+  alignas(storage_ref_type::alignment) __shared__ typename Ref::value_type sm_buffer[ValidSize];
 
   auto g          = cuco::test::cg::this_thread_block();
   auto insert_ref = sets[set_id].make_copy(g, sm_buffer, cuco::thread_scope_block);