Adds support for large number of items in DeviceSelect and DevicePartition

cub/benchmarks/bench/partition/flagged.cu

@@ -27,6 +27,10 @@
 
 #include <cub/device/device_partition.cuh>
 
+#include <thrust/count.h>
+
+#include <cuda/std/type_traits>
+
 #include <look_back_helper.cuh>
 #include <nvbench_helper.cuh>
 
@@ -77,16 +81,35 @@ struct policy_hub_t
 };
 #endif // TUNE_BASE
 
-template <typename T, typename OffsetT>
-void flagged(nvbench::state& state, nvbench::type_list<T, OffsetT>)
+template <typename FlagsItT, typename T, typename OffsetT>
+void init_output_partition_buffer(
+  FlagsItT d_flags,
+  OffsetT num_items,
+  T* d_out,
+  cub::detail::partition_distinct_output_t<T*, T*>& d_partition_out_buffer)
 {
-  using input_it_t        = const T*;
-  using flag_it_t         = const bool*;
-  using output_it_t       = T*;
-  using num_selected_it_t = OffsetT*;
-  using select_op_t       = cub::NullType;
-  using equality_op_t     = cub::NullType;
-  using offset_t          = OffsetT;
+  const auto selected_elements = thrust::count(d_flags, d_flags + num_items, true);
+  d_partition_out_buffer       = cub::detail::partition_distinct_output_t<T*, T*>{d_out, d_out + selected_elements};
+}
+
+template <typename FlagsItT, typename T, typename OffsetT>
+void init_output_partition_buffer(FlagsItT, OffsetT, T* d_out, T*& d_partition_out_buffer)
+{
+  d_partition_out_buffer = d_out;
+}
+
+template <typename T, typename OffsetT, typename UseDistinctPartitionT>
+void flagged(nvbench::state& state, nvbench::type_list<T, OffsetT, UseDistinctPartitionT>)
+{
+  using input_it_t                           = const T*;
+  using flag_it_t                            = const bool*;
+  using num_selected_it_t                    = OffsetT*;
+  using select_op_t                          = cub::NullType;
+  using equality_op_t                        = cub::NullType;
+  using offset_t                             = OffsetT;
+  constexpr bool use_distinct_out_partitions = UseDistinctPartitionT::value;
+  using output_it_t                          = typename ::cuda::std::
+    conditional<use_distinct_out_partitions, cub::detail::partition_distinct_output_t<T*, T*>, T*>::type;
 
 #if !TUNE_BASE
   using policy_t   = policy_hub_t<T>;
@@ -127,8 +150,9 @@ void flagged(nvbench::state& state, nvbench::type_list<T, OffsetT>)
 
   input_it_t d_in                  = thrust::raw_pointer_cast(in.data());
   flag_it_t d_flags                = thrust::raw_pointer_cast(flags.data());
-  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());
+  output_it_t d_out{};
+  init_output_partition_buffer(flags.cbegin(), elements, thrust::raw_pointer_cast(out.data()), d_out);
 
   state.add_element_count(elements);
   state.add_global_memory_reads<T>(elements);
@@ -158,8 +182,10 @@ void flagged(nvbench::state& state, nvbench::type_list<T, OffsetT>)
   });
 }
 
-NVBENCH_BENCH_TYPES(flagged, NVBENCH_TYPE_AXES(fundamental_types, offset_types))
+using distinct_partitions = nvbench::type_list<::cuda::std::false_type, ::cuda::std::true_type>;
+
+NVBENCH_BENCH_TYPES(flagged, NVBENCH_TYPE_AXES(fundamental_types, offset_types, distinct_partitions))
   .set_name("base")
-  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}", "DistinctPartitions{ct}"})
   .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
   .add_string_axis("Entropy", {"1.000", "0.544", "0.000"});

cub/benchmarks/bench/partition/if.cu

@@ -27,6 +27,10 @@
 
 #include <cub/device/device_partition.cuh>
 
+#include <thrust/count.h>
+
+#include <cuda/std/type_traits>
+
 #include <look_back_helper.cuh>
 #include <nvbench_helper.cuh>
 
@@ -102,16 +106,36 @@ T value_from_entropy(double percentage)
   return static_cast<T>(result);
 }
 
-template <typename T, typename OffsetT>
-void partition(nvbench::state& state, nvbench::type_list<T, OffsetT>)
+template <typename InItT, typename T, typename OffsetT, typename SelectOpT>
+void init_output_partition_buffer(
+  InItT d_in,
+  OffsetT num_items,
+  T* d_out,
+  SelectOpT select_op,
+  cub::detail::partition_distinct_output_t<T*, T*>& d_partition_out_buffer)
 {
-  using input_it_t        = const T*;
-  using flag_it_t         = cub::NullType*;
-  using output_it_t       = T*;
-  using num_selected_it_t = OffsetT*;
-  using select_op_t       = less_then_t<T>;
-  using equality_op_t     = cub::NullType;
-  using offset_t          = OffsetT;
+  const auto selected_elements = thrust::count_if(d_in, d_in + num_items, select_op);
+  d_partition_out_buffer       = cub::detail::partition_distinct_output_t<T*, T*>{d_out, d_out + selected_elements};
+}
+
+template <typename InItT, typename T, typename OffsetT, typename SelectOpT>
+void init_output_partition_buffer(InItT, OffsetT, T* d_out, SelectOpT, T*& d_partition_out_buffer)
+{
+  d_partition_out_buffer = d_out;
+}
+
+template <typename T, typename OffsetT, typename UseDistinctPartitionT>
+void partition(nvbench::state& state, nvbench::type_list<T, OffsetT, UseDistinctPartitionT>)
+{
+  using input_it_t                           = const T*;
+  using flag_it_t                            = cub::NullType*;
+  using num_selected_it_t                    = OffsetT*;
+  using select_op_t                          = less_then_t<T>;
+  using equality_op_t                        = cub::NullType;
+  using offset_t                             = OffsetT;
+  constexpr bool use_distinct_out_partitions = UseDistinctPartitionT::value;
+  using output_it_t                          = typename ::cuda::std::
+    conditional<use_distinct_out_partitions, cub::detail::partition_distinct_output_t<T*, T*>, T*>::type;
 
 #if !TUNE_BASE
   using policy_t   = policy_hub_t<T>;
@@ -153,8 +177,9 @@ void partition(nvbench::state& state, nvbench::type_list<T, OffsetT>)
 
   input_it_t d_in                  = thrust::raw_pointer_cast(in.data());
   flag_it_t d_flags                = nullptr;
-  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());
+  output_it_t d_out{};
+  init_output_partition_buffer(in.cbegin(), elements, thrust::raw_pointer_cast(out.data()), select_op, d_out);
 
   state.add_element_count(elements);
   state.add_global_memory_reads<T>(elements);
@@ -183,8 +208,10 @@ void partition(nvbench::state& state, nvbench::type_list<T, OffsetT>)
   });
 }
 
-NVBENCH_BENCH_TYPES(partition, NVBENCH_TYPE_AXES(fundamental_types, offset_types))
+using distinct_partitions = nvbench::type_list<::cuda::std::false_type, ::cuda::std::true_type>;
+
+NVBENCH_BENCH_TYPES(partition, NVBENCH_TYPE_AXES(fundamental_types, offset_types, distinct_partitions))
   .set_name("base")
-  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}", "DistinctPartitions{ct}"})
   .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
   .add_string_axis("Entropy", {"1.000", "0.544", "0.000"});