Add b200 policies for cub.device.partition.flagged,if (#3617)

Co-authored-by: gonidelis <ggonidelis@nvidia.com>
NVIDIA · Feb 7, 2025 · 8072358 · 8072358
1 parent f4248dd
commit 8072358
Show file tree

Hide file tree

Showing 3 changed files with 775 additions and 58 deletions.
diff --git a/cub/cub/agent/agent_select_if.cuh b/cub/cub/agent/agent_select_if.cuh
@@ -137,6 +137,15 @@ struct partition_distinct_output_t
   rejected_iterator_t rejected_it;
 };
 
+template <typename OutputIterator>
+struct is_partition_distinct_output_t : ::cuda::std::false_type
+{};
+
+template <typename SelectedOutputItT, typename RejectedOutputItT>
+struct is_partition_distinct_output_t<partition_distinct_output_t<SelectedOutputItT, RejectedOutputItT>>
+    : ::cuda::std::true_type
+{};
+
 /**
  * @brief AgentSelectIf implements a stateful abstraction of CUDA thread blocks for participating in
  * device-wide selection

diff --git a/cub/cub/device/dispatch/dispatch_select_if.cuh b/cub/cub/device/dispatch/dispatch_select_if.cuh
@@ -407,23 +407,27 @@ __launch_bounds__(int(
  * @tparam OffsetT
  *   Signed integer type for global offsets
  *
- * @tparam SelectImpl SelectionOpt
+ * @tparam SelectionOpt
  *   SelectImpl indicating whether to partition, just selection or selection where the memory for the input and
  *   output may alias each other.
  */
-template <typename InputIteratorT,
-          typename FlagsInputIteratorT,
-          typename SelectedOutputIteratorT,
-          typename NumSelectedIteratorT,
-          typename SelectOpT,
-          typename EqualityOpT,
-          typename OffsetT,
-          SelectImpl SelectionOpt,
-          typename PolicyHub = detail::select::policy_hub<cub::detail::value_t<InputIteratorT>,
-                                                          cub::detail::value_t<FlagsInputIteratorT>,
-                                                          detail::select::per_partition_offset_t,
-                                                          (SelectionOpt == SelectImpl::SelectPotentiallyInPlace),
-                                                          (SelectionOpt == SelectImpl::Partition)>>
+template <
+  typename InputIteratorT,
+  typename FlagsInputIteratorT,
+  typename SelectedOutputIteratorT,
+  typename NumSelectedIteratorT,
+  typename SelectOpT,
+  typename EqualityOpT,
+  typename OffsetT,
+  SelectImpl SelectionOpt,
+  typename PolicyHub = detail::select::policy_hub<
+    detail::value_t<InputIteratorT>,
+    detail::value_t<FlagsInputIteratorT>,
+    // if/flagged/unique only have a single code path for different offset types, partition has different code paths
+    ::cuda::std::conditional_t<SelectionOpt == SelectImpl::Partition, OffsetT, detail::select::per_partition_offset_t>,
+    detail::select::is_partition_distinct_output_t<SelectedOutputIteratorT>::value,
+    (SelectionOpt == SelectImpl::SelectPotentiallyInPlace),
+    (SelectionOpt == SelectImpl::Partition)>>
 struct DispatchSelectIf
 {
   /******************************************************************************