From 9a38c1e66d405f32cf9aa403c2a67c1bf8bdaa15 Mon Sep 17 00:00:00 2001
From: Georgy Evtushenko <evtushenko.georgy@gmail.com>
Date: Sat, 6 May 2023 14:35:49 +0400
Subject: [PATCH] Almost there

---
 benchmarks/bench/select/flagged.cu          |  7 +--
 benchmarks/nvbench_helper/nvbench_helper.cu | 51 +++++++++++++++++++++
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/benchmarks/bench/select/flagged.cu b/benchmarks/bench/select/flagged.cu
index d104d5467..3473f7672 100644
--- a/benchmarks/bench/select/flagged.cu
+++ b/benchmarks/bench/select/flagged.cu
@@ -34,6 +34,9 @@ void select(nvbench::state &state, nvbench::type_list<T, OffsetT>)
   thrust::device_vector<bool> flags(elements);
   thrust::device_vector<offset_t> num_selected(1);
 
+  gen(seed_t{}, in);
+  gen(seed_t{1}, flags, entropy);
+
   // TODO Extract into helper TU
   const auto selected_elements = thrust::count(flags.cbegin(), flags.cend(), true);
   thrust::device_vector<T> out(selected_elements);
@@ -43,8 +46,6 @@ void select(nvbench::state &state, nvbench::type_list<T, OffsetT>)
   output_it_t d_out = thrust::raw_pointer_cast(out.data());
   num_selected_it_t d_num_selected = thrust::raw_pointer_cast(num_selected.data());
 
-  gen(seed_t{}, in, entropy);
-
   state.add_element_count(elements);
   state.add_global_memory_reads<T>(elements);
   state.add_global_memory_reads<bool>(elements);
@@ -84,4 +85,4 @@ NVBENCH_BENCH_TYPES(select, NVBENCH_TYPE_AXES(fundamental_types, offset_types))
   .set_name("cub::DeviceSelect::Flagged")
   .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
   .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
-  .add_string_axis("Entropy", {"1.000", "0.544", "0.201"});
+  .add_string_axis("Entropy", {"1.000", "0.544", "0.000"});
diff --git a/benchmarks/nvbench_helper/nvbench_helper.cu b/benchmarks/nvbench_helper/nvbench_helper.cu
index 0ba237364..665e58f71 100644
--- a/benchmarks/nvbench_helper/nvbench_helper.cu
+++ b/benchmarks/nvbench_helper/nvbench_helper.cu
@@ -265,6 +265,55 @@ void generator_t::operator()(seed_t seed,
   };
 }
 
+double entropy_to_probability(bit_entropy entropy)
+{
+  switch (entropy)
+  {
+    case bit_entropy::_0_000: return 0.0;
+    case bit_entropy::_0_811: return 0.811;
+    case bit_entropy::_0_544: return 0.544;
+    case bit_entropy::_0_337: return 0.337;
+    case bit_entropy::_0_201: return 0.201;
+    case bit_entropy::_1_000: return 1.0;
+    default: return 0.0;
+  }
+}
+
+struct random_to_probability_t
+{
+  double m_probability;
+
+  __host__ __device__ bool operator()(double random_value)
+  {
+    return random_value < m_probability;
+  }
+};
+
+template <>
+void generator_t::operator()(seed_t seed,
+                             thrust::device_vector<bool> &data,
+                             bit_entropy entropy,
+                             bool /* min */,
+                             bool /* max */)
+{
+  if (entropy == bit_entropy::_0_000)
+  {
+    thrust::fill(data.begin(), data.end(), false);
+  }
+  else if (entropy == bit_entropy::_1_000)
+  {
+    thrust::fill(data.begin(), data.end(), true);
+  }
+  else
+  {
+    prepare_random_generator(seed, data.size());
+    thrust::transform(m_distribution.begin(),
+                      m_distribution.end(),
+                      data.begin(),
+                      random_to_probability_t{entropy_to_probability(entropy)});
+  }
+}
+
 template <typename T>
 void gen(seed_t seed, thrust::device_vector<T> &data, bit_entropy entropy, T min, T max)
 {
@@ -280,6 +329,8 @@ void gen(seed_t seed, thrust::device_vector<T> &data, bit_entropy entropy, T min
 
 #define INSTANTIATE(TYPE) INSTANTIATE_RND(TYPE);
 
+INSTANTIATE(bool);
+
 INSTANTIATE(uint8_t);
 INSTANTIATE(uint16_t);
 INSTANTIATE(uint32_t);