[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<double>{1.0, 2.0, 4.0, 8.0, 16.0};
+//! 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,303 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 <thrust/execution_policy.h>
+#include <thrust/random.h>
+#include <thrust/sequence.h>
+#include <thrust/shuffle.h>
+#include <thrust/transform.h>
+#include <thrust/type_traits/is_contiguous_iterator.h>
+
+#include <cuda/iterator>
+#include <cuda/std/cmath>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/iterator>
+#include <cuda/std/limits>
+#include <cuda/std/memory>
+#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 finite or is less than 1.0.
+  explicit uniform(double multiplicity)
+      : multiplicity{multiplicity}
+  {
+    if (!cuda::std::isfinite(multiplicity) || multiplicity < 1.0)
+    {
+      throw ::std::invalid_argument{"Multiplicity must be finite and at least 1"};
+    }
+  }
+
+  //! Average number of generated keys that map to each unique key value.
+  double multiplicity;
+};
+} // namespace distribution
+
+namespace detail
+{
+template <typename ExecPolicy, typename RandomIt, typename Rng>
+void shuffle(ExecPolicy exec_policy, RandomIt begin, RandomIt end, Rng& rng)
+{
+  const auto num_items = cuda::std::distance(begin, end);
+  if constexpr (thrust::is_contiguous_iterator_v<RandomIt>)
+  {
+    // TODO(NVIDIA/cccl#9755): Temporary workaround for thrust::shuffle's &*it contiguous unwrap.
+    auto shuffle_begin = cuda::std::to_address(begin);
+    thrust::shuffle(exec_policy, shuffle_begin, shuffle_begin + num_items, rng);
+  }
+  else
+  {
+    thrust::shuffle(exec_policy, begin, end, rng);
+  }
+}
+
+template <typename Key, typename Distribution, typename Rng>
+struct generate_uniform_fn
+{
+  __host__ __device__ constexpr generate_uniform_fn(cuda::std::size_t num, Distribution dist, cuda::std::size_t seed)
+      : num{num}
+      , dist{dist}
+      , seed{seed}
+  {}
+
+  __host__ __device__ constexpr Key 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) / dist.multiplicity));
+    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<Key> key_dist{Key{0}, static_cast<Key>(num_unique_keys - 1)};
+    return key_dist(rng);
+  }
+
+  cuda::std::size_t num;
+  Distribution dist;
+  cuda::std::size_t seed;
+};
+
+template <typename Key, typename Rng>
+struct dropout_fn
+{
+  __host__ __device__ constexpr explicit dropout_fn(cuda::std::size_t num)
+      : num{num}
+  {}
+
+  __host__ __device__ Key operator()(cuda::std::size_t seed) const noexcept
+  {
+    Rng rng;
+    thrust::uniform_int_distribution<Key> dist{static_cast<Key>(num), cuda::std::numeric_limits<Key>::max()};
+    rng.seed(seed);
+    return dist(rng);
+  }
+
+  cuda::std::size_t num;
+};
+
+template <typename Rng>
+struct dropout_pred
+{
+  __host__ __device__ constexpr explicit dropout_pred(double keep_prob)
+      : keep_prob{keep_prob}
+  {}
+
+  __host__ __device__ 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 Distribution 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 `Distribution` is not a supported distribution tag.
+  template <typename Distribution, typename OutputIt>
+  void generate(Distribution 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 Distribution 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 `Distribution` is not a supported distribution tag.
+  template <typename Distribution, typename OutputIt, typename ExecPolicy>
+  void generate(Distribution 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<Distribution, distribution::unique>)
+    {
+      thrust::sequence(exec_policy, out_begin, out_end, value_type{0});
+      detail::shuffle(exec_policy, out_begin, out_end, rng);
+    }
+    else if constexpr (cuda::std::is_same_v<Distribution, 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, Distribution, 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.
+  //!
+  //! The full range is shuffled afterward, even when all keys are kept.
+  //!
+  //! @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.
+  //!
+  //! The full range is shuffled afterward, even when all keys are kept.
+  //!
+  //! @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});
+    }
+
+    detail::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 Distribution 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 `Distribution` is not a supported distribution tag.
+template <typename Distribution>
+Distribution dist_from_state(nvbench::state const& state)
+{
+  if constexpr (cuda::std::is_same_v<Distribution, distribution::unique>)
+  {
+    return Distribution{};
+  }
+  else if constexpr (cuda::std::is_same_v<Distribution, distribution::uniform>)
+  {
+    return Distribution{state.get_float64("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");