Support segmented reductions and null mask reductions

cpp/CMakeLists.txt

@@ -359,12 +359,19 @@ add_library(
   src/reductions/nth_element.cu
   src/reductions/product.cu
   src/reductions/reductions.cpp
+  src/reductions/segmented_reductions.cpp
   src/reductions/scan/rank_scan.cu
   src/reductions/scan/scan.cpp
   src/reductions/scan/scan_exclusive.cu
   src/reductions/scan/scan_inclusive.cu
   src/reductions/std.cu
   src/reductions/sum.cu
+  src/reductions/segmented_all.cu
+  src/reductions/segmented_any.cu
+  src/reductions/segmented_sum.cu
+  src/reductions/segmented_product.cu
+  src/reductions/segmented_max.cu
+  src/reductions/segmented_min.cu
   src/reductions/sum_of_squares.cu
   src/reductions/var.cu
   src/replace/clamp.cu

cpp/benchmarks/CMakeLists.txt

@@ -24,7 +24,7 @@ target_compile_options(
 
 target_link_libraries(
   cudf_datagen PUBLIC GTest::gmock GTest::gtest GTest::gmock_main GTest::gtest_main
-                      benchmark::benchmark nvbench::nvbench Threads::Threads cudf
+                      benchmark::benchmark nvbench::nvbench Threads::Threads cudf cudftestutil
 )
 
 target_include_directories(
@@ -156,8 +156,13 @@ ConfigureBench(TYPE_DISPATCHER_BENCH type_dispatcher/type_dispatcher_benchmark.c
 # ##################################################################################################
 # * reduction benchmark ---------------------------------------------------------------------------
 ConfigureBench(
-  REDUCTION_BENCH reduction/anyall_benchmark.cpp reduction/dictionary_benchmark.cpp
-  reduction/reduce_benchmark.cpp reduction/scan_benchmark.cpp reduction/minmax_benchmark.cpp
+  REDUCTION_BENCH
+  reduction/anyall_benchmark.cpp
+  reduction/dictionary_benchmark.cpp
+  reduction/reduce_benchmark.cpp
+  reduction/scan_benchmark.cpp
+  reduction/minmax_benchmark.cpp
+  reduction/segment_reduce_benchmark.cpp
 )
 
 # ##################################################################################################

cpp/benchmarks/reduction/segment_reduce_benchmark.cpp

@@ -0,0 +1,86 @@
+#include "cudf/column/column.hpp"
+#include <cudf/aggregation.hpp>
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/iterator.cuh>
+#include <cudf/reduction.hpp>
+#include <cudf/types.hpp>
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_utilities.hpp>
+#include <cudf_test/column_wrapper.hpp>
+#include <fixture/benchmark_fixture.hpp>
+#include <memory>
+#include <synchronization/synchronization.hpp>
+#include <unordered_map>
+#include <utility>
+
+namespace cudf {
+
+class SegmentReduction : public cudf::benchmark {
+};
+
+template <typename InputType>
+void BM_Segment_Reduction_Simple(::benchmark::State& state, std::unique_ptr<aggregation> const& agg)
+{
+  auto const column_size{size_type(state.range(0))};
+  auto const num_segments{size_type(state.range(1))};
+  auto const segment_length = column_size / num_segments;
+
+  test::UniformRandomGenerator<InputType> rand_gen(0, 100);
+  auto data_it = detail::make_counting_transform_iterator(
+    0, [&rand_gen](auto i) { return rand_gen.generate(); });
+
+  auto offset_it = detail::make_counting_transform_iterator(
+    0,
+    [&column_size, &segment_length](auto i) { return std::min(column_size, i * segment_length); });
+
+  test::fixed_width_column_wrapper<InputType> input(data_it, data_it + column_size);
+  test::fixed_width_column_wrapper<size_type> offsets(offset_it, offset_it + num_segments + 1);
+
+  for (auto _ : state) {
+    cuda_event_timer timer(state, true);
+    auto result = cudf::segmented_reduce(input,
+                                         offsets,
+                                         make_sum_aggregation(),
+                                         data_type{type_to_id<InputType>()},
+                                         null_policy::EXCLUDE);
+  }
+}
+
+#define concat(a, b, c) a##b##c
+#define get_agg(op)     concat(cudf::make_, op, _aggregation())
+
+#define SRBM_DEFINE(name, type, op)                                                            \
+  BENCHMARK_DEFINE_F(SegmentReduction, name)(::benchmark::State & state)                       \
+  {                                                                                            \
+    BM_Segment_Reduction_Simple<type>(state, get_agg(op));                                     \
+  }                                                                                            \
+  BENCHMARK_REGISTER_F(SegmentReduction, name)                                                 \
+    ->UseManualTime()                                                                          \
+    ->ArgsProduct(                                                                             \
+      {{1 << 14, /*10k*/ 1 << 17, /*100k*/ 1 << 20, /*1M*/ 1 << 24, /*10M*/ 1 << 27 /*100M*/}, \
+       {                                                                                       \
+         128,                                                                                  \
+         1024,                                                                                 \
+       }});
+
+#define REDUCE_BENCHMARK_DEFINE(type, aggregation) \
+  SRBM_DEFINE(concat(type, _, aggregation), type, aggregation)
+
+#define REDUCE_BENCHMARK_NUMERIC(aggregation)    \
+  REDUCE_BENCHMARK_DEFINE(bool, aggregation);    \
+  REDUCE_BENCHMARK_DEFINE(int8_t, aggregation);  \
+  REDUCE_BENCHMARK_DEFINE(int16_t, aggregation); \
+  REDUCE_BENCHMARK_DEFINE(int32_t, aggregation); \
+  REDUCE_BENCHMARK_DEFINE(int64_t, aggregation); \
+  REDUCE_BENCHMARK_DEFINE(float, aggregation);   \
+  REDUCE_BENCHMARK_DEFINE(double, aggregation);
+
+REDUCE_BENCHMARK_NUMERIC(sum)
+REDUCE_BENCHMARK_NUMERIC(product)
+REDUCE_BENCHMARK_NUMERIC(min)
+REDUCE_BENCHMARK_NUMERIC(max)
+REDUCE_BENCHMARK_NUMERIC(any)
+REDUCE_BENCHMARK_NUMERIC(all)
+
+}  // namespace cudf

cpp/include/cudf/detail/null_mask.cuh

@@ -18,6 +18,7 @@
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/detail/utilities/cuda.cuh>
 #include <cudf/detail/utilities/vector_factories.hpp>
+#include <cudf/detail/valid_if.cuh>
 #include <cudf/null_mask.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/span.hpp>
@@ -279,7 +280,8 @@ rmm::device_uvector<size_type> segmented_count_bits(bitmask_type const* bitmask,
                                                     OffsetIterator first_bit_indices_end,
                                                     OffsetIterator last_bit_indices_begin,
                                                     count_bits_policy count_bits,
-                                                    rmm::cuda_stream_view stream)
+                                                    rmm::cuda_stream_view stream,
+                                                    rmm::mr::device_memory_resource* mr)
 {
   auto const num_ranges =
     static_cast<size_type>(std::distance(first_bit_indices_begin, first_bit_indices_end));
@@ -329,14 +331,15 @@ rmm::device_uvector<size_type> segmented_count_bits(bitmask_type const* bitmask,
     // set bits from the length of the segment.
     auto segments_begin =
       thrust::make_zip_iterator(first_bit_indices_begin, last_bit_indices_begin);
-    auto segments_size = thrust::transform_iterator(segments_begin, [] __device__(auto segment) {
-      auto const begin = thrust::get<0>(segment);
-      auto const end   = thrust::get<1>(segment);
-      return end - begin;
-    });
+    auto segment_length_iterator =
+      thrust::transform_iterator(segments_begin, [] __device__(auto const& segment) {
+        auto const begin = thrust::get<0>(segment);
+        auto const end   = thrust::get<1>(segment);
+        return end - begin;
+      });
     thrust::transform(rmm::exec_policy(stream),
-                      segments_size,
-                      segments_size + num_ranges,
+                      segment_length_iterator,
+                      segment_length_iterator + num_ranges,
                       d_bit_counts.data(),
                       d_bit_counts.data(),
                       [] __device__(auto segment_size, auto segment_bit_count) {
@@ -438,7 +441,8 @@ std::vector<size_type> segmented_count_bits(bitmask_type const* bitmask,
                                        first_bit_indices_end,
                                        last_bit_indices_begin,
                                        count_bits,
-                                       stream);
+                                       stream,
+                                       rmm::mr::get_current_device_resource());
 
   // Copy the results back to the host.
   return make_std_vector_sync(d_bit_counts, stream);
@@ -501,6 +505,111 @@ std::vector<size_type> segmented_null_count(bitmask_type const* bitmask,
   return detail::segmented_count_unset_bits(bitmask, indices_begin, indices_end, stream);
 }
 
+/**
+ * @brief Reduce an input null mask using segments defined by offset indices
+ * into an output null mask.
+ *
+ * @tparam OffsetIterator Random-access input iterator type.
+ * @param bitmask Null mask residing in device memory whose segments will be
+ * reduced into a new mask.
+ * @param first_bit_indices_begin Random-access input iterator to the beginning
+ * of a sequence of indices of the first bit in each segment (inclusive).
+ * @param first_bit_indices_end Random-access input iterator to the end of a
+ * sequence of indices of the first bit in each segment (inclusive).
+ * @param last_bit_indices_begin Random-access input iterator to the beginning
+ * of a sequence of indices of the last bit in each segment (exclusive).
+ * @param null_handling If `null_policy::INCLUDE`, all elements in a segment
+ * must be valid for the reduced value to be valid. If `null_policy::EXCLUDE`,
+ * the reduction is valid if any element in the segment is valid.
+ * @param stream CUDA stream used for device memory operations and kernel launches.
+ * @param mr Device memory resource used to allocate the returned buffer's device memory.
+ * @return A pair containing the reduced null mask and number of nulls.
+ */
+template <typename OffsetIterator>
+std::pair<rmm::device_buffer, size_type> segmented_null_mask_reduction(
+  bitmask_type const* bitmask,
+  OffsetIterator first_bit_indices_begin,
+  OffsetIterator first_bit_indices_end,
+  OffsetIterator last_bit_indices_begin,
+  null_policy null_handling,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr)
+{
+  auto const segments_begin =
+    thrust::make_zip_iterator(first_bit_indices_begin, last_bit_indices_begin);
+  auto const segment_length_iterator =
+    thrust::make_transform_iterator(segments_begin, [] __device__(auto const& segment) {
+      auto const begin = thrust::get<0>(segment);
+      auto const end   = thrust::get<1>(segment);
+      return end - begin;
+    });
+
+  // Empty segments are always null in the output mask
+  auto const num_segments =
+    static_cast<size_type>(std::distance(first_bit_indices_begin, first_bit_indices_end));
+  auto [output_null_mask, output_null_count] = cudf::detail::valid_if(
+    segment_length_iterator,
+    segment_length_iterator + num_segments,
+    [] __device__(auto const& len) { return len > 0; },
+    stream,
+    mr);
+
+  if (bitmask != nullptr) {
+    [[maybe_unused]] auto const [null_policy_bitmask, _] = [&]() {
+      if (null_handling == null_policy::EXCLUDE) {
+        // Output null mask should be valid if any element in the segment is
+        // valid and the segment is non-empty.
+        auto const valid_counts =
+          cudf::detail::segmented_count_bits(bitmask,
+                                             first_bit_indices_begin,
+                                             first_bit_indices_end,
+                                             last_bit_indices_begin,
+                                             cudf::detail::count_bits_policy::SET_BITS,
+                                             stream,
+                                             rmm::mr::get_current_device_resource());
+        return cudf::detail::valid_if(
+          valid_counts.begin(),
+          valid_counts.end(),
+          [] __device__(auto const valid_count) { return valid_count > 0; },
+          stream);
+      } else {
+        // Output null mask should be valid if all elements in the segment are
+        // valid and the segment is non-empty.
+        auto const null_counts =
+          cudf::detail::segmented_count_bits(bitmask,
+                                             first_bit_indices_begin,
+                                             first_bit_indices_end,
+                                             last_bit_indices_begin,
+                                             cudf::detail::count_bits_policy::UNSET_BITS,
+                                             stream,
+                                             rmm::mr::get_current_device_resource());
+        return cudf::detail::valid_if(
+          null_counts.begin(),
+          null_counts.end(),
+          [] __device__(auto const null_count) { return null_count == 0; },
+          stream);
+      }
+    }();
+
+    std::vector<bitmask_type const*> masks{
+      reinterpret_cast<bitmask_type const*>(output_null_mask.data()),
+      reinterpret_cast<bitmask_type const*>(null_policy_bitmask.data())};
+    std::vector<size_type> begin_bits{0, 0};
+    cudf::detail::inplace_bitmask_and(
+      device_span<bitmask_type>(reinterpret_cast<bitmask_type*>(output_null_mask.data()),
+                                num_bitmask_words(num_segments)),
+      masks,
+      begin_bits,
+      num_segments,
+      stream,
+      mr);
+
+    // TODO: inplace_bitmask_and should return its null count (PR 9904)
+    output_null_count = cudf::UNKNOWN_NULL_COUNT;
+  }
+  return std::make_pair(std::move(output_null_mask), output_null_count);
+}
+
 }  // namespace detail
 
 }  // namespace cudf

cpp/include/cudf/detail/reduction.cuh

@@ -18,6 +18,7 @@
 
 #include "reduction_operators.cuh"
 
+#include <cudf/column/column_factories.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -26,6 +27,7 @@
 #include <rmm/exec_policy.hpp>
 
 #include <cub/device/device_reduce.cuh>
+#include <cub/device/device_segmented_reduce.cuh>
 
 #include <thrust/for_each.h>
 #include <thrust/iterator/iterator_traits.h>
@@ -218,6 +220,85 @@ std::unique_ptr<scalar> reduce(InputIterator d_in,
   return std::unique_ptr<scalar>(result);
 }
 
+/**
+ * @brief Compute the specified simple reduction over each of the segments in the
+ * input range of elements.
+ *
+ * @param[in] d_in          the begin iterator to input
+ * @param[in] d_offset      the begin iterator to offset
+ * @param[in] num_segments  the number of segments
+ * @param[in] op            the reduction operator
+ * @param[in] stream        CUDA stream used for device memory operations and kernel launches.
+ * @returns   Output scalar in device memory
+ *
+ * @tparam Op               the reduction operator with device binary operator
+ * @tparam InputIterator    the input column iterator
+ * @tparam OutputType       the output type of reduction
+ */
+template <typename Op,
+          typename InputIterator,
+          typename OffsetIterator,
+          typename OutputType = typename thrust::iterator_value<InputIterator>::type,
+          typename std::enable_if_t<is_fixed_width<OutputType>() &&
+                                    not cudf::is_fixed_point<OutputType>()>* = nullptr>
+std::unique_ptr<column> segmented_reduce(InputIterator d_in,
+                                         OffsetIterator d_offset,
+                                         cudf::size_type num_segments,
+                                         op::simple_op<Op> sop,
+                                         rmm::cuda_stream_view stream,
+                                         rmm::mr::device_memory_resource* mr)
+{
+  auto binary_op  = sop.get_binary_op();
+  auto identity   = sop.template get_identity<OutputType>();
+  auto dev_result = make_fixed_width_column(
+    data_type{type_to_id<OutputType>()}, num_segments, mask_state::UNALLOCATED, stream, mr);
+  auto dev_result_mview = dev_result->mutable_view();
+
+  // Allocate temporary storage
+  rmm::device_buffer d_temp_storage;
+  size_t temp_storage_bytes = 0;
+  cub::DeviceSegmentedReduce::Reduce(d_temp_storage.data(),
+                                     temp_storage_bytes,
+                                     d_in,
+                                     dev_result_mview.data<OutputType>(),
+                                     num_segments,
+                                     d_offset,
+                                     d_offset + 1,
+                                     binary_op,
+                                     identity,
+                                     stream.value());
+  d_temp_storage = rmm::device_buffer{temp_storage_bytes, stream};
+
+  // Run reduction
+  cub::DeviceSegmentedReduce::Reduce(d_temp_storage.data(),
+                                     temp_storage_bytes,
+                                     d_in,
+                                     dev_result_mview.data<OutputType>(),
+                                     num_segments,
+                                     d_offset,
+                                     d_offset + 1,
+                                     binary_op,
+                                     identity,
+                                     stream.value());
+
+  return dev_result;
+}
+
+template <typename Op,
+          typename InputIterator,
+          typename OffsetIterator,
+          typename OutputType = typename thrust::iterator_value<InputIterator>::type,
+          typename std::enable_if_t<std::is_same_v<OutputType, string_view>>* = nullptr>
+std::unique_ptr<column> segmented_reduce(InputIterator,
+                                         OffsetIterator,
+                                         cudf::size_type,
+                                         op::simple_op<Op>,
+                                         rmm::cuda_stream_view,
+                                         rmm::mr::device_memory_resource*)
+{
+  CUDF_FAIL("Segment reduction for string type is unsupported.");
+}
+
 }  // namespace detail
 }  // namespace reduction
 }  // namespace cudf