Refactor bucket storage to 1D layout for improved performance (#694)

This PR refactors the bucket storage implementation, switching from the
previous 2D layout to a 1D layout. Benchmark results show this change
reduces branching instructions by 30%, leading to roughly a 10%
improvement in runtime performance.

---------

Co-authored-by: Daniel Jünger <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

README.md

@@ -50,7 +50,7 @@ std::enable_if_t<(sizeof(Key) == sizeof(Value)), void> static_multimap_count(
 
   thrust::device_vector<pair_type> pairs(num_keys);
   thrust::transform(keys.begin(), keys.end(), pairs.begin(), [] __device__(Key const& key) {
-    return pair_type(key, {});
+    return pair_type{key, {}};
   });
 
   gen.dropout(keys.begin(), keys.end(), matching_rate);

benchmarks/static_multimap/count_bench.cu

@@ -79,4 +79,4 @@ int main(void)
   }
 
   return 0;
-}
+}

examples/static_multiset/host_bulk_example.cu

@@ -15,7 +15,6 @@
  */
 
 #include <cuco/static_set_ref.cuh>
-#include <cuco/storage.cuh>
 
 #include <cuda/std/array>
 #include <cuda/std/functional>
@@ -143,24 +142,24 @@ int main()
   valid_sizes.reserve(num);
 
   for (size_t i = 0; i < num; ++i) {
-    valid_sizes.emplace_back(
-      static_cast<std::size_t>(cuco::make_bucket_extent<ref_type>(subset_sizes[i])));
+    valid_sizes.emplace_back(static_cast<std::size_t>(
+      cuco::make_valid_extent<probing_scheme_type, storage_ref_type>(subset_sizes[i])));
   }
 
   std::vector<std::size_t> offsets(num + 1, 0);
 
-  // prefix sum to compute offsets and total number of buckets
+  // prefix sum to compute offsets and total number of slots
   std::size_t current_sum = 0;
   for (std::size_t i = 0; i < valid_sizes.size(); ++i) {
     current_sum += valid_sizes[i];
     offsets[i + 1] = current_sum;
   }
 
-  // total number of buckets is located at the back of the offsets array
-  auto const total_num_buckets = offsets.back();
+  // total number of slots is located at the back of the offsets array
+  auto const total_num_slots = offsets.back();
 
   // Create a single bulk storage used by all subsets
-  auto set_storage = storage_type{total_num_buckets};
+  auto set_storage = storage_type{total_num_slots};
   // Initializes the storage with the given sentinel
   set_storage.initialize(empty_key_sentinel);
 

examples/static_set/device_subsets_example.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,16 +26,16 @@
  */
 
 template <class SetRef>
-__global__ void shmem_set_kernel(typename SetRef::extent_type bucket_extent,
+__global__ void shmem_set_kernel(typename SetRef::extent_type valid_extent,
                                  cuco::empty_key<typename SetRef::key_type> empty_key_sentinel)
 {
   // We first allocate the shared memory storage for the `set`.
   // The storage is comprised of contiguous buckets of slots,
   // which allow for vectorized loads.
-  __shared__ typename SetRef::bucket_type buckets[bucket_extent.value()];
+  __shared__ typename SetRef::value_type slots[valid_extent.value()];
 
   // Next, we construct the actual storage object from the raw array.
-  auto storage = SetRef::storage_ref_type(bucket_extent, buckets);
+  auto storage = SetRef::storage_ref_type(valid_extent, slots);
   // Now we can instantiate the set from the storage.
   auto set = SetRef(empty_key_sentinel, {}, {}, {}, storage);
 
@@ -103,9 +103,10 @@ int main(void)
   // Cuco imposes a number of non-trivial contraints on the capacity value.
   // This function will take the requested capacity (1000) and return the next larger
   // valid extent.
-  auto constexpr bucket_extent = cuco::make_bucket_extent<set_ref_type>(extent_type{});
+  auto constexpr valid_extent =
+    cuco::make_valid_extent<probing_scheme_type, cuco::storage<bucket_size>>(extent_type{});
 
   // Launch the kernel with a single thread block.
-  shmem_set_kernel<set_ref_type><<<1, 128>>>(bucket_extent, empty_key_sentinel);
+  shmem_set_kernel<set_ref_type><<<1, 128>>>(valid_extent, empty_key_sentinel);
   cudaDeviceSynchronize();
 }

examples/static_set/shared_memory_example.cu

@@ -16,11 +16,12 @@
 
 #pragma once
 
-#include <cuco/detail/storage/bucket_storage_base.cuh>
+#include <cuco/detail/storage/storage_base.cuh>
 #include <cuco/extent.cuh>
 #include <cuco/utility/allocator.hpp>
 
 #include <cuda/std/array>
+#include <cuda/std/functional>
 #include <cuda/stream_ref>
 
 #include <cstddef>
@@ -29,49 +30,36 @@
 #include <memory>
 
 namespace cuco {
-/// Bucket type alias
-template <typename T, int32_t BucketSize>
-using bucket = detail::bucket<T, BucketSize>;
-
 /**
- * @brief Non-owning array of buckets storage reference type.
+ * @brief Non-owning array of slots storage reference type.
  *
  * @tparam T Storage element type
  * @tparam BucketSize Number of slots in each bucket
  * @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 : public detail::bucket_storage_base<T, BucketSize, Extent> {
+class bucket_storage_ref {
  public:
-  /// Array of buckets base class type
-  using base_type = detail::bucket_storage_base<T, BucketSize, Extent>;
-
-  using base_type::bucket_size;  ///< Number of elements processed per bucket
-
-  using extent_type = typename base_type::extent_type;  ///< Storage extent type
-  using size_type   = typename base_type::size_type;    ///< Storage size type
-  using value_type  = typename base_type::value_type;   ///< Slot type
-  using bucket_type = typename base_type::bucket_type;  ///< Slot bucket type
+  static constexpr int32_t bucket_size = BucketSize;  ///< Number of elements processed per bucket
+  static constexpr std::size_t alignment =
+    cuda::std::min(sizeof(T) * bucket_size, std::size_t{16});  ///< Required alignment
 
-  using base_type::capacity;
-  using base_type::num_buckets;
+  using extent_type = Extent;                            ///< Storage extent type
+  using size_type   = typename extent_type::value_type;  ///< Storage size type
+  using value_type  = T;                                 ///< Slot type
+  using bucket_type = cuda::std::array<T, BucketSize>;   ///< Slot bucket type
 
   /**
-   * @brief Constructor of AoS storage ref.
+   * @brief Constructor of slot storage ref.
    *
-   * @param size Number of buckets
-   * @param buckets Pointer to the buckets array
+   * @param size Number of slots
+   * @param slots Pointer to the slots array
    */
   __host__ __device__ explicit constexpr bucket_storage_ref(Extent size,
-                                                            bucket_type* buckets) noexcept;
+                                                            value_type* slots) noexcept;
 
-  /**
-   * @brief Custom un-incrementable input iterator for the convenience of `find` operations.
-   *
-   * @note This iterator is for read only and NOT incrementable.
-   */
-  struct iterator;
-  using const_iterator = iterator const;  ///< Const forward iterator type
+  using iterator       = value_type*;        ///< Iterator type
+  using const_iterator = value_type const*;  ///< Const forward iterator type
 
   /**
    * @brief Returns an iterator to one past the last slot.
@@ -91,76 +79,90 @@ class bucket_storage_ref : public detail::bucket_storage_base<T, BucketSize, Ext
    *
    * @return A const_iterator to one past the last slot
    */
-  [[nodiscard]] __device__ constexpr const_iterator end() const noexcept;
+  [[nodiscard]] __device__ constexpr iterator end() const noexcept;
 
   /**
-   * @brief Gets buckets array.
+   * @brief Gets slots array.
    *
-   * @return Pointer to the first bucket
+   * @return Pointer to the first slot
    */
-  [[nodiscard]] __device__ constexpr bucket_type* data() noexcept;
+  [[nodiscard]] __device__ constexpr value_type* data() noexcept;
 
   /**
-   * @brief Gets bucket array.
+   * @brief Gets slots array.
    *
-   * @return Pointer to the first bucket
+   * @return Pointer to the first slot
    */
-  [[nodiscard]] __device__ constexpr bucket_type* data() const noexcept;
+  [[nodiscard]] __device__ constexpr value_type* data() const noexcept;
 
   /**
    * @brief Returns an array of slots (or a bucket) for a given index.
    *
-   * @param index Index of the bucket
+   * @param index Index of the slot
    * @return An array of slots
    */
   [[nodiscard]] __device__ constexpr bucket_type operator[](size_type index) const noexcept;
 
+  /**
+   * @brief Gets the total number of slot buckets in the current storage.
+   *
+   * @return The total number of slot buckets
+   */
+  [[nodiscard]] __host__ __device__ constexpr size_type num_buckets() const noexcept;
+
+  /**
+   * @brief Gets the total number of slots in the current storage.
+   *
+   * @return The total number of slots
+   */
+  [[nodiscard]] __host__ __device__ constexpr size_type capacity() const noexcept;
+
+  /**
+   * @brief Gets the bucket extent of the current storage.
+   *
+   * @return The bucket extent.
+   */
+  [[nodiscard]] __host__ __device__ constexpr extent_type extent() const noexcept;
+
  private:
-  bucket_type* buckets_;  ///< Pointer to the buckets array
+  extent_type extent_;  ///< Storage extent
+  value_type* slots_;   ///< Pointer to the slots array
 };
 
 /**
- * @brief Array of buckets open addressing storage class.
+ * @brief Array of slots open addressing storage class.
  *
  * @tparam T Slot type
  * @tparam BucketSize Number of slots in each bucket
- * @tparam Extent Type of extent denoting number of buckets
+ * @tparam Extent Type of extent denoting number of slots
  * @tparam Allocator Type of allocator used for device storage (de)allocation
  */
 template <typename T,
           int32_t BucketSize,
           typename Extent    = cuco::extent<std::size_t>,
-          typename Allocator = cuco::cuda_allocator<cuco::bucket<T, BucketSize>>>
-class bucket_storage : public detail::bucket_storage_base<T, BucketSize, Extent> {
+          typename Allocator = cuco::cuda_allocator<T>>
+class bucket_storage {
  public:
-  /// Array of buckets base class type
-  using base_type = detail::bucket_storage_base<T, BucketSize, Extent>;
-
-  using base_type::bucket_size;  ///< Number of elements processed per bucket
+  static constexpr int32_t bucket_size = BucketSize;  ///< Number of elements processed per bucket
 
-  using extent_type = typename base_type::extent_type;  ///< Storage extent type
-  using size_type   = typename base_type::size_type;    ///< Storage size type
-  using value_type  = typename base_type::value_type;   ///< Slot type
-  using bucket_type = typename base_type::bucket_type;  ///< Slot bucket type
-
-  using base_type::capacity;
-  using base_type::num_buckets;
+  using extent_type = Extent;                            ///< Storage extent type
+  using size_type   = typename extent_type::value_type;  ///< Storage size type
+  using value_type  = T;                                 ///< Slot type
+  using bucket_type = cuda::std::array<T, BucketSize>;   ///< Slot bucket type
 
   /// Type of the allocator to (de)allocate buckets
   using allocator_type =
-    typename std::allocator_traits<Allocator>::template rebind_alloc<bucket_type>;
-  using bucket_deleter_type =
-    detail::custom_deleter<size_type, allocator_type>;  ///< Type of bucket deleter
+    typename std::allocator_traits<Allocator>::template rebind_alloc<value_type>;
   using ref_type = bucket_storage_ref<value_type, bucket_size, extent_type>;  ///< Storage ref type
 
   /**
-   * @brief Constructor of bucket storage.
+   * @brief Constructor of bucket slot storage.
    *
    * @note The input `size` should be exclusively determined by the return value of
-   * `make_bucket_extent` since it depends on the requested low-bound value, the probing scheme, and
+   * `make_valid_extent` since it depends on the requested low-bound value, the probing scheme, and
    * the storage.
    *
-   * @param size Number of buckets to (de)allocate
+   * @param size Number of slots to (de)allocate
    * @param allocator Allocator used for (de)allocating device storage
    */
   explicit constexpr bucket_storage(Extent size, Allocator const& allocator = {});
@@ -178,11 +180,11 @@ class bucket_storage : public detail::bucket_storage_base<T, BucketSize, Extent>
   bucket_storage& operator=(bucket_storage const&) = delete;
 
   /**
-   * @brief Gets buckets array.
+   * @brief Gets bucket slots array.
    *
-   * @return Pointer to the first bucket
+   * @return Pointer to the first slot
    */
-  [[nodiscard]] constexpr bucket_type* data() const noexcept;
+  [[nodiscard]] constexpr value_type* data() const noexcept;
 
   /**
    * @brief Gets the storage allocator.
@@ -199,7 +201,7 @@ class bucket_storage : public detail::bucket_storage_base<T, BucketSize, Extent>
   [[nodiscard]] constexpr ref_type ref() const noexcept;
 
   /**
-   * @brief Initializes each slot in the bucket storage to contain `key`.
+   * @brief Initializes each slot in the bucket slot storage to contain `key`.
    *
    * @param key Key to which all keys in `slots` are initialized
    * @param stream Stream used for executing the kernel
@@ -212,13 +214,37 @@ class bucket_storage : public detail::bucket_storage_base<T, BucketSize, Extent>
    * @param key Key to which all keys in `slots` are initialized
    * @param stream Stream used for executing the kernel
    */
-  void initialize_async(value_type key, cuda::stream_ref stream = {}) noexcept;
+  void initialize_async(value_type key, cuda::stream_ref stream = {});
+
+  /**
+   * @brief Gets the total number of slot buckets in the current storage.
+   *
+   * @return The total number of slot buckets
+   */
+  [[nodiscard]] __host__ __device__ constexpr size_type num_buckets() const noexcept;
+
+  /**
+   * @brief Gets the total number of slots in the current storage.
+   *
+   * @return The total number of slots
+   */
+  [[nodiscard]] __host__ __device__ constexpr size_type capacity() const noexcept;
+
+  /**
+   * @brief Gets the bucket extent of the current storage.
+   *
+   * @return The bucket extent.
+   */
+  [[nodiscard]] __host__ __device__ constexpr extent_type extent() const noexcept;
 
  private:
-  allocator_type allocator_;            ///< Allocator used to (de)allocate buckets
-  bucket_deleter_type bucket_deleter_;  ///< Custom buckets deleter
-  /// Pointer to the bucket storage
-  std::unique_ptr<bucket_type, bucket_deleter_type> buckets_;
+  using slot_deleter_type =
+    detail::custom_deleter<size_type, allocator_type>;  ///< Type of slot deleter
+
+  extent_type extent_;        ///< Storage extent
+  allocator_type allocator_;  ///< Allocator used to (de)allocate slots
+  /// Pointer to the slot storage
+  std::unique_ptr<value_type, slot_deleter_type> slots_;
 };
 }  // namespace cuco