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>
 
 #include <cstddef>
@@ -24,20 +25,24 @@ namespace cuco {
  * @brief Non-owning array of slots storage reference type.
  *
  * @tparam T Storage element type
- * @tparam BucketSize Number of slots in each bucket
+ * @tparam BucketSize Number of slots in each bucket; must be a positive power of two
  * @tparam Extent Type of extent denoting storage capacity
  */
 template <typename T, int32_t BucketSize, typename Extent = cuco::extent<std::size_t>>
 class bucket_storage_ref {
+  static_assert(BucketSize > 0 && (BucketSize & (BucketSize - 1)) == 0,
+                "Bucket size must be a positive power of two");
+
  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 +51,8 @@ class bucket_storage_ref {
   /**
    * @brief Constructor of slot storage ref.
    *
+   * @note `slots` must be aligned to `alignment` bytes.
+   *
    * @param size Number of slots
    * @param slots Pointer to the slots array
    */
@@ -92,11 +99,25 @@ class bucket_storage_ref {
   /**
    * @brief Returns an array of slots (or a bucket) for a given index.
    *
+   * @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`.
+   *
+   * @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

@@ -379,7 +379,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 = storage_ref_.load_bucket(*probing_iter);
 
       for (auto& slot_content : bucket_slots) {
         auto const eq_res = this->predicate_.template operator()<is_insert::YES>(
@@ -438,9 +438,36 @@ 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 = storage_ref_.load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
+        if constexpr (bucket_size > 1) {
+          // Resolve availability before the ordered key comparisons.
+          int32_t first_available = bucket_size;
+          cuda::static_for<bucket_size>([&] __device__(auto i) {
+            auto const slot_key = this->extract_key(bucket_slots[i()]);
+            if (cuco::detail::bitwise_compare(slot_key, this->predicate_.empty_sentinel_) ||
+                cuco::detail::bitwise_compare(slot_key, this->predicate_.erased_sentinel_)) {
+              first_available = cuda::std::min(first_available, int32_t{i()});
+            }
+          });
+          bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
+          if constexpr (!allows_duplicates) {
+            cuda::static_for<bucket_size>([&] __device__(auto i) {
+              if (i() < first_available && result.state_ == detail::equal_result::UNEQUAL) {
+                auto const res =
+                  this->predicate_.equal_to(key, this->extract_key(bucket_slots[i()]));
+                if (res != detail::equal_result::UNEQUAL) {
+                  result = bucket_probing_results{res, i()};
+                }
+              }
+            });
+          }
+          if (result.state_ == detail::equal_result::UNEQUAL && first_available < bucket_size) {
+            return bucket_probing_results{detail::equal_result::AVAILABLE, first_available};
+          }
+          return result;
+        }
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
         cuda::static_for<bucket_size>([&] __device__(auto i) {
           if (result.state_ == detail::equal_result::UNEQUAL) {
@@ -473,9 +500,13 @@ class open_addressing_ref_impl
         auto status         = insert_result::CONTINUE;
         if (group.thread_rank() == src_lane) {
           if constexpr (SupportsErase) {
-            status = attempt_insert(this->get_slot_ptr(*probing_iter, intra_bucket_index),
-                                    bucket_slots[intra_bucket_index],
-                                    val);
+            // Use constant indices to keep the bucket snapshot out of local memory.
+            auto slot_content = bucket_slots[0];
+            cuda::static_for<bucket_size>([&, slot_index = intra_bucket_index] __device__(auto i) {
+              if (slot_index == i()) { slot_content = bucket_slots[i()]; }
+            });
+            status = attempt_insert(
+              this->get_slot_ptr(*probing_iter, intra_bucket_index), slot_content, val);
           } else {
             status = attempt_insert(this->get_slot_ptr(*probing_iter, intra_bucket_index),
                                     this->empty_slot_sentinel(),
@@ -540,7 +571,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 = storage_ref_.load_bucket(*probing_iter);
 
       for (auto i = 0; i < bucket_size; ++i) {
         auto const eq_res = this->predicate_.template operator()<is_insert::YES>(
@@ -612,7 +643,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 = storage_ref_.load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -644,7 +675,11 @@ class open_addressing_ref_impl
         auto const res      = group.shfl(reinterpret_cast<intptr_t>(slot_ptr), src_lane);
         auto const status   = [&, target_idx = intra_bucket_index]() {
           if (group.thread_rank() != src_lane) { return insert_result::CONTINUE; }
-          return this->attempt_insert_stable(slot_ptr, bucket_slots[target_idx], val);
+          auto slot_content = bucket_slots[0];
+          cuda::static_for<bucket_size>([&, slot_index = target_idx] __device__(auto i) {
+            if (slot_index == i()) { slot_content = bucket_slots[i()]; }
+          });
+          return this->attempt_insert_stable(slot_ptr, slot_content, val);
         }();
 
         switch (group.shfl(status, src_lane)) {
@@ -692,7 +727,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 = storage_ref_.load_bucket(*probing_iter);
 
       for (auto& slot_content : bucket_slots) {
         auto const eq_res =
@@ -740,7 +775,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 = storage_ref_.load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
@@ -804,16 +839,24 @@ 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 = storage_ref_.load_bucket(*probing_iter);
 
-      for (auto i = 0; i < bucket_size; ++i) {
-        switch (this->predicate_.template operator()<is_insert::NO>(
-          key, this->extract_key(bucket_slots[i]))) {
-          case detail::equal_result::UNEQUAL: continue;
-          case detail::equal_result::EMPTY: return false;
-          case detail::equal_result::EQUAL: return true;
+      int32_t first_empty = bucket_size;
+      cuda::static_for<bucket_size>([&] __device__(auto i) {
+        if (cuco::detail::bitwise_compare(this->extract_key(bucket_slots[i()]),
+                                          this->predicate_.empty_sentinel_)) {
+          first_empty = cuda::std::min(first_empty, int32_t{i()});
         }
-      }
+      });
+      bool found = false;
+      cuda::static_for<bucket_size>([&] __device__(auto i) {
+        if (i() < first_empty && !found) {
+          found = this->predicate_.equal_to(key, this->extract_key(bucket_slots[i()])) ==
+                  detail::equal_result::EQUAL;
+        }
+      });
+      if (found) { return true; }
+      if (first_empty < bucket_size) { return false; }
       ++probing_iter;
       if (*probing_iter == init_idx) { return false; }
     }
@@ -845,20 +888,24 @@ class open_addressing_ref_impl
     auto const init_idx = *probing_iter;
 
     while (true) {
-      auto const bucket_slots = storage_ref_[*probing_iter];
-
-      auto const state = [&]() {
-        auto res = detail::equal_result::UNEQUAL;
-        for (auto i = 0; i < bucket_size; ++i) {
-          res = this->predicate_.template operator()<is_insert::NO>(
-            key, this->extract_key(bucket_slots[i]));
-          if (res != detail::equal_result::UNEQUAL) { return res; }
-        }
-        return res;
-      }();
+      auto const bucket_slots = storage_ref_.load_bucket(*probing_iter);
 
-      if (group.any(state == detail::equal_result::EQUAL)) { return true; }
-      if (group.any(state == detail::equal_result::EMPTY)) { return false; }
+      int32_t first_empty = bucket_size;
+      cuda::static_for<bucket_size>([&] __device__(auto i) {
+        if (cuco::detail::bitwise_compare(this->extract_key(bucket_slots[i()]),
+                                          this->predicate_.empty_sentinel_)) {
+          first_empty = cuda::std::min(first_empty, int32_t{i()});
+        }
+      });
+      bool found = false;
+      cuda::static_for<bucket_size>([&] __device__(auto i) {
+        if (i() < first_empty && !found) {
+          found = this->predicate_.equal_to(key, this->extract_key(bucket_slots[i()])) ==
+                  detail::equal_result::EQUAL;
+        }
+      });
+      if (group.any(found)) { return true; }
+      if (group.any(first_empty < bucket_size)) { return false; }
 
       ++probing_iter;
       if (*probing_iter == init_idx) { return false; }
@@ -890,7 +937,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 = storage_ref_.load_bucket(*probing_iter);
 
       for (auto i = 0; i < bucket_size; ++i) {
         switch (this->predicate_.template operator()<is_insert::NO>(
@@ -935,17 +982,27 @@ 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 = storage_ref_.load_bucket(*probing_iter);
 
       auto const [state, intra_bucket_index] = [&]() {
+        // Classify empty slots first so independent bucket loads can overlap.
+        int32_t first_empty = bucket_size;
+        cuda::static_for<bucket_size>([&] __device__(auto i) {
+          if (cuco::detail::bitwise_compare(this->extract_key(bucket_slots[i()]),
+                                            this->predicate_.empty_sentinel_)) {
+            first_empty = cuda::std::min(first_empty, int32_t{i()});
+          }
+        });
         bucket_probing_results result{detail::equal_result::UNEQUAL, -1};
         cuda::static_for<bucket_size>([&] __device__(auto i) {
-          if (result.state_ == detail::equal_result::UNEQUAL) {
-            auto res = this->predicate_.template operator()<is_insert::NO>(
-              key, this->extract_key(bucket_slots[i()]));
+          if (i() < first_empty && result.state_ == detail::equal_result::UNEQUAL) {
+            auto const res = this->predicate_.equal_to(key, this->extract_key(bucket_slots[i()]));
             if (res != detail::equal_result::UNEQUAL) { result = bucket_probing_results{res, i()}; }
           }
         });
+        if (result.state_ == detail::equal_result::UNEQUAL && first_empty < bucket_size) {
+          return bucket_probing_results{detail::equal_result::EMPTY, first_empty};
+        }
         return result;
       }();
 
@@ -991,7 +1048,7 @@ class open_addressing_ref_impl
       size_type count     = 0;
 
       while (true) {
-        auto const bucket_slots                = storage_ref_[*probing_iter];
+        auto const bucket_slots                = storage_ref_.load_bucket(*probing_iter);
         cuda::std::int32_t equals[bucket_size] = {0};
         bool empty_found                       = false;
 
@@ -1036,7 +1093,7 @@ class open_addressing_ref_impl
     size_type count     = 0;
 
     while (true) {
-      auto const bucket_slots                = storage_ref_[*probing_iter];
+      auto const bucket_slots                = storage_ref_.load_bucket(*probing_iter);
       cuda::std::int32_t equals[bucket_size] = {0};
       bool should_return                     = false;
 
@@ -1286,7 +1343,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 = storage_ref_.load_bucket(*probing_iter);
 
             cuda::static_for<bucket_size>([&] __device__(auto i) {
               equals[i()] = false;
@@ -1397,7 +1454,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 = storage_ref_.load_bucket(*probing_iter);
 
       bool should_return = false;
       cuda::static_for<bucket_size>([&] __device__(auto i) {
@@ -1458,7 +1515,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 = storage_ref_.load_bucket(*probing_iter);
 
       for (cuda::std::int32_t i = 0; i < bucket_size and !should_return; ++i) {
         switch (this->predicate_.template operator()<is_insert::NO>(
@@ -1528,7 +1585,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 = storage_ref_.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>(