[cudax][cuco] Port fixed_capacity_map benchmarks from cuCollections

cudax/benchmarks/bench/cuco/common/defaults.cuh

@@ -0,0 +1,44 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <cuda/std/cstdint>
+
+#include <vector>
+
+#include <nvbench/nvbench.cuh>
+#include <nvbench/range.cuh>
+
+namespace cuda::experimental::cuco::benchmark::defaults
+{
+//! Key types covered by the default CUCO benchmark type axes.
+using key_type_range = ::nvbench::type_list<::nvbench::int32_t, ::nvbench::int64_t>;
+//! Value types covered by the default CUCO benchmark type axes.
+using value_type_range = ::nvbench::type_list<::nvbench::int32_t, ::nvbench::int64_t>;
+
+//! Default number of inputs used when sweeping another benchmark axis.
+inline constexpr auto n = ::nvbench::int64_t{100'000'000};
+//! Default fixed-capacity map target occupancy.
+inline constexpr auto occupancy = 0.5;
+//! Default lookup matching rate for contains-style benchmarks.
+inline constexpr auto matching_rate = 1.0;
+//! Default deterministic seed used by benchmark data generators.
+inline constexpr auto seed = ::cuda::std::uint32_t{42};
+
+//! Input-size sweep that remains cacheable for direct comparisons with CUCO benchmarks.
+inline const auto n_range_cache = ::std::vector<::nvbench::int64_t>{8'000, 80'000, 800'000, 8'000'000, 80'000'000};
+//! Occupancy sweep used by fixed-capacity container benchmarks.
+inline const auto occupancy_range = ::nvbench::range(0.1, 0.9, 0.1);
+//! Average multiplicity sweep for duplicate-key distributions.
+inline const auto multiplicity_range = ::std::vector<::nvbench::int64_t>{1, 2, 4, 8, 16};
+//! Matching-rate sweep used by contains-style benchmarks.
+inline const auto matching_rate_range = ::nvbench::range(0.1, 1.0, 0.1);
+} // namespace cuda::experimental::cuco::benchmark::defaults

cudax/benchmarks/bench/cuco/common/key_generator.cuh

@@ -0,0 +1,283 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <cuda/std/detail/__config>
+
+#include <thrust/execution_policy.h>
+#include <thrust/random.h>
+#include <thrust/sequence.h>
+#include <thrust/shuffle.h>
+#include <thrust/transform.h>
+
+#include <cuda/iterator>
+#include <cuda/std/cmath>
+#include <cuda/std/cstddef>
+#include <cuda/std/iterator>
+#include <cuda/std/limits>
+#include <cuda/std/type_traits>
+
+#include <stdexcept>
+
+#include "defaults.cuh"
+#include <nvbench/nvbench.cuh>
+
+namespace cuda::experimental::cuco::benchmark
+{
+namespace distribution
+{
+//! Distribution tag for unique keys generated by shuffling the sequence `[0, N)`.
+struct unique
+{};
+
+//! Distribution tag for uniformly sampled keys with controlled average multiplicity.
+struct uniform
+{
+  //! Constructs a uniform distribution tag with the requested average key multiplicity.
+  //!
+  //! @param __multiplicity Average number of generated keys that map to each unique key value.
+  //! @throws std::invalid_argument if `__multiplicity` is not positive.
+  explicit uniform(::nvbench::int64_t __multiplicity)
+      : __value{__multiplicity}
+  {
+    if (__multiplicity <= 0)
+    {
+      throw ::std::invalid_argument{"Multiplicity must be greater than 0"};
+    }
+  }
+
+  //! Average number of generated keys that map to each unique key value.
+  ::nvbench::int64_t __value;
+};
+} // namespace distribution
+
+namespace detail
+{
+template <typename _Tp, typename _Dist, typename _Rng>
+struct __generate_uniform_fn
+{
+  _CCCL_HOST_DEVICE_API constexpr __generate_uniform_fn(
+    ::cuda::std::size_t __num, _Dist __dist, ::cuda::std::size_t __seed)
+      : __num{__num}
+      , __dist{__dist}
+      , __seed{__seed}
+  {}
+
+  _CCCL_HOST_DEVICE_API constexpr _Tp operator()(::cuda::std::size_t __idx) const noexcept
+  {
+    _Rng __rng;
+    __rng.seed(__seed + __idx * 1664525ull + 1013904223ull);
+
+    const auto __num_unique_keys_unclamped = static_cast<::cuda::std::size_t>(
+      ::cuda::std::ceil(static_cast<double>(__num) / static_cast<double>(__dist.__value)));
+    const auto __num_unique_keys =
+      __num_unique_keys_unclamped < ::cuda::std::size_t{1} ? ::cuda::std::size_t{1} : __num_unique_keys_unclamped;
+    ::thrust::uniform_int_distribution<_Tp> __key_dist{_Tp{0}, static_cast<_Tp>(__num_unique_keys - 1)};
+    return __key_dist(__rng);
+  }
+
+  ::cuda::std::size_t __num;
+  _Dist __dist;
+  ::cuda::std::size_t __seed;
+};
+
+template <typename _Tp, typename _Rng>
+struct __dropout_fn
+{
+  _CCCL_HOST_DEVICE_API constexpr explicit __dropout_fn(::cuda::std::size_t __num)
+      : __num{__num}
+  {}
+
+  _CCCL_HOST_DEVICE_API _Tp operator()(::cuda::std::size_t __seed) const noexcept
+  {
+    _Rng __rng;
+    ::thrust::uniform_int_distribution<_Tp> __dist{static_cast<_Tp>(__num), ::cuda::std::numeric_limits<_Tp>::max()};
+    __rng.seed(__seed);
+    return __dist(__rng);
+  }
+
+  ::cuda::std::size_t __num;
+};
+
+template <typename _Rng>
+struct __dropout_pred
+{
+  _CCCL_HOST_DEVICE_API constexpr explicit __dropout_pred(double __keep_prob)
+      : __keep_prob{__keep_prob}
+  {}
+
+  _CCCL_HOST_DEVICE_API bool operator()(::cuda::std::size_t __seed) const noexcept
+  {
+    _Rng __rng;
+    ::thrust::uniform_real_distribution<double> __dist{0.0, 1.0};
+    __rng.seed(__seed);
+    return __dist(__rng) > __keep_prob;
+  }
+
+  double __keep_prob;
+};
+} // namespace detail
+
+//! Random key generator used by CUCO benchmarks.
+//!
+//! The generator defaults to `defaults::seed` to keep benchmark data reproducible across runs.
+//!
+//! @tparam _Rng Pseudo-random number generator type compatible with Thrust random distributions.
+template <typename _Rng = ::thrust::default_random_engine>
+class key_generator
+{
+public:
+  //! Constructs a key generator with the given seed.
+  //!
+  //! @param __seed Seed used to initialize the generator state.
+  explicit key_generator(::cuda::std::uint32_t __seed = defaults::seed)
+      : __rng{__seed}
+  {}
+
+  //! Generates keys according to the given distribution using the default device execution policy.
+  //!
+  //! @tparam _Dist Distribution tag type.
+  //! @tparam _OutputIt Output iterator type whose value type is the generated key type.
+  //! @param __dist Distribution tag controlling how keys are generated.
+  //! @param __out_begin Beginning of the output key range.
+  //! @param __out_end End of the output key range.
+  //! @throws std::invalid_argument if `_Dist` is not a supported distribution tag.
+  template <typename _Dist, typename _OutputIt>
+  void generate(_Dist __dist, _OutputIt __out_begin, _OutputIt __out_end)
+  {
+    generate(__dist, __out_begin, __out_end, ::thrust::device);
+  }
+
+  //! Generates keys according to the given distribution using the provided execution policy.
+  //!
+  //! @tparam _Dist Distribution tag type.
+  //! @tparam _OutputIt Output iterator type whose value type is the generated key type.
+  //! @tparam _ExecPolicy Thrust execution policy type.
+  //! @param __dist Distribution tag controlling how keys are generated.
+  //! @param __out_begin Beginning of the output key range.
+  //! @param __out_end End of the output key range.
+  //! @param __exec_policy Execution policy used for the underlying Thrust algorithms.
+  //! @throws std::invalid_argument if `_Dist` is not a supported distribution tag.
+  template <typename _Dist, typename _OutputIt, typename _ExecPolicy>
+  void generate(_Dist __dist, _OutputIt __out_begin, _OutputIt __out_end, _ExecPolicy __exec_policy)
+  {
+    using __value_type = typename ::cuda::std::iterator_traits<_OutputIt>::value_type;
+
+    if constexpr (::cuda::std::is_same_v<_Dist, distribution::unique>)
+    {
+      ::thrust::sequence(__exec_policy, __out_begin, __out_end, __value_type{0});
+      ::thrust::shuffle(__exec_policy, __out_begin, __out_end, __rng);
+    }
+    else if constexpr (::cuda::std::is_same_v<_Dist, distribution::uniform>)
+    {
+      const auto __num_keys = static_cast<::cuda::std::size_t>(::cuda::std::distance(__out_begin, __out_end));
+      const auto __seed     = static_cast<::cuda::std::size_t>(__rng());
+
+      ::thrust::transform(
+        __exec_policy,
+        ::cuda::counting_iterator<::cuda::std::size_t>{0},
+        ::cuda::counting_iterator<::cuda::std::size_t>{__num_keys},
+        __out_begin,
+        detail::__generate_uniform_fn<__value_type, _Dist, _Rng>{__num_keys, __dist, __seed});
+    }
+    else
+    {
+      throw ::std::invalid_argument{"Unexpected distribution type"};
+    }
+  }
+
+  //! Drops keys with probability `1 - keep_prob` using the default device execution policy.
+  //!
+  //! Replaced keys are sampled from `[N, max_key]`, where `N` is the number of keys in the range.
+  //!
+  //! @tparam _InOutIt Mutable iterator type whose value type is the key type.
+  //! @param __begin Beginning of the key range to update in place.
+  //! @param __end End of the key range to update in place.
+  //! @param __keep_prob Probability of keeping each original key. Must be in `[0, 1]`.
+  //! @throws std::invalid_argument if `__keep_prob` is outside `[0, 1]`.
+  template <typename _InOutIt>
+  void dropout(_InOutIt __begin, _InOutIt __end, double __keep_prob)
+  {
+    dropout(__begin, __end, __keep_prob, ::thrust::device);
+  }
+
+  //! Drops keys with probability `1 - keep_prob` using the provided execution policy.
+  //!
+  //! Replaced keys are sampled from `[N, max_key]`, where `N` is the number of keys in the range.
+  //!
+  //! @tparam _InOutIt Mutable iterator type whose value type is the key type.
+  //! @tparam _ExecPolicy Thrust execution policy type.
+  //! @param __begin Beginning of the key range to update in place.
+  //! @param __end End of the key range to update in place.
+  //! @param __keep_prob Probability of keeping each original key. Must be in `[0, 1]`.
+  //! @param __exec_policy Execution policy used for the underlying Thrust algorithms.
+  //! @throws std::invalid_argument if `__keep_prob` is outside `[0, 1]`.
+  template <typename _InOutIt, typename _ExecPolicy>
+  void dropout(_InOutIt __begin, _InOutIt __end, double __keep_prob, _ExecPolicy __exec_policy)
+  {
+    using __value_type = typename ::cuda::std::iterator_traits<_InOutIt>::value_type;
+
+    if (__keep_prob < 0.0 || __keep_prob > 1.0)
+    {
+      throw ::std::invalid_argument{"Probability needs to be between 0 and 1"};
+    }
+
+    if (__keep_prob < 1.0)
+    {
+      const auto __num_keys = static_cast<::cuda::std::size_t>(::cuda::std::distance(__begin, __end));
+      ::cuda::counting_iterator<::cuda::std::size_t> __seeds{static_cast<::cuda::std::size_t>(__rng())};
+
+      ::thrust::transform_if(
+        __exec_policy,
+        __seeds,
+        __seeds + __num_keys,
+        __begin,
+        detail::__dropout_fn<__value_type, _Rng>{__num_keys},
+        detail::__dropout_pred<_Rng>{__keep_prob});
+    }
+
+    ::thrust::shuffle(__exec_policy, __begin, __end, __rng);
+  }
+
+private:
+  _Rng __rng;
+};
+
+//! Constructs the requested distribution tag from NVBench axis values.
+//!
+//! `distribution::uniform` reads the `Multiplicity` axis from `state`.
+//!
+//! @tparam _Dist Distribution tag type to construct.
+//! @param __state NVBench state containing distribution-specific axis values.
+//! @return Distribution tag initialized from the benchmark state.
+//! @throws std::invalid_argument if `_Dist` is not a supported distribution tag.
+template <typename _Dist>
+_Dist dist_from_state(::nvbench::state const& __state)
+{
+  if constexpr (::cuda::std::is_same_v<_Dist, distribution::unique>)
+  {
+    return _Dist{};
+  }
+  else if constexpr (::cuda::std::is_same_v<_Dist, distribution::uniform>)
+  {
+    return _Dist{__state.get_int64("Multiplicity")};
+  }
+  else
+  {
+    throw ::std::invalid_argument{"Unexpected distribution type"};
+  }
+}
+} // namespace cuda::experimental::cuco::benchmark
+
+NVBENCH_DECLARE_TYPE_STRINGS(
+  ::cuda::experimental::cuco::benchmark::distribution::unique, "UNIQUE", "distribution::unique");
+NVBENCH_DECLARE_TYPE_STRINGS(
+  ::cuda::experimental::cuco::benchmark::distribution::uniform, "UNIFORM", "distribution::uniform");