Change streaming algorithms to use operator+= from using operator+ (NVIDIA#4428)

Author: oleksandr-pavlyk

cub/cub/device/dispatch/dispatch_advance_iterators.cuh

@@ -0,0 +1,109 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#pragma once
+
+#include <cub/config.cuh>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/std/type_traits>
+
+CUB_NAMESPACE_BEGIN
+
+namespace detail
+{
+template <typename T, typename U, typename = void>
+struct has_plus_operator : ::cuda::std::false_type
+{};
+
+template <typename T, typename U>
+struct has_plus_operator<T, U, ::cuda::std::void_t<decltype(::cuda::std::declval<T>() + ::cuda::std::declval<U>())>>
+    : ::cuda::std::true_type
+{};
+
+template <typename T, typename U>
+constexpr bool has_plus_operator_v = has_plus_operator<T, U>::value;
+
+// Helper function that advances a given iterator only if it supports being advanced by the given offset
+template <typename IteratorT, typename OffsetT>
+CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE IteratorT
+advance_iterators_if_supported(IteratorT iter, [[maybe_unused]] OffsetT offset)
+{
+  if constexpr (has_plus_operator_v<IteratorT, OffsetT>)
+  {
+    // If operator+ is valid, advance the iterator.
+    return iter + offset;
+  }
+  else
+  {
+    // Otherwise, return iter unmodified.
+    return iter;
+  }
+}
+
+template <typename T, typename U, typename = void>
+struct has_add_assign_operator : ::cuda::std::false_type
+{};
+
+template <typename T, typename U>
+struct has_add_assign_operator<T,
+                               U,
+                               ::cuda::std::void_t<decltype(::cuda::std::declval<T&>() += ::cuda::std::declval<U>())>>
+    : ::cuda::std::true_type
+{};
+
+template <typename T, typename U>
+constexpr bool has_add_assign_operator_v = has_add_assign_operator<T, U>::value;
+
+// Helper function that advances a given iterator only if it supports being advanced by the given offset
+template <typename IteratorT, typename OffsetT>
+CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE void
+advance_iterators_inplace_if_supported(IteratorT& iter, [[maybe_unused]] OffsetT offset)
+{
+  if constexpr (has_add_assign_operator_v<IteratorT, OffsetT>)
+  {
+    // If operator+ is valid, advance the iterator.
+    iter += offset;
+  }
+}
+
+// Helper function that checks whether all of the given iterators support the + operator with the given offset
+template <typename OffsetT, typename... Iterators>
+CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE bool
+all_iterators_support_plus_operator(OffsetT /*offset*/, Iterators... /*iters*/)
+{
+  if constexpr ((has_plus_operator_v<Iterators, OffsetT> && ...))
+  {
+    return true;
+  }
+  else
+  {
+    return false;
+  }
+}
+
+// Helper function that checks whether all of the given iterators support the + operator with the given offset
+template <typename OffsetT, typename... Iterators>
+CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE bool
+all_iterators_support_add_assign_operator(OffsetT /*offset*/, Iterators... /*iters*/)
+{
+  if constexpr ((has_add_assign_operator_v<Iterators, OffsetT> && ...))
+  {
+    return true;
+  }
+  else
+  {
+    return false;
+  }
+}
+
+} // namespace detail
+
+CUB_NAMESPACE_END

cub/cub/device/dispatch/dispatch_common.cuh

@@ -42,50 +42,4 @@ enum class SelectImpl
   Partition
 };
 
-namespace detail
-{
-template <typename T, typename U, typename = void>
-struct has_plus_operator : ::cuda::std::false_type
-{};
-
-template <typename T, typename U>
-struct has_plus_operator<T, U, ::cuda::std::void_t<decltype(::cuda::std::declval<T>() + ::cuda::std::declval<U>())>>
-    : ::cuda::std::true_type
-{};
-
-template <typename T, typename U>
-constexpr bool has_plus_operator_v = has_plus_operator<T, U>::value;
-
-// Helper function that advances a given iterator only if it supports being advanced by the given offset
-template <typename IteratorT, typename OffsetT>
-_CCCL_HOST_DEVICE IteratorT advance_iterators_if_supported(IteratorT iter, [[maybe_unused]] OffsetT offset)
-{
-  if constexpr (has_plus_operator_v<IteratorT, OffsetT>)
-  {
-    // If operator+ is valid, advance the iterator.
-    return iter + offset;
-  }
-  else
-  {
-    // Otherwise, return iter unmodified.
-    return iter;
-  }
-}
-
-// Helper function that checks whether all of the given iterators support the + operator with the given offset
-template <typename OffsetT, typename... Iterators>
-_CCCL_HOST_DEVICE bool all_iterators_support_plus_operator(OffsetT /*offset*/, Iterators... /*iters*/)
-{
-  if constexpr ((has_plus_operator_v<Iterators, OffsetT> && ...))
-  {
-    return true;
-  }
-  else
-  {
-    return false;
-  }
-}
-
-} // namespace detail
-
 CUB_NAMESPACE_END

cub/cub/device/dispatch/dispatch_radix_sort.cuh

@@ -44,6 +44,7 @@
 #  pragma system_header
 #endif // no system header
 
+#include <cub/device/dispatch/dispatch_advance_iterators.cuh>
 #include <cub/device/dispatch/kernels/radix_sort.cuh>
 #include <cub/device/dispatch/tuning/tuning_radix_sort.cuh>
 #include <cub/util_debug.cuh>
@@ -1405,11 +1406,14 @@ struct DispatchSegmentedRadixSort
     // If d_begin_offsets and d_end_offsets do not support operator+ then we can't have more than
     // max_num_segments_per_invocation segments per invocation
     if (num_invocations > 1
-        && !detail::all_iterators_support_plus_operator(::cuda::std::int64_t{}, d_begin_offsets, d_end_offsets))
+        && !detail::all_iterators_support_add_assign_operator(::cuda::std::int64_t{}, d_begin_offsets, d_end_offsets))
     {
       return cudaErrorInvalidValue;
     }
 
+    BeginOffsetIteratorT begin_offsets_current_it = d_begin_offsets;
+    EndOffsetIteratorT end_offsets_current_it     = d_end_offsets;
+
     // Iterate over chunks of segments
     for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++)
     {
@@ -1440,8 +1444,8 @@ struct DispatchSegmentedRadixSort
               d_keys_out,
               d_values_in,
               d_values_out,
-              detail::advance_iterators_if_supported(d_begin_offsets, current_segment_offset),
-              detail::advance_iterators_if_supported(d_end_offsets, current_segment_offset),
+              begin_offsets_current_it,
+              end_offsets_current_it,
               current_bit,
               pass_bits,
               decomposer);
@@ -1453,6 +1457,12 @@ struct DispatchSegmentedRadixSort
         return error;
       }
 
+      if (invocation_index + 1 < num_invocations)
+      {
+        detail::advance_iterators_inplace_if_supported(begin_offsets_current_it, num_current_segments);
+        detail::advance_iterators_inplace_if_supported(end_offsets_current_it, num_current_segments);
+      }
+
       // Sync the stream if specified to flush runtime errors
       error = CubDebug(detail::DebugSyncStream(stream));
       if (cudaSuccess != error)

cub/cub/device/dispatch/dispatch_reduce.cuh

@@ -46,7 +46,7 @@
 
 #include <cub/detail/launcher/cuda_runtime.cuh>
 #include <cub/detail/type_traits.cuh> // for cub::detail::invoke_result_t
-#include <cub/device/dispatch/dispatch_common.cuh>
+#include <cub/device/dispatch/dispatch_advance_iterators.cuh>
 #include <cub/device/dispatch/kernels/reduce.cuh>
 #include <cub/device/dispatch/kernels/segmented_reduce.cuh>
 #include <cub/device/dispatch/tuning/tuning_reduce.cuh>
@@ -823,7 +823,8 @@ struct DispatchSegmentedReduce
       // indirect_arg_t as the iterator type, which does not support the + operator.
       // TODO (elstehle): Remove this check once https://github.com/NVIDIA/cccl/issues/4148 is resolved.
       if (num_invocations > 1
-          && !detail::all_iterators_support_plus_operator(::cuda::std::int64_t{}, d_out, d_begin_offsets, d_end_offsets))
+          && !detail::all_iterators_support_add_assign_operator(
+            ::cuda::std::int64_t{}, d_out, d_begin_offsets, d_end_offsets))
       {
         return cudaErrorInvalidValue;
       }
@@ -848,13 +849,7 @@ struct DispatchSegmentedReduce
         // Invoke DeviceReduceKernel
         launcher_factory(
           static_cast<::cuda::std::uint32_t>(num_current_segments), policy.SegmentedReduce().BlockThreads(), 0, stream)
-          .doit(segmented_reduce_kernel,
-                d_in,
-                detail::advance_iterators_if_supported(d_out, current_seg_offset),
-                detail::advance_iterators_if_supported(d_begin_offsets, current_seg_offset),
-                detail::advance_iterators_if_supported(d_end_offsets, current_seg_offset),
-                reduction_op,
-                init);
+          .doit(segmented_reduce_kernel, d_in, d_out, d_begin_offsets, d_end_offsets, reduction_op, init);
 
         // Check for failure to launch
         error = CubDebug(cudaPeekAtLastError());
@@ -863,6 +858,13 @@ struct DispatchSegmentedReduce
           break;
         }
 
+        if (invocation_index + 1 < num_invocations)
+        {
+          detail::advance_iterators_inplace_if_supported(d_out, num_current_segments);
+          detail::advance_iterators_inplace_if_supported(d_begin_offsets, num_current_segments);
+          detail::advance_iterators_inplace_if_supported(d_end_offsets, num_current_segments);
+        }
+
         // Sync the stream if specified to flush runtime errors
         error = CubDebug(detail::DebugSyncStream(stream));
         if (cudaSuccess != error)