From f024f46a6b3dabacdebedf98bea4f839463581c7 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 11 Aug 2025 22:36:00 +0000 Subject: [PATCH 001/100] Move segmented sort kernels to separate header --- .../dispatch/dispatch_segmented_sort.cuh | 653 +---------------- .../dispatch/kernels/segmented_sort.cuh | 672 ++++++++++++++++++ cub/test/catch2_test_nvrtc.cu | 1 + 3 files changed, 674 insertions(+), 652 deletions(-) create mode 100644 cub/cub/device/dispatch/kernels/segmented_sort.cuh diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 69c146c850b..8f9f12a4491 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -37,24 +37,16 @@ # pragma system_header #endif // no system header -#include -#include -#include -#include -#include -#include #include #include #include +#include #include -#include #include #include #include #include -#include -#include #include #include #include @@ -70,649 +62,6 @@ CUB_NAMESPACE_BEGIN namespace detail::segmented_sort { -// Type used to index within segments within a single invocation -using local_segment_index_t = ::cuda::std::uint32_t; -// Type used for total number of segments and to index within segments globally -using global_segment_offset_t = ::cuda::std::int64_t; - -/** - * @brief Fallback kernel, in case there's not enough segments to - * take advantage of partitioning. - * - * In this case, the sorting method is still selected based on the segment size. - * If a single warp can sort the segment, the algorithm will use the sub-warp - * merge sort. Otherwise, the algorithm will use the in-shared-memory version of - * block radix sort. If data don't fit into shared memory, the algorithm will - * use in-global-memory radix sort. - * - * @param[in] d_keys_in_orig - * Input keys buffer - * - * @param[out] d_keys_out_orig - * Output keys buffer - * - * @param[in,out] d_keys_double_buffer - * Double keys buffer - * - * @param[in] d_values_in_orig - * Input values buffer - * - * @param[out] d_values_out_orig - * Output values buffer - * - * @param[in,out] d_values_double_buffer - * Double values buffer - * - * @param[in] d_begin_offsets - * Random-access input iterator to the sequence of beginning offsets of length - * `num_segments`, such that `d_begin_offsets[i]` is the first element of the - * i-th data segment in `d_keys_*` and `d_values_*` - * - * @param[in] d_end_offsets - * Random-access input iterator to the sequence of ending offsets of length - * `num_segments`, such that `d_end_offsets[i]-1` is the last element of the - * i-th data segment in `d_keys_*` and `d_values_*`. - * If `d_end_offsets[i]-1 <= d_begin_offsets[i]`, the i-th segment is - * considered empty. - */ -template -__launch_bounds__(ChainedPolicyT::ActivePolicy::LargeSegmentPolicy::BLOCK_THREADS) - CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortFallbackKernel( - const KeyT* d_keys_in_orig, - KeyT* d_keys_out_orig, - device_double_buffer d_keys_double_buffer, - const ValueT* d_values_in_orig, - ValueT* d_values_out_orig, - device_double_buffer d_values_double_buffer, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets) -{ - using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; - using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - using MediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT::MediumPolicyT; - - const auto segment_id = static_cast(blockIdx.x); - OffsetT segment_begin = d_begin_offsets[segment_id]; - OffsetT segment_end = d_end_offsets[segment_id]; - OffsetT num_items = segment_end - segment_begin; - - if (num_items <= 0) - { - return; - } - - using AgentSegmentedRadixSortT = - radix_sort::AgentSegmentedRadixSort; - - using WarpReduceT = cub::WarpReduce; - - using AgentWarpMergeSortT = - sub_warp_merge_sort::AgentSubWarpSort; - - __shared__ union - { - typename AgentSegmentedRadixSortT::TempStorage block_sort; - typename WarpReduceT::TempStorage warp_reduce; - typename AgentWarpMergeSortT::TempStorage medium_warp_sort; - } temp_storage; - - constexpr bool keys_only = ::cuda::std::is_same_v; - AgentSegmentedRadixSortT agent(num_items, temp_storage.block_sort); - - constexpr int begin_bit = 0; - constexpr int end_bit = sizeof(KeyT) * 8; - - constexpr int cacheable_tile_size = LargeSegmentPolicyT::BLOCK_THREADS * LargeSegmentPolicyT::ITEMS_PER_THREAD; - - d_keys_in_orig += segment_begin; - d_keys_out_orig += segment_begin; - - if (!keys_only) - { - d_values_in_orig += segment_begin; - d_values_out_orig += segment_begin; - } - - if (num_items <= MediumPolicyT::ITEMS_PER_TILE) - { - // Sort by a single warp - if (threadIdx.x < MediumPolicyT::WARP_THREADS) - { - AgentWarpMergeSortT(temp_storage.medium_warp_sort) - .ProcessSegment(num_items, d_keys_in_orig, d_keys_out_orig, d_values_in_orig, d_values_out_orig); - } - } - else if (num_items < cacheable_tile_size) - { - // Sort by a CTA if data fits into shared memory - agent.ProcessSinglePass(begin_bit, end_bit, d_keys_in_orig, d_values_in_orig, d_keys_out_orig, d_values_out_orig); - } - else - { - // Sort by a CTA with multiple reads from global memory - int current_bit = begin_bit; - int pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); - - d_keys_double_buffer = device_double_buffer( - d_keys_double_buffer.current() + segment_begin, d_keys_double_buffer.alternate() + segment_begin); - - if (!keys_only) - { - d_values_double_buffer = device_double_buffer( - d_values_double_buffer.current() + segment_begin, d_values_double_buffer.alternate() + segment_begin); - } - - agent.ProcessIterative( - current_bit, - pass_bits, - d_keys_in_orig, - d_values_in_orig, - d_keys_double_buffer.current(), - d_values_double_buffer.current()); - current_bit += pass_bits; - - _CCCL_PRAGMA_NOUNROLL() - while (current_bit < end_bit) - { - pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); - - __syncthreads(); - agent.ProcessIterative( - current_bit, - pass_bits, - d_keys_double_buffer.current(), - d_values_double_buffer.current(), - d_keys_double_buffer.alternate(), - d_values_double_buffer.alternate()); - - d_keys_double_buffer.swap(); - d_values_double_buffer.swap(); - current_bit += pass_bits; - } - } -} - -/** - * @brief Single kernel for moderate size (less than a few thousand items) - * segments. - * - * This kernel allocates a sub-warp per segment. Therefore, this kernel assigns - * a single thread block to multiple segments. Segments fall into two - * categories. An architectural warp usually sorts segments in the medium-size - * category, while a few threads sort segments in the small-size category. Since - * segments are partitioned, we know the last thread block index assigned to - * sort medium-size segments. A particular thread block can check this number to - * find out which category it was assigned to sort. In both cases, the - * merge sort is used. - * - * @param[in] small_segments - * Number of segments that can be sorted by a warp part - * - * @param[in] medium_segments - * Number of segments that can be sorted by a warp - * - * @param[in] medium_blocks - * Number of CTAs assigned to process medium segments - * - * @param[in] d_small_segments_indices - * Small segments mapping of length @p small_segments, such that - * `d_small_segments_indices[i]` is the input segment index - * - * @param[in] d_medium_segments_indices - * Medium segments mapping of length @p medium_segments, such that - * `d_medium_segments_indices[i]` is the input segment index - * - * @param[in] d_keys_in_orig - * Input keys buffer - * - * @param[out] d_keys_out_orig - * Output keys buffer - * - * @param[in] d_values_in_orig - * Input values buffer - * - * @param[out] d_values_out_orig - * Output values buffer - * - * @param[in] d_begin_offsets - * Random-access input iterator to the sequence of beginning offsets of length - * `num_segments`, such that `d_begin_offsets[i]` is the first element of the - * ith data segment in `d_keys_*` and `d_values_*` - * - * @param[in] d_end_offsets - * Random-access input iterator to the sequence of ending offsets of length - * `num_segments`, such that `d_end_offsets[i]-1` is the last element of the - * ith data segment in `d_keys_*` and `d_values_*`. If - * `d_end_offsets[i]-1 <= d_begin_offsets[i]`, the ith is - * considered empty. - */ -template -__launch_bounds__(ChainedPolicyT::ActivePolicy::SmallAndMediumSegmentedSortPolicyT::BLOCK_THREADS) - CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortKernelSmall( - local_segment_index_t small_segments, - local_segment_index_t medium_segments, - local_segment_index_t medium_blocks, - const local_segment_index_t* d_small_segments_indices, - const local_segment_index_t* d_medium_segments_indices, - const KeyT* d_keys_in, - KeyT* d_keys_out, - const ValueT* d_values_in, - ValueT* d_values_out, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets) -{ - using local_segment_index_t = local_segment_index_t; - - const local_segment_index_t tid = threadIdx.x; - const local_segment_index_t bid = blockIdx.x; - - using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; - using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - using MediumPolicyT = typename SmallAndMediumPolicyT::MediumPolicyT; - using SmallPolicyT = typename SmallAndMediumPolicyT::SmallPolicyT; - - constexpr auto threads_per_medium_segment = static_cast(MediumPolicyT::WARP_THREADS); - constexpr auto threads_per_small_segment = static_cast(SmallPolicyT::WARP_THREADS); - - using MediumAgentWarpMergeSortT = - sub_warp_merge_sort::AgentSubWarpSort; - - using SmallAgentWarpMergeSortT = - sub_warp_merge_sort::AgentSubWarpSort; - - constexpr auto segments_per_medium_block = - static_cast(SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); - - constexpr auto segments_per_small_block = - static_cast(SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); - - __shared__ union - { - typename MediumAgentWarpMergeSortT::TempStorage medium_storage[segments_per_medium_block]; - - typename SmallAgentWarpMergeSortT::TempStorage small_storage[segments_per_small_block]; - } temp_storage; - - if (bid < medium_blocks) - { - const local_segment_index_t sid_within_block = tid / threads_per_medium_segment; - const local_segment_index_t medium_segment_id = bid * segments_per_medium_block + sid_within_block; - - if (medium_segment_id < medium_segments) - { - const local_segment_index_t global_segment_id = d_medium_segments_indices[medium_segment_id]; - - const OffsetT segment_begin = d_begin_offsets[global_segment_id]; - const OffsetT segment_end = d_end_offsets[global_segment_id]; - const OffsetT num_items = segment_end - segment_begin; - - MediumAgentWarpMergeSortT(temp_storage.medium_storage[sid_within_block]) - .ProcessSegment(num_items, - d_keys_in + segment_begin, - d_keys_out + segment_begin, - d_values_in + segment_begin, - d_values_out + segment_begin); - } - } - else - { - const local_segment_index_t sid_within_block = tid / threads_per_small_segment; - const local_segment_index_t small_segment_id = (bid - medium_blocks) * segments_per_small_block + sid_within_block; - - if (small_segment_id < small_segments) - { - const local_segment_index_t global_segment_id = d_small_segments_indices[small_segment_id]; - - const OffsetT segment_begin = d_begin_offsets[global_segment_id]; - const OffsetT segment_end = d_end_offsets[global_segment_id]; - const OffsetT num_items = segment_end - segment_begin; - - SmallAgentWarpMergeSortT(temp_storage.small_storage[sid_within_block]) - .ProcessSegment(num_items, - d_keys_in + segment_begin, - d_keys_out + segment_begin, - d_values_in + segment_begin, - d_values_out + segment_begin); - } - } -} - -/** - * @brief Single kernel for large size (more than a few thousand items) segments. - * - * @param[in] d_keys_in_orig - * Input keys buffer - * - * @param[out] d_keys_out_orig - * Output keys buffer - * - * @param[in] d_values_in_orig - * Input values buffer - * - * @param[out] d_values_out_orig - * Output values buffer - * - * @param[in] d_begin_offsets - * Random-access input iterator to the sequence of beginning offsets of length - * `num_segments`, such that `d_begin_offsets[i]` is the first element of the - * ith data segment in `d_keys_*` and `d_values_*` - * - * @param[in] d_end_offsets - * Random-access input iterator to the sequence of ending offsets of length - * `num_segments`, such that `d_end_offsets[i]-1` is the last element of the - * ith data segment in `d_keys_*` and `d_values_*`. If - * `d_end_offsets[i]-1 <= d_begin_offsets[i]`, the ith is - * considered empty. - */ -template -__launch_bounds__(ChainedPolicyT::ActivePolicy::LargeSegmentPolicy::BLOCK_THREADS) - CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortKernelLarge( - const local_segment_index_t* d_segments_indices, - const KeyT* d_keys_in_orig, - KeyT* d_keys_out_orig, - device_double_buffer d_keys_double_buffer, - const ValueT* d_values_in_orig, - ValueT* d_values_out_orig, - device_double_buffer d_values_double_buffer, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets) -{ - using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; - using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - using local_segment_index_t = local_segment_index_t; - - constexpr int small_tile_size = LargeSegmentPolicyT::BLOCK_THREADS * LargeSegmentPolicyT::ITEMS_PER_THREAD; - - using AgentSegmentedRadixSortT = - radix_sort::AgentSegmentedRadixSort; - - __shared__ typename AgentSegmentedRadixSortT::TempStorage storage; - - const local_segment_index_t bid = blockIdx.x; - - constexpr int begin_bit = 0; - constexpr int end_bit = sizeof(KeyT) * 8; - - const local_segment_index_t global_segment_id = d_segments_indices[bid]; - const OffsetT segment_begin = d_begin_offsets[global_segment_id]; - const OffsetT segment_end = d_end_offsets[global_segment_id]; - const OffsetT num_items = segment_end - segment_begin; - - constexpr bool keys_only = ::cuda::std::is_same_v; - AgentSegmentedRadixSortT agent(num_items, storage); - - d_keys_in_orig += segment_begin; - d_keys_out_orig += segment_begin; - - if (!keys_only) - { - d_values_in_orig += segment_begin; - d_values_out_orig += segment_begin; - } - - if (num_items < small_tile_size) - { - // Sort in shared memory if the segment fits into it - agent.ProcessSinglePass(begin_bit, end_bit, d_keys_in_orig, d_values_in_orig, d_keys_out_orig, d_values_out_orig); - } - else - { - // Sort reading global memory multiple times - int current_bit = begin_bit; - int pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); - - d_keys_double_buffer = device_double_buffer( - d_keys_double_buffer.current() + segment_begin, d_keys_double_buffer.alternate() + segment_begin); - - if (!keys_only) - { - d_values_double_buffer = device_double_buffer( - d_values_double_buffer.current() + segment_begin, d_values_double_buffer.alternate() + segment_begin); - } - - agent.ProcessIterative( - current_bit, - pass_bits, - d_keys_in_orig, - d_values_in_orig, - d_keys_double_buffer.current(), - d_values_double_buffer.current()); - current_bit += pass_bits; - - _CCCL_PRAGMA_NOUNROLL() - while (current_bit < end_bit) - { - pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); - - __syncthreads(); - agent.ProcessIterative( - current_bit, - pass_bits, - d_keys_double_buffer.current(), - d_values_double_buffer.current(), - d_keys_double_buffer.alternate(), - d_values_double_buffer.alternate()); - - d_keys_double_buffer.swap(); - d_values_double_buffer.swap(); - current_bit += pass_bits; - } - } -} - -/* - * Continuation is called after the partitioning stage. It launches kernels - * to sort large and small segments using the partitioning results. Separation - * of this stage is required to eliminate device-side synchronization in - * the CDP mode. - */ -template -CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortContinuation( - LargeKernelT large_kernel, - SmallKernelT small_kernel, - int num_segments, - KeyT* d_current_keys, - KeyT* d_final_keys, - device_double_buffer d_keys_double_buffer, - ValueT* d_current_values, - ValueT* d_final_values, - device_double_buffer d_values_double_buffer, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets, - local_segment_index_t* group_sizes, - local_segment_index_t* large_and_medium_segments_indices, - local_segment_index_t* small_segments_indices, - cudaStream_t stream) -{ - using local_segment_index_t = local_segment_index_t; - - cudaError error = cudaSuccess; - - const local_segment_index_t large_segments = group_sizes[0]; - - if (large_segments > 0) - { - // One CTA per segment - const local_segment_index_t blocks_in_grid = large_segments; - -#ifdef CUB_DEBUG_LOG - _CubLog("Invoking " - "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", - static_cast(blocks_in_grid), - LargeSegmentPolicyT::BLOCK_THREADS, - (long long) stream); -#endif // CUB_DEBUG_LOG - - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(blocks_in_grid, LargeSegmentPolicyT::BLOCK_THREADS, 0, stream) - .doit(large_kernel, - large_and_medium_segments_indices, - d_current_keys, - d_final_keys, - d_keys_double_buffer, - d_current_values, - d_final_values, - d_values_double_buffer, - d_begin_offsets, - d_end_offsets); - - // Check for failure to launch - error = CubDebug(cudaPeekAtLastError()); - if (cudaSuccess != error) - { - return error; - } - - // Sync the stream if specified to flush runtime errors - error = CubDebug(DebugSyncStream(stream)); - if (cudaSuccess != error) - { - return error; - } - } - - const local_segment_index_t small_segments = group_sizes[1]; - const local_segment_index_t medium_segments = - static_cast(num_segments) - (large_segments + small_segments); - - const local_segment_index_t small_blocks = - ::cuda::ceil_div(small_segments, SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); - - const local_segment_index_t medium_blocks = - ::cuda::ceil_div(medium_segments, SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); - - const local_segment_index_t small_and_medium_blocks_in_grid = small_blocks + medium_blocks; - - if (small_and_medium_blocks_in_grid) - { -#ifdef CUB_DEBUG_LOG - _CubLog("Invoking " - "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", - static_cast(small_and_medium_blocks_in_grid), - SmallAndMediumPolicyT::BLOCK_THREADS, - (long long) stream); -#endif // CUB_DEBUG_LOG - - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( - small_and_medium_blocks_in_grid, SmallAndMediumPolicyT::BLOCK_THREADS, 0, stream) - .doit(small_kernel, - small_segments, - medium_segments, - medium_blocks, - small_segments_indices, - large_and_medium_segments_indices + num_segments - medium_segments, - d_current_keys, - d_final_keys, - d_current_values, - d_final_values, - d_begin_offsets, - d_end_offsets); - - // Check for failure to launch - error = CubDebug(cudaPeekAtLastError()); - if (cudaSuccess != error) - { - return error; - } - - // Sync the stream if specified to flush runtime errors - error = CubDebug(DebugSyncStream(stream)); - if (cudaSuccess != error) - { - return error; - } - } - - return error; -} - -#ifdef CUB_RDC_ENABLED -/* - * Continuation kernel is used only in the CDP mode. It's used to - * launch DeviceSegmentedSortContinuation as a separate kernel. - */ -template -__launch_bounds__(1) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContinuationKernel( - LargeKernelT large_kernel, - SmallKernelT small_kernel, - local_segment_index_t num_segments, - KeyT* d_current_keys, - KeyT* d_final_keys, - device_double_buffer d_keys_double_buffer, - ValueT* d_current_values, - ValueT* d_final_values, - device_double_buffer d_values_double_buffer, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets, - local_segment_index_t* group_sizes, - local_segment_index_t* large_and_medium_segments_indices, - local_segment_index_t* small_segments_indices) -{ - using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; - using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - - // In case of CDP: - // 1. each CTA has a different main stream - // 2. all streams are non-blocking - // 3. child grid always completes before the parent grid - // 4. streams can be used only from the CTA in which they were created - // 5. streams created on the host cannot be used on the device - // - // Due to (4, 5), we can't pass the user-provided stream in the continuation. - // Due to (1, 2, 3) it's safe to pass the main stream. - cudaError_t error = - detail::segmented_sort::DeviceSegmentedSortContinuation( - large_kernel, - small_kernel, - num_segments, - d_current_keys, - d_final_keys, - d_keys_double_buffer, - d_current_values, - d_final_values, - d_values_double_buffer, - d_begin_offsets, - d_end_offsets, - group_sizes, - large_and_medium_segments_indices, - small_segments_indices, - 0); // always launching on the main stream (see motivation above) - - error = CubDebug(error); -} -#endif // CUB_RDC_ENABLED - } // namespace detail::segmented_sort template + +#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 +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN +namespace detail::segmented_sort +{ + +// Type used to index within segments within a single invocation +using local_segment_index_t = ::cuda::std::uint32_t; +// Type used for total number of segments and to index within segments globally +using global_segment_offset_t = ::cuda::std::int64_t; + +/** + * @brief Fallback kernel, in case there's not enough segments to + * take advantage of partitioning. + * + * In this case, the sorting method is still selected based on the segment size. + * If a single warp can sort the segment, the algorithm will use the sub-warp + * merge sort. Otherwise, the algorithm will use the in-shared-memory version of + * block radix sort. If data don't fit into shared memory, the algorithm will + * use in-global-memory radix sort. + * + * @param[in] d_keys_in_orig + * Input keys buffer + * + * @param[out] d_keys_out_orig + * Output keys buffer + * + * @param[in,out] d_keys_double_buffer + * Double keys buffer + * + * @param[in] d_values_in_orig + * Input values buffer + * + * @param[out] d_values_out_orig + * Output values buffer + * + * @param[in,out] d_values_double_buffer + * Double values buffer + * + * @param[in] d_begin_offsets + * Random-access input iterator to the sequence of beginning offsets of length + * `num_segments`, such that `d_begin_offsets[i]` is the first element of the + * i-th data segment in `d_keys_*` and `d_values_*` + * + * @param[in] d_end_offsets + * Random-access input iterator to the sequence of ending offsets of length + * `num_segments`, such that `d_end_offsets[i]-1` is the last element of the + * i-th data segment in `d_keys_*` and `d_values_*`. + * If `d_end_offsets[i]-1 <= d_begin_offsets[i]`, the i-th segment is + * considered empty. + */ +template +__launch_bounds__(ChainedPolicyT::ActivePolicy::LargeSegmentPolicy::BLOCK_THREADS) + CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortFallbackKernel( + const KeyT* d_keys_in_orig, + KeyT* d_keys_out_orig, + device_double_buffer d_keys_double_buffer, + const ValueT* d_values_in_orig, + ValueT* d_values_out_orig, + device_double_buffer d_values_double_buffer, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets) +{ + using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; + using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; + using MediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT::MediumPolicyT; + + const auto segment_id = static_cast(blockIdx.x); + OffsetT segment_begin = d_begin_offsets[segment_id]; + OffsetT segment_end = d_end_offsets[segment_id]; + OffsetT num_items = segment_end - segment_begin; + + if (num_items <= 0) + { + return; + } + + using AgentSegmentedRadixSortT = + radix_sort::AgentSegmentedRadixSort; + + using WarpReduceT = cub::WarpReduce; + + using AgentWarpMergeSortT = + sub_warp_merge_sort::AgentSubWarpSort; + + __shared__ union + { + typename AgentSegmentedRadixSortT::TempStorage block_sort; + typename WarpReduceT::TempStorage warp_reduce; + typename AgentWarpMergeSortT::TempStorage medium_warp_sort; + } temp_storage; + + constexpr bool keys_only = ::cuda::std::is_same_v; + AgentSegmentedRadixSortT agent(num_items, temp_storage.block_sort); + + constexpr int begin_bit = 0; + constexpr int end_bit = sizeof(KeyT) * 8; + + constexpr int cacheable_tile_size = LargeSegmentPolicyT::BLOCK_THREADS * LargeSegmentPolicyT::ITEMS_PER_THREAD; + + d_keys_in_orig += segment_begin; + d_keys_out_orig += segment_begin; + + if (!keys_only) + { + d_values_in_orig += segment_begin; + d_values_out_orig += segment_begin; + } + + if (num_items <= MediumPolicyT::ITEMS_PER_TILE) + { + // Sort by a single warp + if (threadIdx.x < MediumPolicyT::WARP_THREADS) + { + AgentWarpMergeSortT(temp_storage.medium_warp_sort) + .ProcessSegment(num_items, d_keys_in_orig, d_keys_out_orig, d_values_in_orig, d_values_out_orig); + } + } + else if (num_items < cacheable_tile_size) + { + // Sort by a CTA if data fits into shared memory + agent.ProcessSinglePass(begin_bit, end_bit, d_keys_in_orig, d_values_in_orig, d_keys_out_orig, d_values_out_orig); + } + else + { + // Sort by a CTA with multiple reads from global memory + int current_bit = begin_bit; + int pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); + + d_keys_double_buffer = device_double_buffer( + d_keys_double_buffer.current() + segment_begin, d_keys_double_buffer.alternate() + segment_begin); + + if (!keys_only) + { + d_values_double_buffer = device_double_buffer( + d_values_double_buffer.current() + segment_begin, d_values_double_buffer.alternate() + segment_begin); + } + + agent.ProcessIterative( + current_bit, + pass_bits, + d_keys_in_orig, + d_values_in_orig, + d_keys_double_buffer.current(), + d_values_double_buffer.current()); + current_bit += pass_bits; + + _CCCL_PRAGMA_NOUNROLL() + while (current_bit < end_bit) + { + pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); + + __syncthreads(); + agent.ProcessIterative( + current_bit, + pass_bits, + d_keys_double_buffer.current(), + d_values_double_buffer.current(), + d_keys_double_buffer.alternate(), + d_values_double_buffer.alternate()); + + d_keys_double_buffer.swap(); + d_values_double_buffer.swap(); + current_bit += pass_bits; + } + } +} + +/** + * @brief Single kernel for moderate size (less than a few thousand items) + * segments. + * + * This kernel allocates a sub-warp per segment. Therefore, this kernel assigns + * a single thread block to multiple segments. Segments fall into two + * categories. An architectural warp usually sorts segments in the medium-size + * category, while a few threads sort segments in the small-size category. Since + * segments are partitioned, we know the last thread block index assigned to + * sort medium-size segments. A particular thread block can check this number to + * find out which category it was assigned to sort. In both cases, the + * merge sort is used. + * + * @param[in] small_segments + * Number of segments that can be sorted by a warp part + * + * @param[in] medium_segments + * Number of segments that can be sorted by a warp + * + * @param[in] medium_blocks + * Number of CTAs assigned to process medium segments + * + * @param[in] d_small_segments_indices + * Small segments mapping of length @p small_segments, such that + * `d_small_segments_indices[i]` is the input segment index + * + * @param[in] d_medium_segments_indices + * Medium segments mapping of length @p medium_segments, such that + * `d_medium_segments_indices[i]` is the input segment index + * + * @param[in] d_keys_in_orig + * Input keys buffer + * + * @param[out] d_keys_out_orig + * Output keys buffer + * + * @param[in] d_values_in_orig + * Input values buffer + * + * @param[out] d_values_out_orig + * Output values buffer + * + * @param[in] d_begin_offsets + * Random-access input iterator to the sequence of beginning offsets of length + * `num_segments`, such that `d_begin_offsets[i]` is the first element of the + * ith data segment in `d_keys_*` and `d_values_*` + * + * @param[in] d_end_offsets + * Random-access input iterator to the sequence of ending offsets of length + * `num_segments`, such that `d_end_offsets[i]-1` is the last element of the + * ith data segment in `d_keys_*` and `d_values_*`. If + * `d_end_offsets[i]-1 <= d_begin_offsets[i]`, the ith is + * considered empty. + */ +template +__launch_bounds__(ChainedPolicyT::ActivePolicy::SmallAndMediumSegmentedSortPolicyT::BLOCK_THREADS) + CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortKernelSmall( + local_segment_index_t small_segments, + local_segment_index_t medium_segments, + local_segment_index_t medium_blocks, + const local_segment_index_t* d_small_segments_indices, + const local_segment_index_t* d_medium_segments_indices, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets) +{ + using local_segment_index_t = local_segment_index_t; + + const local_segment_index_t tid = threadIdx.x; + const local_segment_index_t bid = blockIdx.x; + + using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; + using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; + using MediumPolicyT = typename SmallAndMediumPolicyT::MediumPolicyT; + using SmallPolicyT = typename SmallAndMediumPolicyT::SmallPolicyT; + + constexpr auto threads_per_medium_segment = static_cast(MediumPolicyT::WARP_THREADS); + constexpr auto threads_per_small_segment = static_cast(SmallPolicyT::WARP_THREADS); + + using MediumAgentWarpMergeSortT = + sub_warp_merge_sort::AgentSubWarpSort; + + using SmallAgentWarpMergeSortT = + sub_warp_merge_sort::AgentSubWarpSort; + + constexpr auto segments_per_medium_block = + static_cast(SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); + + constexpr auto segments_per_small_block = + static_cast(SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); + + __shared__ union + { + typename MediumAgentWarpMergeSortT::TempStorage medium_storage[segments_per_medium_block]; + + typename SmallAgentWarpMergeSortT::TempStorage small_storage[segments_per_small_block]; + } temp_storage; + + if (bid < medium_blocks) + { + const local_segment_index_t sid_within_block = tid / threads_per_medium_segment; + const local_segment_index_t medium_segment_id = bid * segments_per_medium_block + sid_within_block; + + if (medium_segment_id < medium_segments) + { + const local_segment_index_t global_segment_id = d_medium_segments_indices[medium_segment_id]; + + const OffsetT segment_begin = d_begin_offsets[global_segment_id]; + const OffsetT segment_end = d_end_offsets[global_segment_id]; + const OffsetT num_items = segment_end - segment_begin; + + MediumAgentWarpMergeSortT(temp_storage.medium_storage[sid_within_block]) + .ProcessSegment(num_items, + d_keys_in + segment_begin, + d_keys_out + segment_begin, + d_values_in + segment_begin, + d_values_out + segment_begin); + } + } + else + { + const local_segment_index_t sid_within_block = tid / threads_per_small_segment; + const local_segment_index_t small_segment_id = (bid - medium_blocks) * segments_per_small_block + sid_within_block; + + if (small_segment_id < small_segments) + { + const local_segment_index_t global_segment_id = d_small_segments_indices[small_segment_id]; + + const OffsetT segment_begin = d_begin_offsets[global_segment_id]; + const OffsetT segment_end = d_end_offsets[global_segment_id]; + const OffsetT num_items = segment_end - segment_begin; + + SmallAgentWarpMergeSortT(temp_storage.small_storage[sid_within_block]) + .ProcessSegment(num_items, + d_keys_in + segment_begin, + d_keys_out + segment_begin, + d_values_in + segment_begin, + d_values_out + segment_begin); + } + } +} + +/** + * @brief Single kernel for large size (more than a few thousand items) segments. + * + * @param[in] d_keys_in_orig + * Input keys buffer + * + * @param[out] d_keys_out_orig + * Output keys buffer + * + * @param[in] d_values_in_orig + * Input values buffer + * + * @param[out] d_values_out_orig + * Output values buffer + * + * @param[in] d_begin_offsets + * Random-access input iterator to the sequence of beginning offsets of length + * `num_segments`, such that `d_begin_offsets[i]` is the first element of the + * ith data segment in `d_keys_*` and `d_values_*` + * + * @param[in] d_end_offsets + * Random-access input iterator to the sequence of ending offsets of length + * `num_segments`, such that `d_end_offsets[i]-1` is the last element of the + * ith data segment in `d_keys_*` and `d_values_*`. If + * `d_end_offsets[i]-1 <= d_begin_offsets[i]`, the ith is + * considered empty. + */ +template +__launch_bounds__(ChainedPolicyT::ActivePolicy::LargeSegmentPolicy::BLOCK_THREADS) + CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortKernelLarge( + const local_segment_index_t* d_segments_indices, + const KeyT* d_keys_in_orig, + KeyT* d_keys_out_orig, + device_double_buffer d_keys_double_buffer, + const ValueT* d_values_in_orig, + ValueT* d_values_out_orig, + device_double_buffer d_values_double_buffer, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets) +{ + using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; + using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; + using local_segment_index_t = local_segment_index_t; + + constexpr int small_tile_size = LargeSegmentPolicyT::BLOCK_THREADS * LargeSegmentPolicyT::ITEMS_PER_THREAD; + + using AgentSegmentedRadixSortT = + radix_sort::AgentSegmentedRadixSort; + + __shared__ typename AgentSegmentedRadixSortT::TempStorage storage; + + const local_segment_index_t bid = blockIdx.x; + + constexpr int begin_bit = 0; + constexpr int end_bit = sizeof(KeyT) * 8; + + const local_segment_index_t global_segment_id = d_segments_indices[bid]; + const OffsetT segment_begin = d_begin_offsets[global_segment_id]; + const OffsetT segment_end = d_end_offsets[global_segment_id]; + const OffsetT num_items = segment_end - segment_begin; + + constexpr bool keys_only = ::cuda::std::is_same_v; + AgentSegmentedRadixSortT agent(num_items, storage); + + d_keys_in_orig += segment_begin; + d_keys_out_orig += segment_begin; + + if (!keys_only) + { + d_values_in_orig += segment_begin; + d_values_out_orig += segment_begin; + } + + if (num_items < small_tile_size) + { + // Sort in shared memory if the segment fits into it + agent.ProcessSinglePass(begin_bit, end_bit, d_keys_in_orig, d_values_in_orig, d_keys_out_orig, d_values_out_orig); + } + else + { + // Sort reading global memory multiple times + int current_bit = begin_bit; + int pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); + + d_keys_double_buffer = device_double_buffer( + d_keys_double_buffer.current() + segment_begin, d_keys_double_buffer.alternate() + segment_begin); + + if (!keys_only) + { + d_values_double_buffer = device_double_buffer( + d_values_double_buffer.current() + segment_begin, d_values_double_buffer.alternate() + segment_begin); + } + + agent.ProcessIterative( + current_bit, + pass_bits, + d_keys_in_orig, + d_values_in_orig, + d_keys_double_buffer.current(), + d_values_double_buffer.current()); + current_bit += pass_bits; + + _CCCL_PRAGMA_NOUNROLL() + while (current_bit < end_bit) + { + pass_bits = (::cuda::std::min) (int{LargeSegmentPolicyT::RADIX_BITS}, (end_bit - current_bit)); + + __syncthreads(); + agent.ProcessIterative( + current_bit, + pass_bits, + d_keys_double_buffer.current(), + d_values_double_buffer.current(), + d_keys_double_buffer.alternate(), + d_values_double_buffer.alternate()); + + d_keys_double_buffer.swap(); + d_values_double_buffer.swap(); + current_bit += pass_bits; + } + } +} +/* + * Continuation is called after the partitioning stage. It launches kernels + * to sort large and small segments using the partitioning results. Separation + * of this stage is required to eliminate device-side synchronization in + * the CDP mode. + */ +template +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortContinuation( + LargeKernelT large_kernel, + SmallKernelT small_kernel, + int num_segments, + KeyT* d_current_keys, + KeyT* d_final_keys, + device_double_buffer d_keys_double_buffer, + ValueT* d_current_values, + ValueT* d_final_values, + device_double_buffer d_values_double_buffer, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + local_segment_index_t* group_sizes, + local_segment_index_t* large_and_medium_segments_indices, + local_segment_index_t* small_segments_indices, + cudaStream_t stream) +{ + using local_segment_index_t = local_segment_index_t; + + cudaError error = cudaSuccess; + + const local_segment_index_t large_segments = group_sizes[0]; + + if (large_segments > 0) + { + // One CTA per segment + const local_segment_index_t blocks_in_grid = large_segments; + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking " + "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", + static_cast(blocks_in_grid), + LargeSegmentPolicyT::BLOCK_THREADS, + (long long) stream); +#endif // CUB_DEBUG_LOG + + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(blocks_in_grid, LargeSegmentPolicyT::BLOCK_THREADS, 0, stream) + .doit(large_kernel, + large_and_medium_segments_indices, + d_current_keys, + d_final_keys, + d_keys_double_buffer, + d_current_values, + d_final_values, + d_values_double_buffer, + d_begin_offsets, + d_end_offsets); + + // Check for failure to launch + error = CubDebug(cudaPeekAtLastError()); + if (cudaSuccess != error) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + } + + const local_segment_index_t small_segments = group_sizes[1]; + const local_segment_index_t medium_segments = + static_cast(num_segments) - (large_segments + small_segments); + + const local_segment_index_t small_blocks = + ::cuda::ceil_div(small_segments, SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); + + const local_segment_index_t medium_blocks = + ::cuda::ceil_div(medium_segments, SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); + + const local_segment_index_t small_and_medium_blocks_in_grid = small_blocks + medium_blocks; + + if (small_and_medium_blocks_in_grid) + { +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking " + "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", + static_cast(small_and_medium_blocks_in_grid), + SmallAndMediumPolicyT::BLOCK_THREADS, + (long long) stream); +#endif // CUB_DEBUG_LOG + + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + small_and_medium_blocks_in_grid, SmallAndMediumPolicyT::BLOCK_THREADS, 0, stream) + .doit(small_kernel, + small_segments, + medium_segments, + medium_blocks, + small_segments_indices, + large_and_medium_segments_indices + num_segments - medium_segments, + d_current_keys, + d_final_keys, + d_current_values, + d_final_values, + d_begin_offsets, + d_end_offsets); + + // Check for failure to launch + error = CubDebug(cudaPeekAtLastError()); + if (cudaSuccess != error) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + } + + return error; +} + +#ifdef CUB_RDC_ENABLED +/* + * Continuation kernel is used only in the CDP mode. It's used to + * launch DeviceSegmentedSortContinuation as a separate kernel. + */ +template +__launch_bounds__(1) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContinuationKernel( + LargeKernelT large_kernel, + SmallKernelT small_kernel, + local_segment_index_t num_segments, + KeyT* d_current_keys, + KeyT* d_final_keys, + device_double_buffer d_keys_double_buffer, + ValueT* d_current_values, + ValueT* d_final_values, + device_double_buffer d_values_double_buffer, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + local_segment_index_t* group_sizes, + local_segment_index_t* large_and_medium_segments_indices, + local_segment_index_t* small_segments_indices) +{ + using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; + using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; + using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; + + // In case of CDP: + // 1. each CTA has a different main stream + // 2. all streams are non-blocking + // 3. child grid always completes before the parent grid + // 4. streams can be used only from the CTA in which they were created + // 5. streams created on the host cannot be used on the device + // + // Due to (4, 5), we can't pass the user-provided stream in the continuation. + // Due to (1, 2, 3) it's safe to pass the main stream. + cudaError_t error = + detail::segmented_sort::DeviceSegmentedSortContinuation( + large_kernel, + small_kernel, + num_segments, + d_current_keys, + d_final_keys, + d_keys_double_buffer, + d_current_values, + d_final_values, + d_values_double_buffer, + d_begin_offsets, + d_end_offsets, + group_sizes, + large_and_medium_segments_indices, + small_segments_indices, + 0); // always launching on the main stream (see motivation above) + + error = CubDebug(error); +} +#endif // CUB_RDC_ENABLED + +} // namespace detail::segmented_sort +CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_nvrtc.cu b/cub/test/catch2_test_nvrtc.cu index a7b78f42d17..c7c9ccabeb9 100644 --- a/cub/test/catch2_test_nvrtc.cu +++ b/cub/test/catch2_test_nvrtc.cu @@ -84,6 +84,7 @@ TEST_CASE("Test nvrtc", "[test][nvrtc]") #include #include #include + #include #include #include #include From ca729cfa49100a23086eabaa47df044e8770b9e5 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 11 Aug 2025 22:41:18 +0000 Subject: [PATCH 002/100] Remove unused code --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 4 ---- 1 file changed, 4 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 8f9f12a4491..9cad13aea5c 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -60,10 +60,6 @@ CUB_NAMESPACE_BEGIN -namespace detail::segmented_sort -{ -} // namespace detail::segmented_sort - template Date: Tue, 12 Aug 2025 13:39:08 +0000 Subject: [PATCH 003/100] Move continuation back to dispatch file --- .../dispatch/dispatch_segmented_sort.cuh | 198 ++++++++++++++++++ .../dispatch/kernels/segmented_sort.cuh | 197 ----------------- 2 files changed, 198 insertions(+), 197 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 9cad13aea5c..cda370d9e03 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -60,6 +60,204 @@ CUB_NAMESPACE_BEGIN +namespace detail::segmented_sort +{ +/* + * Continuation is called after the partitioning stage. It launches kernels + * to sort large and small segments using the partitioning results. Separation + * of this stage is required to eliminate device-side synchronization in + * the CDP mode. + */ +template +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortContinuation( + LargeKernelT large_kernel, + SmallKernelT small_kernel, + int num_segments, + KeyT* d_current_keys, + KeyT* d_final_keys, + device_double_buffer d_keys_double_buffer, + ValueT* d_current_values, + ValueT* d_final_values, + device_double_buffer d_values_double_buffer, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + local_segment_index_t* group_sizes, + local_segment_index_t* large_and_medium_segments_indices, + local_segment_index_t* small_segments_indices, + cudaStream_t stream) +{ + using local_segment_index_t = local_segment_index_t; + + cudaError error = cudaSuccess; + + const local_segment_index_t large_segments = group_sizes[0]; + + if (large_segments > 0) + { + // One CTA per segment + const local_segment_index_t blocks_in_grid = large_segments; + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking " + "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", + static_cast(blocks_in_grid), + LargeSegmentPolicyT::BLOCK_THREADS, + (long long) stream); +#endif // CUB_DEBUG_LOG + + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(blocks_in_grid, LargeSegmentPolicyT::BLOCK_THREADS, 0, stream) + .doit(large_kernel, + large_and_medium_segments_indices, + d_current_keys, + d_final_keys, + d_keys_double_buffer, + d_current_values, + d_final_values, + d_values_double_buffer, + d_begin_offsets, + d_end_offsets); + + // Check for failure to launch + error = CubDebug(cudaPeekAtLastError()); + if (cudaSuccess != error) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + } + + const local_segment_index_t small_segments = group_sizes[1]; + const local_segment_index_t medium_segments = + static_cast(num_segments) - (large_segments + small_segments); + + const local_segment_index_t small_blocks = + ::cuda::ceil_div(small_segments, SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); + + const local_segment_index_t medium_blocks = + ::cuda::ceil_div(medium_segments, SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); + + const local_segment_index_t small_and_medium_blocks_in_grid = small_blocks + medium_blocks; + + if (small_and_medium_blocks_in_grid) + { +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking " + "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", + static_cast(small_and_medium_blocks_in_grid), + SmallAndMediumPolicyT::BLOCK_THREADS, + (long long) stream); +#endif // CUB_DEBUG_LOG + + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + small_and_medium_blocks_in_grid, SmallAndMediumPolicyT::BLOCK_THREADS, 0, stream) + .doit(small_kernel, + small_segments, + medium_segments, + medium_blocks, + small_segments_indices, + large_and_medium_segments_indices + num_segments - medium_segments, + d_current_keys, + d_final_keys, + d_current_values, + d_final_values, + d_begin_offsets, + d_end_offsets); + + // Check for failure to launch + error = CubDebug(cudaPeekAtLastError()); + if (cudaSuccess != error) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + } + + return error; +} + +#ifdef CUB_RDC_ENABLED +/* + * Continuation kernel is used only in the CDP mode. It's used to + * launch DeviceSegmentedSortContinuation as a separate kernel. + */ +template +__launch_bounds__(1) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContinuationKernel( + LargeKernelT large_kernel, + SmallKernelT small_kernel, + local_segment_index_t num_segments, + KeyT* d_current_keys, + KeyT* d_final_keys, + device_double_buffer d_keys_double_buffer, + ValueT* d_current_values, + ValueT* d_final_values, + device_double_buffer d_values_double_buffer, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + local_segment_index_t* group_sizes, + local_segment_index_t* large_and_medium_segments_indices, + local_segment_index_t* small_segments_indices) +{ + using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; + using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; + using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; + + // In case of CDP: + // 1. each CTA has a different main stream + // 2. all streams are non-blocking + // 3. child grid always completes before the parent grid + // 4. streams can be used only from the CTA in which they were created + // 5. streams created on the host cannot be used on the device + // + // Due to (4, 5), we can't pass the user-provided stream in the continuation. + // Due to (1, 2, 3) it's safe to pass the main stream. + cudaError_t error = + detail::segmented_sort::DeviceSegmentedSortContinuation( + large_kernel, + small_kernel, + num_segments, + d_current_keys, + d_final_keys, + d_keys_double_buffer, + d_current_values, + d_final_values, + d_values_double_buffer, + d_begin_offsets, + d_end_offsets, + group_sizes, + large_and_medium_segments_indices, + small_segments_indices, + 0); // always launching on the main stream (see motivation above) + + error = CubDebug(error); +} +#endif // CUB_RDC_ENABLED +} // namespace detail::segmented_sort + template #include -#include - CUB_NAMESPACE_BEGIN namespace detail::segmented_sort { @@ -473,200 +471,5 @@ __launch_bounds__(ChainedPolicyT::ActivePolicy::LargeSegmentPolicy::BLOCK_THREAD } } } -/* - * Continuation is called after the partitioning stage. It launches kernels - * to sort large and small segments using the partitioning results. Separation - * of this stage is required to eliminate device-side synchronization in - * the CDP mode. - */ -template -CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortContinuation( - LargeKernelT large_kernel, - SmallKernelT small_kernel, - int num_segments, - KeyT* d_current_keys, - KeyT* d_final_keys, - device_double_buffer d_keys_double_buffer, - ValueT* d_current_values, - ValueT* d_final_values, - device_double_buffer d_values_double_buffer, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets, - local_segment_index_t* group_sizes, - local_segment_index_t* large_and_medium_segments_indices, - local_segment_index_t* small_segments_indices, - cudaStream_t stream) -{ - using local_segment_index_t = local_segment_index_t; - - cudaError error = cudaSuccess; - - const local_segment_index_t large_segments = group_sizes[0]; - - if (large_segments > 0) - { - // One CTA per segment - const local_segment_index_t blocks_in_grid = large_segments; - -#ifdef CUB_DEBUG_LOG - _CubLog("Invoking " - "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", - static_cast(blocks_in_grid), - LargeSegmentPolicyT::BLOCK_THREADS, - (long long) stream); -#endif // CUB_DEBUG_LOG - - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(blocks_in_grid, LargeSegmentPolicyT::BLOCK_THREADS, 0, stream) - .doit(large_kernel, - large_and_medium_segments_indices, - d_current_keys, - d_final_keys, - d_keys_double_buffer, - d_current_values, - d_final_values, - d_values_double_buffer, - d_begin_offsets, - d_end_offsets); - - // Check for failure to launch - error = CubDebug(cudaPeekAtLastError()); - if (cudaSuccess != error) - { - return error; - } - - // Sync the stream if specified to flush runtime errors - error = CubDebug(DebugSyncStream(stream)); - if (cudaSuccess != error) - { - return error; - } - } - - const local_segment_index_t small_segments = group_sizes[1]; - const local_segment_index_t medium_segments = - static_cast(num_segments) - (large_segments + small_segments); - - const local_segment_index_t small_blocks = - ::cuda::ceil_div(small_segments, SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); - - const local_segment_index_t medium_blocks = - ::cuda::ceil_div(medium_segments, SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); - - const local_segment_index_t small_and_medium_blocks_in_grid = small_blocks + medium_blocks; - - if (small_and_medium_blocks_in_grid) - { -#ifdef CUB_DEBUG_LOG - _CubLog("Invoking " - "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", - static_cast(small_and_medium_blocks_in_grid), - SmallAndMediumPolicyT::BLOCK_THREADS, - (long long) stream); -#endif // CUB_DEBUG_LOG - - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( - small_and_medium_blocks_in_grid, SmallAndMediumPolicyT::BLOCK_THREADS, 0, stream) - .doit(small_kernel, - small_segments, - medium_segments, - medium_blocks, - small_segments_indices, - large_and_medium_segments_indices + num_segments - medium_segments, - d_current_keys, - d_final_keys, - d_current_values, - d_final_values, - d_begin_offsets, - d_end_offsets); - - // Check for failure to launch - error = CubDebug(cudaPeekAtLastError()); - if (cudaSuccess != error) - { - return error; - } - - // Sync the stream if specified to flush runtime errors - error = CubDebug(DebugSyncStream(stream)); - if (cudaSuccess != error) - { - return error; - } - } - - return error; -} - -#ifdef CUB_RDC_ENABLED -/* - * Continuation kernel is used only in the CDP mode. It's used to - * launch DeviceSegmentedSortContinuation as a separate kernel. - */ -template -__launch_bounds__(1) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContinuationKernel( - LargeKernelT large_kernel, - SmallKernelT small_kernel, - local_segment_index_t num_segments, - KeyT* d_current_keys, - KeyT* d_final_keys, - device_double_buffer d_keys_double_buffer, - ValueT* d_current_values, - ValueT* d_final_values, - device_double_buffer d_values_double_buffer, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets, - local_segment_index_t* group_sizes, - local_segment_index_t* large_and_medium_segments_indices, - local_segment_index_t* small_segments_indices) -{ - using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; - using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - - // In case of CDP: - // 1. each CTA has a different main stream - // 2. all streams are non-blocking - // 3. child grid always completes before the parent grid - // 4. streams can be used only from the CTA in which they were created - // 5. streams created on the host cannot be used on the device - // - // Due to (4, 5), we can't pass the user-provided stream in the continuation. - // Due to (1, 2, 3) it's safe to pass the main stream. - cudaError_t error = - detail::segmented_sort::DeviceSegmentedSortContinuation( - large_kernel, - small_kernel, - num_segments, - d_current_keys, - d_final_keys, - d_keys_double_buffer, - d_current_values, - d_final_values, - d_values_double_buffer, - d_begin_offsets, - d_end_offsets, - group_sizes, - large_and_medium_segments_indices, - small_segments_indices, - 0); // always launching on the main stream (see motivation above) - - error = CubDebug(error); -} -#endif // CUB_RDC_ENABLED - } // namespace detail::segmented_sort CUB_NAMESPACE_END From fbcbc834b388723a2a59d9a8cbff0424be39cc12 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 12 Aug 2025 22:09:23 +0000 Subject: [PATCH 004/100] Begin working on dynamic dispatch for segmented sort --- cub/cub/device/device_partition.cuh | 4 +- .../dispatch/dispatch_segmented_sort.cuh | 368 +++++++++--------- .../dispatch/tuning/tuning_segmented_sort.cuh | 89 +++++ 3 files changed, 271 insertions(+), 190 deletions(-) diff --git a/cub/cub/device/device_partition.cuh b/cub/cub/device/device_partition.cuh index 54fad690408..259b2d40d63 100644 --- a/cub/cub/device/device_partition.cuh +++ b/cub/cub/device/device_partition.cuh @@ -382,7 +382,9 @@ private: typename OffsetT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT, - typename PolicyHub> + typename PolicyHub, + typename KernelSource, + typename KernelLauncherFactory> friend class DispatchSegmentedSort; // Internal version without NVTX range diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index cda370d9e03..3f27e96c706 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -68,14 +68,14 @@ namespace detail::segmented_sort * of this stage is required to eliminate device-side synchronization in * the CDP mode. */ -template + typename EndOffsetIteratorT, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortContinuation( LargeKernelT large_kernel, SmallKernelT small_kernel, @@ -91,7 +91,9 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont local_segment_index_t* group_sizes, local_segment_index_t* large_and_medium_segments_indices, local_segment_index_t* small_segments_indices, - cudaStream_t stream) + cudaStream_t stream, + KernelLauncherFactory launcher_factory, + WrappedPolicyT wrapped_policy) { using local_segment_index_t = local_segment_index_t; @@ -108,11 +110,11 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont _CubLog("Invoking " "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", static_cast(blocks_in_grid), - LargeSegmentPolicyT::BLOCK_THREADS, + wrapped_policy.LargeSegment().BlockThreads(), (long long) stream); #endif // CUB_DEBUG_LOG - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(blocks_in_grid, LargeSegmentPolicyT::BLOCK_THREADS, 0, stream) + launcher_factory(blocks_in_grid, wrapped_policy.LargeSegment().BlockThreads(), 0, stream) .doit(large_kernel, large_and_medium_segments_indices, d_current_keys, @@ -143,11 +145,10 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont const local_segment_index_t medium_segments = static_cast(num_segments) - (large_segments + small_segments); - const local_segment_index_t small_blocks = - ::cuda::ceil_div(small_segments, SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); + const local_segment_index_t small_blocks = ::cuda::ceil_div(small_segments, wrapped_policy.SegmentsPerSmallBlock()); const local_segment_index_t medium_blocks = - ::cuda::ceil_div(medium_segments, SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); + ::cuda::ceil_div(medium_segments, wrapped_policy.SegmentsPerMediumBlock()); const local_segment_index_t small_and_medium_blocks_in_grid = small_blocks + medium_blocks; @@ -157,12 +158,12 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont _CubLog("Invoking " "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", static_cast(small_and_medium_blocks_in_grid), - SmallAndMediumPolicyT::BLOCK_THREADS, + small_and_medium_policy.BlockThreads(), (long long) stream); #endif // CUB_DEBUG_LOG - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( - small_and_medium_blocks_in_grid, SmallAndMediumPolicyT::BLOCK_THREADS, 0, stream) + launcher_factory( + small_and_medium_blocks_in_grid, wrapped_policy.SmallAndMediumSegmentedSort().BlockThreads(), 0, stream) .doit(small_kernel, small_segments, medium_segments, @@ -199,13 +200,14 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont * Continuation kernel is used only in the CDP mode. It's used to * launch DeviceSegmentedSortContinuation as a separate kernel. */ -template + typename EndOffsetIteratorT, + typename KernelLauncherFactory> __launch_bounds__(1) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContinuationKernel( LargeKernelT large_kernel, SmallKernelT small_kernel, @@ -220,12 +222,10 @@ __launch_bounds__(1) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContin EndOffsetIteratorT d_end_offsets, local_segment_index_t* group_sizes, local_segment_index_t* large_and_medium_segments_indices, - local_segment_index_t* small_segments_indices) + local_segment_index_t* small_segments_indices, + KernelLauncherFactory launcher_factory, + WrappedPolicyT wrapped_policy) { - using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; - using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - // In case of CDP: // 1. each CTA has a different main stream // 2. all streams are non-blocking @@ -235,27 +235,49 @@ __launch_bounds__(1) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContin // // Due to (4, 5), we can't pass the user-provided stream in the continuation. // Due to (1, 2, 3) it's safe to pass the main stream. - cudaError_t error = - detail::segmented_sort::DeviceSegmentedSortContinuation( - large_kernel, - small_kernel, - num_segments, - d_current_keys, - d_final_keys, - d_keys_double_buffer, - d_current_values, - d_final_values, - d_values_double_buffer, - d_begin_offsets, - d_end_offsets, - group_sizes, - large_and_medium_segments_indices, - small_segments_indices, - 0); // always launching on the main stream (see motivation above) + cudaError_t error = detail::segmented_sort::DeviceSegmentedSortContinuation( + large_kernel, + small_kernel, + num_segments, + d_current_keys, + d_final_keys, + d_keys_double_buffer, + d_current_values, + d_final_values, + d_values_double_buffer, + d_begin_offsets, + d_end_offsets, + group_sizes, + large_and_medium_segments_indices, + small_segments_indices, + 0, // always launching on the main stream (see motivation above) + launcher_factory, + wrapped_policy); error = CubDebug(error); } #endif // CUB_RDC_ENABLED +template +struct DeviceSegmentedSortKernelSource +{ + CUB_DEFINE_KERNEL_GETTER( + SegmentedSortFallbackKernel, + DeviceSegmentedSortFallbackKernel); + + CUB_DEFINE_KERNEL_GETTER( + SegmentedSortKernelSmall, + DeviceSegmentedSortKernelSmall); + + CUB_DEFINE_KERNEL_GETTER( + SegmentedSortKernelLarge, + DeviceSegmentedSortKernelLarge); +}; } // namespace detail::segmented_sort template > + typename PolicyHub = detail::segmented_sort::policy_hub, + typename KernelSource = detail::segmented_sort::DeviceSegmentedSortKernelSource< + typename PolicyHub::MaxPolicy, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> struct DispatchSegmentedSort { using local_segment_index_t = detail::segmented_sort::local_segment_index_t; @@ -369,48 +400,39 @@ struct DispatchSegmentedSort /// CUDA stream to launch kernels within. cudaStream_t stream; - CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchSegmentedSort( - void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - DoubleBuffer& d_values, - ::cuda::std::int64_t num_items, - global_segment_offset_t num_segments, - BeginOffsetIteratorT d_begin_offsets, - EndOffsetIteratorT d_end_offsets, - bool is_overwrite_okay, - cudaStream_t stream) - : d_temp_storage(d_temp_storage) - , temp_storage_bytes(temp_storage_bytes) - , d_keys(d_keys) - , d_values(d_values) - , num_items(num_items) - , num_segments(num_segments) - , d_begin_offsets(d_begin_offsets) - , d_end_offsets(d_end_offsets) - , is_overwrite_okay(is_overwrite_okay) - , stream(stream) - {} + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; template - CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke() + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT policy = {}) { - using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; + auto wrapped_policy = detail::segmented_sort::MakeSegmentedSortPolicyWrapper(policy); - static_assert(LargeSegmentPolicyT::LOAD_MODIFIER != CacheLoadModifier::LOAD_LDG, - "The memory consistency model does not apply to texture accesses"); + // using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; + // using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - static_assert(KEYS_ONLY || LargeSegmentPolicyT::LOAD_ALGORITHM != BLOCK_LOAD_STRIPED - || SmallAndMediumPolicyT::MediumPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED - || SmallAndMediumPolicyT::SmallPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED, - "Striped load will make this algorithm unstable"); + wrapped_policy.CheckLoadModifierIsNotLDG(); - static_assert(SmallAndMediumPolicyT::MediumPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED - || SmallAndMediumPolicyT::SmallPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED, - "Striped stores will produce unsorted results"); + // static_assert(LargeSegmentPolicyT::LOAD_MODIFIER != CacheLoadModifier::LOAD_LDG, + // "The memory consistency model does not apply to texture accesses"); - constexpr int radix_bits = LargeSegmentPolicyT::RADIX_BITS; + if constexpr (!KEYS_ONLY) + { + wrapped_policy.CheckLoadAlgorithmIsNotStriped(); + } + // static_assert(KEYS_ONLY || LargeSegmentPolicyT::LOAD_ALGORITHM != BLOCK_LOAD_STRIPED + // || SmallAndMediumPolicyT::MediumPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED + // || SmallAndMediumPolicyT::SmallPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED, + // "Striped load will make this algorithm unstable"); + + wrapped_policy.CheckStoreAlgorithmIsNotStriped(); + // static_assert(SmallAndMediumPolicyT::MediumPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED + // || SmallAndMediumPolicyT::SmallPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED, + // "Striped stores will produce unsorted results"); + + const int radix_bits = wrapped_policy.LargeSegmentRadixBits(); + // constexpr int radix_bits = LargeSegmentPolicyT::RADIX_BITS; cudaError error = cudaSuccess; @@ -420,7 +442,8 @@ struct DispatchSegmentedSort // Prepare temporary storage layout //------------------------------------------------------------------------ - const bool partition_segments = num_segments > ActivePolicyT::PARTITIONING_THRESHOLD; + const bool partition_segments = num_segments > wrapped_policy.PartitioningThreshold(); + // const bool partition_segments = num_segments > ActivePolicyT::PARTITIONING_THRESHOLD; cub::detail::temporary_storage::layout<5> temporary_storage_layout; @@ -451,10 +474,10 @@ struct DispatchSegmentedSort size_t three_way_partition_temp_storage_bytes{}; LargeSegmentsSelectorT large_segments_selector( - SmallAndMediumPolicyT::MediumPolicyT::ITEMS_PER_TILE, d_begin_offsets, d_end_offsets); + wrapped_policy.MediumPolicyItemsPerTile(), d_begin_offsets, d_end_offsets); SmallSegmentsSelectorT small_segments_selector( - SmallAndMediumPolicyT::SmallPolicyT::ITEMS_PER_TILE + 1, d_begin_offsets, d_end_offsets); + wrapped_policy.SmallPolicyItemsPerTile() + 1, d_begin_offsets, d_end_offsets); auto device_partition_temp_storage = keys_slot->create_alias(); @@ -579,23 +602,9 @@ struct DispatchSegmentedSort { // Partition input segments into size groups and assign specialized // kernels for each of them. - error = SortWithPartitioning( - detail::segmented_sort::DeviceSegmentedSortKernelLarge< - Order, - MaxPolicyT, - KeyT, - ValueT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT>, - detail::segmented_sort::DeviceSegmentedSortKernelSmall< - Order, - MaxPolicyT, - KeyT, - ValueT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT>, + error = SortWithPartitioning( + kernel_source.SegmentedSortKernelLarge(), + kernel_source.SegmentedSortKernelSmall(), three_way_partition_temp_storage_bytes, d_keys_double_buffer, d_values_double_buffer, @@ -604,24 +613,16 @@ struct DispatchSegmentedSort device_partition_temp_storage, large_and_medium_segments_indices, small_segments_indices, - group_sizes); + group_sizes, + wrapped_policy); } else { // If there are not enough segments, there's no reason to spend time // on extra partitioning steps. - error = SortWithoutPartitioning( - detail::segmented_sort::DeviceSegmentedSortFallbackKernel< - Order, - MaxPolicyT, - KeyT, - ValueT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT>, - d_keys_double_buffer, - d_values_double_buffer); + error = SortWithoutPartitioning( + kernel_source.SegmentedSortFallbackKernel(), d_keys_double_buffer, d_values_double_buffer, wrapped_policy); } d_keys.selector = GetFinalSelector(d_keys.selector, radix_bits); @@ -632,6 +633,7 @@ struct DispatchSegmentedSort return error; } + template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -642,42 +644,35 @@ struct DispatchSegmentedSort BeginOffsetIteratorT d_begin_offsets, EndOffsetIteratorT d_end_offsets, bool is_overwrite_okay, - cudaStream_t stream) + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) { - cudaError error = cudaSuccess; - - do + // Get PTX version + int ptx_version = 0; + if (cudaError error = CubDebug(launcher_factory.PtxVersion(ptx_version)); cudaSuccess != error) { - // Get PTX version - int ptx_version = 0; - error = CubDebug(PtxVersion(ptx_version)); - if (cudaSuccess != error) - { - break; - } + return error; + } - // Create dispatch functor - DispatchSegmentedSort dispatch( - d_temp_storage, - temp_storage_bytes, - d_keys, - d_values, - num_items, - num_segments, - d_begin_offsets, - d_end_offsets, - is_overwrite_okay, - stream); - - // Dispatch to chained policy - error = CubDebug(PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch)); - if (cudaSuccess != error) - { - break; - } - } while (false); + // Create dispatch functor + DispatchSegmentedSort dispatch{ + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + is_overwrite_okay, + stream, + kernel_source, + launcher_factory}; - return error; + // Dispatch to chained policy + return CubDebug(max_policy.Invoke(ptx_version, dispatch)); } private: @@ -707,7 +702,7 @@ private: return buffer.d_buffers[final_selector]; } - template + template CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t SortWithPartitioning( LargeKernelT large_kernel, SmallKernelT small_kernel, @@ -719,7 +714,8 @@ private: cub::detail::temporary_storage::alias& device_partition_temp_storage, cub::detail::temporary_storage::alias& large_and_medium_segments_indices, cub::detail::temporary_storage::alias& small_segments_indices, - cub::detail::temporary_storage::alias& group_sizes) + cub::detail::temporary_storage::alias& group_sizes, + WrappedPolicyT wrapped_policy) { cudaError_t error = cudaSuccess; @@ -771,43 +767,35 @@ private: #else // CUB_RDC_ENABLED -# define CUB_TEMP_DEVICE_CODE \ - error = \ - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(1, 1, 0, stream) \ - .doit( \ - detail::segmented_sort::DeviceSegmentedSortContinuationKernel< \ - typename PolicyHub::MaxPolicy, \ - LargeKernelT, \ - SmallKernelT, \ - KeyT, \ - ValueT, \ - BeginOffsetIteratorT, \ - EndOffsetIteratorT>, \ - large_kernel, \ - small_kernel, \ - current_num_segments, \ - d_keys.Current(), \ - GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_keys), \ - d_keys_double_buffer, \ - d_values.Current(), \ - GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_values), \ - d_values_double_buffer, \ - current_begin_offset, \ - current_end_offset, \ - group_sizes.get(), \ - large_and_medium_segments_indices.get(), \ - small_segments_indices.get()); \ - error = CubDebug(error); \ - \ - if (cudaSuccess != error) \ - { \ - return error; \ - } \ - \ - error = CubDebug(detail::DebugSyncStream(stream)); \ - if (cudaSuccess != error) \ - { \ - return error; \ +# define CUB_TEMP_DEVICE_CODE \ + error = \ + launcher_factory(1, 1, 0, stream) \ + .doit(kernel_source.SegmentedSortContinuationKernel(), \ + large_kernel, \ + small_kernel, \ + current_num_segments, \ + d_keys.Current(), \ + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), \ + d_keys_double_buffer, \ + d_values.Current(), \ + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), \ + d_values_double_buffer, \ + current_begin_offset, \ + current_end_offset, \ + group_sizes.get(), \ + large_and_medium_segments_indices.get(), \ + small_segments_indices.get()); \ + error = CubDebug(error); \ + \ + if (cudaSuccess != error) \ + { \ + return error; \ + } \ + \ + error = CubDebug(detail::DebugSyncStream(stream)); \ + if (cudaSuccess != error) \ + { \ + return error; \ } #endif // CUB_RDC_ENABLED @@ -818,7 +806,7 @@ private: NV_IS_HOST, ( local_segment_index_t h_group_sizes[num_selected_groups]; - error = CubDebug(cudaMemcpyAsync(h_group_sizes, + error = CubDebug(launcher_factory.MemcpyAsync(h_group_sizes, group_sizes.get(), num_selected_groups * sizeof(local_segment_index_t), @@ -836,23 +824,24 @@ private: return error; } - error = detail::segmented_sort::DeviceSegmentedSortContinuation( + error = detail::segmented_sort::DeviceSegmentedSortContinuation( large_kernel, small_kernel, current_num_segments, d_keys.Current(), - GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_keys), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), d_keys_double_buffer, d_values.Current(), - GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_values), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), d_values_double_buffer, current_begin_offset, current_end_offset, h_group_sizes, large_and_medium_segments_indices.get(), small_segments_indices.get(), - stream);), + stream, + launcher_factory, + wrapped_policy);), // NV_IS_DEVICE: (CUB_TEMP_DEVICE_CODE)); // clang-format on @@ -862,16 +851,17 @@ private: return error; } - template + template CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t SortWithoutPartitioning( FallbackKernelT fallback_kernel, cub::detail::device_double_buffer& d_keys_double_buffer, - cub::detail::device_double_buffer& d_values_double_buffer) + cub::detail::device_double_buffer& d_values_double_buffer, + WrappedPolicyT wrapped_policy) { cudaError_t error = cudaSuccess; const auto blocks_in_grid = static_cast(num_segments); - constexpr auto threads_in_block = static_cast(LargeSegmentPolicyT::BLOCK_THREADS); + constexpr auto threads_in_block = static_cast(wrapped_policy.LargeSegment().BlockThreads()); // Log kernel configuration #ifdef CUB_DEBUG_LOG @@ -880,18 +870,18 @@ private: blocks_in_grid, threads_in_block, (long long) stream, - LargeSegmentPolicyT::ITEMS_PER_THREAD, - LargeSegmentPolicyT::RADIX_BITS); + wrapped_policy.LargeSegment().ItemsPerThread(), + wrapped_policy.LargeSegmentRadixBits()); #endif // CUB_DEBUG_LOG // Invoke fallback kernel - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(blocks_in_grid, threads_in_block, 0, stream) + launcher_factory(blocks_in_grid, threads_in_block, 0, stream) .doit(fallback_kernel, d_keys.Current(), - GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_keys), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), d_keys_double_buffer, d_values.Current(), - GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_values), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), d_values_double_buffer, d_begin_offsets, d_end_offsets); diff --git a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index aa3c920dbc5..a360ffe4f28 100644 --- a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -47,6 +47,95 @@ namespace detail { namespace segmented_sort { + +template +struct SegmentedSortPolicyWrapper : PolicyT +{ + CUB_RUNTIME_FUNCTION SegmentedSortPolicyWrapper(PolicyT base) + : PolicyT(base) + {} +}; + +template +struct SegmentedSortPolicyWrapper> + : StaticPolicyT +{ + CUB_RUNTIME_FUNCTION SegmentedSortPolicyWrapper(StaticPolicyT base) + : StaticPolicyT(base) + {} + + CUB_RUNTIME_FUNCTION static constexpr auto LargeSegment() + { + return cub::detail::MakePolicyWrapper(typename StaticPolicyT::LargeSegmentPolicy()); + } + + CUB_RUNTIME_FUNCTION static constexpr auto SmallAndMediumSegmentedSort() + { + return cub::detail::MakePolicyWrapper(typename StaticPolicyT::SmallAndMediumSegmentedSortPolicyT()); + } + + CUB_RUNTIME_FUNCTION static constexpr void CheckLoadModifierIsNotLDG() + { + static_assert(StaticPolicyT::LargeSegmentPolicy::LOAD_MODIFIER != CacheLoadModifier::LOAD_LDG, + "The memory consistency model does not apply to texture accesses"); + } + + CUB_RUNTIME_FUNCTION static constexpr void CheckLoadAlgorithmIsNotStriped() + { + static_assert( + StaticPolicyT::LargeSegmentPolicy::LOAD_ALGORITHM != BLOCK_LOAD_STRIPED + || StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::MediumPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED + || StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::SmallPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED, + "Striped load will make this algorithm unstable"); + } + + CUB_RUNTIME_FUNCTION static constexpr void CheckStoreAlgorithmIsNotStriped() + { + static_assert( + StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::MediumPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED + || StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::SmallPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED, + "Striped stores will produce unsorted results"); + } + + CUB_RUNTIME_FUNCTION static constexpr int PartitioningThreshold() + { + return StaticPolicyT::PARTITIONING_THRESHOLD; + } + + CUB_RUNTIME_FUNCTION static constexpr int LargeSegmentRadixBits() + { + return StaticPolicyT::LargeSegmentPolicy::RADIX_BITS; + } + + CUB_RUNTIME_FUNCTION static constexpr int SegmentsPerSmallBlock() + { + return StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::SEGMENTS_PER_SMALL_BLOCK; + } + + CUB_RUNTIME_FUNCTION static constexpr int SegmentsPerMediumBlock() + { + return StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::SEGMENTS_PER_MEDIUM_BLOCK; + } + + CUB_RUNTIME_FUNCTION static constexpr int SmallPolicyItemsPerTile() + { + return StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::SmallPolicyT::ITEMS_PER_TILE; + } + + CUB_RUNTIME_FUNCTION static constexpr int MediumPolicyItemsPerTile() + { + return StaticPolicyT::SmallAndMediumSegmentedSortPolicyT::MediumPolicyT::ITEMS_PER_TILE; + } +}; + +template +CUB_RUNTIME_FUNCTION SegmentedSortPolicyWrapper MakeSegmentedSortPolicyWrapper(PolicyT policy) +{ + return SegmentedSortPolicyWrapper{policy}; +} + template struct policy_hub { From 6be7b9d94aef89f2b83d5bf3db450b8f5c6b9eee Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 13 Aug 2025 00:15:18 +0000 Subject: [PATCH 005/100] Fix compilation errors --- .../dispatch/dispatch_segmented_sort.cuh | 29 ++++++++++--------- .../dispatch/tuning/tuning_segmented_sort.cuh | 12 ++++++++ 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 3f27e96c706..abebf3450e2 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -110,11 +110,11 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont _CubLog("Invoking " "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", static_cast(blocks_in_grid), - wrapped_policy.LargeSegment().BlockThreads(), + wrapped_policy.BlockThreads(wrapped_policy.LargeSegment()), (long long) stream); #endif // CUB_DEBUG_LOG - launcher_factory(blocks_in_grid, wrapped_policy.LargeSegment().BlockThreads(), 0, stream) + launcher_factory(blocks_in_grid, wrapped_policy.BlockThreads(wrapped_policy.LargeSegment()), 0, stream) .doit(large_kernel, large_and_medium_segments_indices, d_current_keys, @@ -158,12 +158,14 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont _CubLog("Invoking " "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", static_cast(small_and_medium_blocks_in_grid), - small_and_medium_policy.BlockThreads(), + wrapped_policy.BlockThreads(wrapped_policy.SmallAndMediumSegmentedSort()), (long long) stream); #endif // CUB_DEBUG_LOG - launcher_factory( - small_and_medium_blocks_in_grid, wrapped_policy.SmallAndMediumSegmentedSort().BlockThreads(), 0, stream) + launcher_factory(small_and_medium_blocks_in_grid, + wrapped_policy.BlockThreads(wrapped_policy.SmallAndMediumSegmentedSort()), + 0, + stream) .doit(small_kernel, small_segments, medium_segments, @@ -807,11 +809,11 @@ private: ( local_segment_index_t h_group_sizes[num_selected_groups]; error = CubDebug(launcher_factory.MemcpyAsync(h_group_sizes, - group_sizes.get(), - num_selected_groups * - sizeof(local_segment_index_t), - cudaMemcpyDeviceToHost, - stream)); + group_sizes.get(), + num_selected_groups * + sizeof(local_segment_index_t), + cudaMemcpyDeviceToHost, + stream)); if (cudaSuccess != error) { @@ -860,8 +862,9 @@ private: { cudaError_t error = cudaSuccess; - const auto blocks_in_grid = static_cast(num_segments); - constexpr auto threads_in_block = static_cast(wrapped_policy.LargeSegment().BlockThreads()); + const auto blocks_in_grid = static_cast(num_segments); + constexpr auto threads_in_block = + static_cast(wrapped_policy.BlockThreads(wrapped_policy.LargeSegment())); // Log kernel configuration #ifdef CUB_DEBUG_LOG @@ -870,7 +873,7 @@ private: blocks_in_grid, threads_in_block, (long long) stream, - wrapped_policy.LargeSegment().ItemsPerThread(), + wrapped_policy.ItemsPerThread(wrapped_policy.LargeSegment()), wrapped_policy.LargeSegmentRadixBits()); #endif // CUB_DEBUG_LOG diff --git a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index a360ffe4f28..3056ddf3bc3 100644 --- a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -128,6 +128,18 @@ struct SegmentedSortPolicyWrapper + CUB_RUNTIME_FUNCTION static constexpr int BlockThreads(PolicyT /*policy*/) + { + return PolicyT::BLOCK_THREADS; + } + + template + CUB_RUNTIME_FUNCTION static constexpr int ItemsPerThread(PolicyT /*policy*/) + { + return PolicyT::ITEMS_PER_THREAD; + } }; template From 9ca779b7449f1862373adf2b3052b3f4d1dbd9b3 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 15:34:06 +0000 Subject: [PATCH 006/100] Add initial segmented sort c parallel implementation --- c/parallel/include/cccl/c/segmented_sort.h | 68 +++ c/parallel/src/segmented_sort.cu | 636 +++++++++++++++++++++ c/parallel/test/test_segmented_sort.cpp | 596 +++++++++++++++++++ 3 files changed, 1300 insertions(+) create mode 100644 c/parallel/include/cccl/c/segmented_sort.h create mode 100644 c/parallel/src/segmented_sort.cu create mode 100644 c/parallel/test/test_segmented_sort.cpp diff --git a/c/parallel/include/cccl/c/segmented_sort.h b/c/parallel/include/cccl/c/segmented_sort.h new file mode 100644 index 00000000000..29a9fbfae9b --- /dev/null +++ b/c/parallel/include/cccl/c/segmented_sort.h @@ -0,0 +1,68 @@ +//===----------------------------------------------------------------------===// +// +// Part of CUDA Experimental in CUDA Core Compute Libraries, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +#pragma once + +#ifndef CCCL_C_EXPERIMENTAL +# error "C exposure is experimental and subject to change. Define CCCL_C_EXPERIMENTAL to acknowledge this notice." +#endif // !CCCL_C_EXPERIMENTAL + +#include +#include + +#include +#include + +CCCL_C_EXTERN_C_BEGIN + +typedef struct cccl_device_segmented_sort_build_result_t +{ + int cc; + void* cubin; + size_t cubin_size; + CUlibrary library; + CUkernel segmented_sort_fallback_kernel; + CUkernel segmented_sort_kernel_small; + CUkernel segmented_sort_kernel_large; + void* runtime_policy; +} cccl_device_segmented_sort_build_result_t; + +// TODO return a union of nvtx/cuda/nvrtc errors or a string? +CCCL_C_API CUresult cccl_device_segmented_sort_build( + cccl_device_segmented_sort_build_result_t* build, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_keys_out, + cccl_iterator_t d_values_in, + cccl_iterator_t d_values_out, + cccl_iterator_t begin_offset_in, + cccl_iterator_t end_offset_in, + int cc_major, + int cc_minor, + const char* cub_path, + const char* thrust_path, + const char* libcudacxx_path, + const char* ctk_path); + +CCCL_C_API CUresult cccl_device_segmented_sort( + cccl_device_segmented_sort_build_result_t build, + void* d_temp_storage, + size_t* temp_storage_bytes, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_keys_out, + cccl_iterator_t d_values_in, + cccl_iterator_t d_values_out, + uint64_t num_segments, + cccl_iterator_t start_offset_in, + cccl_iterator_t end_offset_in, + CUstream stream); + +CCCL_C_API CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_result_t* bld_ptr); + +CCCL_C_EXTERN_C_END diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu new file mode 100644 index 00000000000..0c6617bfa35 --- /dev/null +++ b/c/parallel/src/segmented_sort.cu @@ -0,0 +1,636 @@ +//===----------------------------------------------------------------------===// +// +// Part of CUDA Experimental in CUDA C++ Core Libraries, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +#include // cub::detail::choose_offset_t +#include // cub::detail::CudaDriverLauncherFactory +#include // cub::DispatchSegmentedSort +#include // DeviceSegmentedSort kernels +#include // policy_hub +#include // cub::LoadModifier + +#include // std::exception +#include // std::format +#include // std::string +#include // std::string_view +#include // std::is_same_v + +#include // printf + +#include "jit_templates/templates/input_iterator.h" +#include "jit_templates/templates/output_iterator.h" +#include "jit_templates/traits.h" +#include "util/context.h" +#include "util/errors.h" +#include "util/indirect_arg.h" +#include "util/runtime_policy.h" +#include "util/types.h" +#include +#include // cccl_type_info +#include +#include +#include + +struct device_segmented_sort_policy; +using OffsetT = unsigned long long; +static_assert(std::is_same_v, OffsetT>, "OffsetT must be size_t"); + +// check we can map OffsetT to ::cuda::std::uint64_t +static_assert(std::is_unsigned_v); +static_assert(sizeof(OffsetT) == sizeof(::cuda::std::uint64_t)); + +namespace segmented_sort +{ + +// Runtime policy structure for segmented sort +struct segmented_sort_runtime_policy +{ + int partitioning_threshold; + int large_segment_radix_bits; + int segments_per_small_block; + int segments_per_medium_block; + int small_policy_items_per_tile; + int medium_policy_items_per_tile; + + // Required methods for SegmentedSortPolicyWrapper + constexpr int PartitioningThreshold() const + { + return partitioning_threshold; + } + constexpr int LargeSegmentRadixBits() const + { + return large_segment_radix_bits; + } + constexpr int SegmentsPerSmallBlock() const + { + return segments_per_small_block; + } + constexpr int SegmentsPerMediumBlock() const + { + return segments_per_medium_block; + } + constexpr int SmallPolicyItemsPerTile() const + { + return small_policy_items_per_tile; + } + constexpr int MediumPolicyItemsPerTile() const + { + return medium_policy_items_per_tile; + } + + // Additional methods expected by SegmentedSortPolicyWrapper + constexpr void CheckLoadModifierIsNotLDG() const {} // No-op validation + constexpr void CheckLoadAlgorithmIsNotStriped() const {} // No-op validation + constexpr void CheckStoreAlgorithmIsNotStriped() const {} // No-op validation + + // Policy accessor methods + constexpr int BlockThreads(int /* large_segment_policy */) const + { + return 256; + } // Default block size + constexpr int LargeSegment() const + { + return 0; + } // Return index for large segment policy + constexpr auto SmallAndMediumSegmentedSort() const + { + return *this; + } // Return policy for small/medium segments + + using MaxPolicy = segmented_sort_runtime_policy; + + template + cudaError_t Invoke(int, F& op) + { + return op.template Invoke(*this); + } +}; + +// Function to create runtime policy from JSON +segmented_sort_runtime_policy from_json(const nlohmann::json& j) +{ + return segmented_sort_runtime_policy{ + .partitioning_threshold = j["PartitioningThreshold"].get(), + .large_segment_radix_bits = j["LargeSegmentRadixBits"].get(), + .segments_per_small_block = j["SegmentsPerSmallBlock"].get(), + .segments_per_medium_block = j["SegmentsPerMediumBlock"].get(), + .small_policy_items_per_tile = j["SmallPolicyItemsPerTile"].get(), + .medium_policy_items_per_tile = j["MediumPolicyItemsPerTile"].get()}; +} + +std::string get_device_segmented_sort_fallback_kernel_name( + std::string_view /* key_iterator_t */, + std::string_view /* value_iterator_t */, + std::string_view start_offset_iterator_t, + std::string_view end_offset_iterator_t, + std::string_view key_t, + std::string_view value_t) +{ + std::string chained_policy_t; + check(nvrtcGetTypeName(&chained_policy_t)); + + std::string offset_t; + check(nvrtcGetTypeName(&offset_t)); + + /* + template // 6 + DeviceSegmentedSortFallbackKernel(...); + */ + return std::format( + "cub::detail::segmented_sort::DeviceSegmentedSortFallbackKernel", + chained_policy_t, // 0 + key_t, // 1 + value_t, // 2 + start_offset_iterator_t, // 3 + end_offset_iterator_t, // 4 + offset_t); // 5 +} + +std::string get_device_segmented_sort_kernel_small_name( + std::string_view /* key_iterator_t */, + std::string_view /* value_iterator_t */, + std::string_view start_offset_iterator_t, + std::string_view end_offset_iterator_t, + std::string_view key_t, + std::string_view value_t) +{ + std::string chained_policy_t; + check(nvrtcGetTypeName(&chained_policy_t)); + + std::string offset_t; + check(nvrtcGetTypeName(&offset_t)); + + /* + template // 6 + DeviceSegmentedSortKernelSmall(...); + */ + return std::format( + "cub::detail::segmented_sort::DeviceSegmentedSortKernelSmall", + chained_policy_t, // 0 + key_t, // 1 + value_t, // 2 + start_offset_iterator_t, // 3 + end_offset_iterator_t, // 4 + offset_t); // 5 +} + +std::string get_device_segmented_sort_kernel_large_name( + std::string_view /* key_iterator_t */, + std::string_view /* value_iterator_t */, + std::string_view start_offset_iterator_t, + std::string_view end_offset_iterator_t, + std::string_view key_t, + std::string_view value_t) +{ + std::string chained_policy_t; + check(nvrtcGetTypeName(&chained_policy_t)); + + std::string offset_t; + check(nvrtcGetTypeName(&offset_t)); + + /* + template // 6 + DeviceSegmentedSortKernelLarge(...); + */ + return std::format( + "cub::detail::segmented_sort::DeviceSegmentedSortKernelLarge", + chained_policy_t, // 0 + key_t, // 1 + value_t, // 2 + start_offset_iterator_t, // 3 + end_offset_iterator_t, // 4 + offset_t); // 5 +} + +struct segmented_sort_kernel_source +{ + cccl_device_segmented_sort_build_result_t& build; + + CUkernel SegmentedSortFallbackKernel() const + { + return build.segmented_sort_fallback_kernel; + } + CUkernel SegmentedSortKernelSmall() const + { + return build.segmented_sort_kernel_small; + } + CUkernel SegmentedSortKernelLarge() const + { + return build.segmented_sort_kernel_large; + } +}; +} // namespace segmented_sort + +struct segmented_sort_keys_input_iterator_tag; +struct segmented_sort_keys_output_iterator_tag; +struct segmented_sort_values_input_iterator_tag; +struct segmented_sort_values_output_iterator_tag; +struct segmented_sort_start_offset_iterator_tag; +struct segmented_sort_end_offset_iterator_tag; + +CUresult cccl_device_segmented_sort_build( + cccl_device_segmented_sort_build_result_t* build_ptr, + cccl_iterator_t keys_in_it, + cccl_iterator_t keys_out_it, + cccl_iterator_t values_in_it, + cccl_iterator_t values_out_it, + cccl_iterator_t start_offset_it, + cccl_iterator_t end_offset_it, + int cc_major, + int cc_minor, + const char* cub_path, + const char* thrust_path, + const char* libcudacxx_path, + const char* ctk_path) +{ + CUresult error = CUDA_SUCCESS; + + try + { + const char* name = "device_segmented_sort"; + + const int cc = cc_major * 10 + cc_minor; + + // Get iterator specializations + const auto [keys_in_iterator_name, keys_in_iterator_src] = + get_specialization(template_id(), keys_in_it); + + const auto [keys_out_iterator_name, keys_out_iterator_src] = + get_specialization( + template_id(), keys_out_it, keys_in_it.value_type); + + // Determine if this is keys-only sorting + const bool keys_only = (values_in_it.state == nullptr || values_out_it.state == nullptr); + + std::string values_in_iterator_name, values_in_iterator_src; + std::string values_out_iterator_name, values_out_iterator_src; + + if (!keys_only) + { + const auto [vi_name, vi_src] = get_specialization( + template_id(), values_in_it); + values_in_iterator_name = vi_name; + values_in_iterator_src = vi_src; + + const auto [vo_name, vo_src] = get_specialization( + template_id(), values_out_it, values_in_it.value_type); + values_out_iterator_name = vo_name; + values_out_iterator_src = vo_src; + } + else + { + // For keys-only sorting, use NullType for values + values_in_iterator_name = "cub::NullType*"; + values_out_iterator_name = "cub::NullType*"; + values_in_iterator_src = ""; + values_out_iterator_src = ""; + } + + const auto [start_offset_iterator_name, start_offset_iterator_src] = + get_specialization( + template_id(), start_offset_it); + + const auto [end_offset_iterator_name, end_offset_iterator_src] = + get_specialization(template_id(), end_offset_it); + + // OffsetT is checked to match have 64-bit size + const auto offset_t = cccl_type_enum_to_name(cccl_type_enum::CCCL_UINT64); + + // Get type names for keys and values + const std::string key_t = cccl_type_enum_to_name(keys_in_it.value_type.type); + const std::string value_t = keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); + + const std::string dependent_definitions_src = std::format( + R"XXX( +struct __align__({1}) storage_t {{ + char data[{0}]; +}}; +{2} +{3} +{4} +{5} +{6} +{7} +)XXX", + keys_in_it.value_type.size, // 0 + keys_in_it.value_type.alignment, // 1 + keys_in_iterator_src, // 2 + keys_out_iterator_src, // 3 + values_in_iterator_src, // 4 + values_out_iterator_src, // 5 + start_offset_iterator_src, // 6 + end_offset_iterator_src); // 7 + + // Runtime parameter tuning + const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); + + constexpr size_t ptx_num_args = 5; + const char* ptx_args[ptx_num_args] = {ptx_arch.c_str(), cub_path, thrust_path, libcudacxx_path, "-rdc=true"}; + + static constexpr std::string_view policy_wrapper_expr_tmpl = + R"XXXX(cub::detail::segmented_sort::MakeSegmentedSortPolicyWrapper(cub::detail::segmented_sort::policy_hub<{0}, {1}>::MaxPolicy::ActivePolicy{{}}))XXXX"; + + const auto policy_wrapper_expr = std::format( + policy_wrapper_expr_tmpl, + key_t, // 0 + value_t); // 1 + + static constexpr std::string_view ptx_query_tu_src_tmpl = R"XXXX( +#include +{0} +{1} +)XXXX"; + + const auto ptx_query_tu_src = + std::format(ptx_query_tu_src_tmpl, jit_template_header_contents, dependent_definitions_src); + + nlohmann::json runtime_policy = get_policy(policy_wrapper_expr, ptx_query_tu_src, ptx_args); + + auto segmented_sort_policy = segmented_sort::from_json(runtime_policy); + + // Extract sub-policy information if available + std::string small_and_medium_policy_str; + if (runtime_policy.contains("SmallAndMediumSegmentedSort")) + { + auto sub_policy = runtime_policy["SmallAndMediumSegmentedSort"]; + auto block_threads = sub_policy["BlockThreads"].get(); + auto segments_per_medium = sub_policy["SegmentsPerMediumBlock"].get(); + auto segments_per_small = sub_policy["SegmentsPerSmallBlock"].get(); + + small_and_medium_policy_str = std::format( + R"XXX( + // Small and Medium Segment Policy + static constexpr int SMALL_MEDIUM_BLOCK_THREADS = {0}; + static constexpr int SMALL_MEDIUM_SEGMENTS_PER_MEDIUM_BLOCK = {1}; + static constexpr int SMALL_MEDIUM_SEGMENTS_PER_SMALL_BLOCK = {2};)XXX", + block_threads, + segments_per_medium, + segments_per_small); + } + + // Build the policy structure manually + const std::string segmented_sort_policy_str = std::format( + R"XXX( + static constexpr int PARTITIONING_THRESHOLD = {0}; + static constexpr int LARGE_SEGMENT_RADIX_BITS = {1}; + static constexpr int SEGMENTS_PER_SMALL_BLOCK = {2}; + static constexpr int SEGMENTS_PER_MEDIUM_BLOCK = {3}; + static constexpr int SMALL_POLICY_ITEMS_PER_TILE = {4}; + static constexpr int MEDIUM_POLICY_ITEMS_PER_TILE = {5};{6} + using MaxPolicy = cub::detail::segmented_sort::policy_hub<{7}, {8}>::MaxPolicy; +)XXX", + segmented_sort_policy.partitioning_threshold, // 0 + segmented_sort_policy.large_segment_radix_bits, // 1 + segmented_sort_policy.segments_per_small_block, // 2 + segmented_sort_policy.segments_per_medium_block, // 3 + segmented_sort_policy.small_policy_items_per_tile, // 4 + segmented_sort_policy.medium_policy_items_per_tile, // 5 + small_and_medium_policy_str, // 6 + key_t, // 7 + value_t); // 8 + + // agent_policy_t is to specify parameters like policy_hub does in dispatch_segmented_sort.cuh + constexpr std::string_view program_preamble_template = R"XXX( +#include +{0} +{1} +struct device_segmented_sort_policy {{ + struct ActivePolicy {{ + {2} + }}; +}}; +)XXX"; + + std::string final_src = std::format( + program_preamble_template, + jit_template_header_contents, // 0 + dependent_definitions_src, // 1 + segmented_sort_policy_str); // 2 + + std::string segmented_sort_fallback_kernel_name = segmented_sort::get_device_segmented_sort_fallback_kernel_name( + keys_in_iterator_name, + values_in_iterator_name, + start_offset_iterator_name, + end_offset_iterator_name, + key_t, + value_t); + + std::string segmented_sort_kernel_small_name = segmented_sort::get_device_segmented_sort_kernel_small_name( + keys_in_iterator_name, + values_in_iterator_name, + start_offset_iterator_name, + end_offset_iterator_name, + key_t, + value_t); + + std::string segmented_sort_kernel_large_name = segmented_sort::get_device_segmented_sort_kernel_large_name( + keys_in_iterator_name, + values_in_iterator_name, + start_offset_iterator_name, + end_offset_iterator_name, + key_t, + value_t); + + std::string segmented_sort_fallback_kernel_lowered_name; + std::string segmented_sort_kernel_small_lowered_name; + std::string segmented_sort_kernel_large_lowered_name; + + const std::string arch = std::format("-arch=sm_{0}{1}", cc_major, cc_minor); + + constexpr size_t num_args = 9; + const char* args[num_args] = { + arch.c_str(), + cub_path, + thrust_path, + libcudacxx_path, + ctk_path, + "-rdc=true", + "-dlto", + "-DCUB_DISABLE_CDP", + "-std=c++20"}; + + constexpr size_t num_lto_args = 2; + const char* lopts[num_lto_args] = {"-lto", arch.c_str()}; + + // Collect all LTO-IRs to be linked. + nvrtc_ltoir_list ltoir_list; + nvrtc_ltoir_list_appender appender{ltoir_list}; + + // add iterator definitions + appender.add_iterator_definition(keys_in_it); + appender.add_iterator_definition(keys_out_it); + if (!keys_only) + { + appender.add_iterator_definition(values_in_it); + appender.add_iterator_definition(values_out_it); + } + appender.add_iterator_definition(start_offset_it); + appender.add_iterator_definition(end_offset_it); + + nvrtc_link_result result = + begin_linking_nvrtc_program(num_lto_args, lopts) + ->add_program(nvrtc_translation_unit{final_src.c_str(), name}) + ->add_expression({segmented_sort_fallback_kernel_name}) + ->add_expression({segmented_sort_kernel_small_name}) + ->add_expression({segmented_sort_kernel_large_name}) + ->compile_program({args, num_args}) + ->get_name({segmented_sort_fallback_kernel_name, segmented_sort_fallback_kernel_lowered_name}) + ->get_name({segmented_sort_kernel_small_name, segmented_sort_kernel_small_lowered_name}) + ->get_name({segmented_sort_kernel_large_name, segmented_sort_kernel_large_lowered_name}) + ->link_program() + ->add_link_list(ltoir_list) + ->finalize_program(); + + // populate build struct members + cuLibraryLoadData(&build_ptr->library, result.data.get(), nullptr, nullptr, 0, nullptr, nullptr, 0); + check(cuLibraryGetKernel(&build_ptr->segmented_sort_fallback_kernel, + build_ptr->library, + segmented_sort_fallback_kernel_lowered_name.c_str())); + check(cuLibraryGetKernel( + &build_ptr->segmented_sort_kernel_small, build_ptr->library, segmented_sort_kernel_small_lowered_name.c_str())); + check(cuLibraryGetKernel( + &build_ptr->segmented_sort_kernel_large, build_ptr->library, segmented_sort_kernel_large_lowered_name.c_str())); + + build_ptr->cc = cc; + build_ptr->cubin = (void*) result.data.release(); + build_ptr->cubin_size = result.size; + // Use the runtime policy extracted via from_json + build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_policy{segmented_sort_policy}; + } + catch (const std::exception& exc) + { + fflush(stderr); + printf("\nEXCEPTION in cccl_device_segmented_sort_build(): %s\n", exc.what()); + fflush(stdout); + error = CUDA_ERROR_UNKNOWN; + } + + return error; +} + +CUresult cccl_device_segmented_sort( + cccl_device_segmented_sort_build_result_t build, + void* d_temp_storage, + size_t* temp_storage_bytes, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_keys_out, + cccl_iterator_t d_values_in, + cccl_iterator_t d_values_out, + uint64_t num_segments, + cccl_iterator_t start_offset_in, + cccl_iterator_t end_offset_in, + CUstream stream) +{ + bool pushed = false; + CUresult error = CUDA_SUCCESS; + try + { + pushed = try_push_context(); + + CUdevice cu_device; + check(cuCtxGetDevice(&cu_device)); + + // Create DoubleBuffer structures for keys and values + // CUB will handle keys-only vs key-value sorting internally + auto d_keys_double_buffer = cub::DoubleBuffer( + static_cast(d_keys_in.state), static_cast(d_keys_out.state)); + auto d_values_double_buffer = cub::DoubleBuffer( + static_cast(d_values_in.state), static_cast(d_values_out.state)); + + auto exec_status = cub::DispatchSegmentedSort< + cub::SortOrder::Ascending, + indirect_arg_t, // KeyT + indirect_arg_t, // ValueT + OffsetT, // OffsetT + indirect_iterator_t, // BeginOffsetIteratorT + indirect_iterator_t, // EndOffsetIteratorT + cub::detail::segmented_sort::policy_hub, // PolicyHub + segmented_sort::segmented_sort_kernel_source, // KernelSource + cub::detail::CudaDriverLauncherFactory>:: // KernelLaunchFactory + Dispatch( + d_temp_storage, + *temp_storage_bytes, + d_keys_double_buffer, + d_values_double_buffer, + 0, // num_items - not used in segmented sort + static_cast(num_segments), + indirect_iterator_t{start_offset_in}, + indirect_iterator_t{end_offset_in}, + true, // is_overwrite_okay + stream, + /* kernel_source */ {build}, + /* launcher_factory */ cub::detail::CudaDriverLauncherFactory{cu_device, build.cc}, + /* policy */ *static_cast(build.runtime_policy)); + + error = static_cast(exec_status); + } + catch (const std::exception& exc) + { + fflush(stderr); + printf("\nEXCEPTION in cccl_device_segmented_sort(): %s\n", exc.what()); + fflush(stdout); + error = CUDA_ERROR_UNKNOWN; + } + + if (pushed) + { + CUcontext dummy; + cuCtxPopCurrent(&dummy); + } + + return error; +} + +CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_result_t* build_ptr) +{ + try + { + if (build_ptr == nullptr) + { + return CUDA_ERROR_INVALID_VALUE; + } + + // allocation behind cubin is owned by unique_ptr with delete[] deleter now + std::unique_ptr cubin(reinterpret_cast(build_ptr->cubin)); + + // Clean up the runtime policy + delete static_cast(build_ptr->runtime_policy); + check(cuLibraryUnload(build_ptr->library)); + } + catch (const std::exception& exc) + { + fflush(stderr); + printf("\nEXCEPTION in cccl_device_segmented_sort_cleanup(): %s\n", exc.what()); + fflush(stdout); + return CUDA_ERROR_UNKNOWN; + } + + return CUDA_SUCCESS; +} diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp new file mode 100644 index 00000000000..16595bb5e3d --- /dev/null +++ b/c/parallel/test/test_segmented_sort.cpp @@ -0,0 +1,596 @@ +//===----------------------------------------------------------------------===// +// +// Part of CUDA Experimental in CUDA C++ Core Libraries, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +#include +#include +#include +#include +#include // std::optional +#include +#include +#include + +#include + +#include "algorithm_execution.h" +#include "build_result_caching.h" +#include "test_util.h" +#include +#include + +using BuildResultT = cccl_device_segmented_sort_build_result_t; + +struct segmented_sort_cleanup +{ + CUresult operator()(BuildResultT* build_data) const noexcept + { + return cccl_device_segmented_sort_cleanup(build_data); + } +}; + +using segmented_sort_deleter = BuildResultDeleter; +using segmented_sort_build_cache_t = build_cache_t>; + +template +auto& get_cache() +{ + return fixture::get_or_create().get_value(); +} + +struct segmented_sort_build +{ + CUresult operator()( + BuildResultT* build_ptr, + cccl_iterator_t keys_in, + cccl_iterator_t keys_out, + cccl_iterator_t values_in, + cccl_iterator_t values_out, + uint64_t, + cccl_iterator_t start_offsets, + cccl_iterator_t end_offsets, + int cc_major, + int cc_minor, + const char* cub_path, + const char* thrust_path, + const char* libcudacxx_path, + const char* ctk_path) const noexcept + { + return cccl_device_segmented_sort_build( + build_ptr, + keys_in, + keys_out, + values_in, + values_out, + start_offsets, + end_offsets, + cc_major, + cc_minor, + cub_path, + thrust_path, + libcudacxx_path, + ctk_path); + } +}; + +struct segmented_sort_run +{ + template + CUresult operator()(Ts... args) const noexcept + { + return cccl_device_segmented_sort(args...); + } +}; + +template +void segmented_sort( + cccl_iterator_t keys_in, + cccl_iterator_t keys_out, + cccl_iterator_t values_in, + cccl_iterator_t values_out, + uint64_t num_segments, + cccl_iterator_t start_offsets, + cccl_iterator_t end_offsets, + std::optional& cache, + const std::optional& lookup_key) +{ + AlgorithmExecute( + cache, lookup_key, keys_in, keys_out, values_in, values_out, num_segments, start_offsets, end_offsets); +} + +// ============== +// Test section +// ============== + +static std::tuple make_step_counting_iterator_sources( + std::string_view index_ty_name, + std::string_view state_name, + std::string_view advance_fn_name, + std::string_view dereference_fn_name) +{ + static constexpr std::string_view it_state_src_tmpl = R"XXX( +struct {0} {{ + {1} linear_id; + {1} row_size; +}}; +)XXX"; + + const std::string it_state_def_src = std::format(it_state_src_tmpl, state_name, index_ty_name); + + static constexpr std::string_view it_def_src_tmpl = R"XXX( +extern "C" __device__ void {0}({1}* state, {2} offset) +{{ + state->linear_id += offset; +}} +)XXX"; + + const std::string it_advance_fn_def_src = + std::format(it_def_src_tmpl, /*0*/ advance_fn_name, state_name, index_ty_name); + + static constexpr std::string_view it_deref_src_tmpl = R"XXX( +extern "C" __device__ {2} {0}({1}* state) +{{ + return (state->linear_id) * (state->row_size); +}} +)XXX"; + + const std::string it_deref_fn_def_src = + std::format(it_deref_src_tmpl, dereference_fn_name, state_name, index_ty_name); + + return std::make_tuple(it_state_def_src, it_advance_fn_def_src, it_deref_fn_def_src); +} + +struct SegmentedSort_KeysOnly_Fixture_Tag; +C2H_TEST_LIST("segmented_sort can sort keys-only with integral types", + "[segmented_sort][keys_only]", + std::int32_t, + std::int64_t, + std::uint32_t, + std::uint64_t) +{ + // generate choices for n_segments: 0, 13 and 2 random samples from [50, 200) + const std::size_t n_segments = GENERATE(0, 13, take(2, random(50, 200))); + // generate choices for segment size: 1, 20 and random samples + const std::size_t segment_size = GENERATE(1, 20, take(2, random(10, 100))); + + const std::size_t n_elems = n_segments * segment_size; + + std::vector host_keys = generate(n_elems); + std::vector host_keys_out(n_elems); + + REQUIRE(host_keys.size() == n_elems); + REQUIRE(host_keys_out.size() == n_elems); + + pointer_t keys_in_ptr(host_keys); // copy from host to device + pointer_t keys_out_ptr(host_keys_out); // copy from host to device + + // Create null value iterators for keys-only sorting + // For keys-only sorting, we create dummy iterators that won't be used + auto dummy_values_it = make_constant_iterator(std::string{"TestType"}); + dummy_values_it.state.value = TestType{}; + cccl_iterator_t values_in = dummy_values_it; + cccl_iterator_t values_out = dummy_values_it; + + using SizeT = unsigned long long; + static constexpr std::string_view index_ty_name = "unsigned long long"; + + struct segment_offset_iterator_state_t + { + SizeT linear_id; + SizeT segment_size; + }; + + static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; + static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; + static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; + + const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = + make_step_counting_iterator_sources( + index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); + + iterator_t start_offset_it = + make_iterator( + {offset_iterator_state_name, offset_iterator_state_src}, + {advance_offset_method_name, offset_iterator_advance_src}, + {deref_offset_method_name, offset_iterator_deref_src}); + + start_offset_it.state.linear_id = 0; + start_offset_it.state.segment_size = segment_size; + + // Create end offset iterator (points to one past start) + iterator_t end_offset_it = + make_iterator( + {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); + + end_offset_it.state.linear_id = 1; + end_offset_it.state.segment_size = segment_size; + + auto& build_cache = get_cache(); + const auto& test_key = make_key(); + + segmented_sort( + keys_in_ptr, keys_out_ptr, values_in, values_out, n_segments, start_offset_it, end_offset_it, build_cache, test_key); + + // Create expected result by sorting each segment + std::vector expected_keys = host_keys; + for (std::size_t i = 0; i < n_segments; ++i) + { + std::size_t segment_start = i * segment_size; + std::size_t segment_end = segment_start + segment_size; + std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end); + } + + REQUIRE(expected_keys == std::vector(keys_out_ptr)); +} + +struct SegmentedSort_KeyValuePairs_Fixture_Tag; +C2H_TEST_LIST("segmented_sort can sort key-value pairs with integral types", + "[segmented_sort][key_value]", + std::int32_t, + std::int64_t, + std::uint32_t, + std::uint64_t) +{ + // generate choices for n_segments: 0, 10 and random samples + const std::size_t n_segments = GENERATE(0, 10, take(2, random(30, 100))); + // generate choices for segment size + const std::size_t segment_size = GENERATE(1, 15, take(2, random(5, 50))); + + const std::size_t n_elems = n_segments * segment_size; + + std::vector host_keys = generate(n_elems); + std::vector host_values = generate(n_elems); + std::vector host_keys_out(n_elems); + std::vector host_values_out(n_elems); + + REQUIRE(host_keys.size() == n_elems); + REQUIRE(host_values.size() == n_elems); + + pointer_t keys_in_ptr(host_keys); + pointer_t keys_out_ptr(host_keys_out); + pointer_t values_in_ptr(host_values); + pointer_t values_out_ptr(host_values_out); + + using SizeT = unsigned long long; + static constexpr std::string_view index_ty_name = "unsigned long long"; + + struct segment_offset_iterator_state_t + { + SizeT linear_id; + SizeT segment_size; + }; + + static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; + static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; + static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; + + const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = + make_step_counting_iterator_sources( + index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); + + iterator_t start_offset_it = + make_iterator( + {offset_iterator_state_name, offset_iterator_state_src}, + {advance_offset_method_name, offset_iterator_advance_src}, + {deref_offset_method_name, offset_iterator_deref_src}); + + start_offset_it.state.linear_id = 0; + start_offset_it.state.segment_size = segment_size; + + iterator_t end_offset_it = + make_iterator( + {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); + + end_offset_it.state.linear_id = 1; + end_offset_it.state.segment_size = segment_size; + + auto& build_cache = get_cache(); + const auto& test_key = make_key(); + + segmented_sort( + keys_in_ptr, + keys_out_ptr, + values_in_ptr, + values_out_ptr, + n_segments, + start_offset_it, + end_offset_it, + build_cache, + test_key); + + // Create expected result by sorting each segment with key-value pairs + std::vector> key_value_pairs; + for (std::size_t i = 0; i < n_elems; ++i) + { + key_value_pairs.emplace_back(host_keys[i], host_values[i]); + } + + std::vector expected_keys(n_elems); + std::vector expected_values(n_elems); + + for (std::size_t i = 0; i < n_segments; ++i) + { + std::size_t segment_start = i * segment_size; + std::size_t segment_end = segment_start + segment_size; + + // Sort this segment by key + std::sort( + key_value_pairs.begin() + segment_start, key_value_pairs.begin() + segment_end, [](const auto& a, const auto& b) { + return a.first < b.first; + }); + + // Extract sorted keys and values + for (std::size_t j = segment_start; j < segment_end; ++j) + { + expected_keys[j] = key_value_pairs[j].first; + expected_values[j] = key_value_pairs[j].second; + } + } + + REQUIRE(expected_keys == std::vector(keys_out_ptr)); + REQUIRE(expected_values == std::vector(values_out_ptr)); +} + +struct custom_pair +{ + int key; + size_t value; + + bool operator==(const custom_pair& other) const + { + return key == other.key && value == other.value; + } + + bool operator<(const custom_pair& other) const + { + return key < other.key; + } +}; + +struct SegmentedSort_CustomTypes_Fixture_Tag; +C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][custom_types]") +{ + using KeyT = custom_pair; + using ValueT = float; + + const std::size_t n_segments = 25; + const std::size_t segment_size = 20; + const std::size_t n_elems = n_segments * segment_size; + + // Generate custom key data + std::vector host_keys(n_elems); + for (std::size_t i = 0; i < n_elems; ++i) + { + host_keys[i] = custom_pair{static_cast(i % 1000), static_cast(i % 100)}; + } + + // Generate float values by first generating ints and then transforming + std::vector host_values_int = generate(n_elems); + std::vector host_values(n_elems); + std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { + return static_cast(x); + }); + std::vector host_keys_out(n_elems); + std::vector host_values_out(n_elems); + + pointer_t keys_in_ptr(host_keys); + pointer_t keys_out_ptr(host_keys_out); + pointer_t values_in_ptr(host_values); + pointer_t values_out_ptr(host_values_out); + + using SizeT = ::cuda::std::size_t; + std::vector segments(n_segments + 1); + for (std::size_t i = 0; i <= n_segments; ++i) + { + segments[i] = i * segment_size; + } + + pointer_t offset_ptr(segments); + + auto start_offset_it = static_cast(offset_ptr); + auto end_offset_it = start_offset_it; + end_offset_it.state = offset_ptr.ptr + 1; + + auto& build_cache = get_cache(); + const auto& test_key = make_key(); + + segmented_sort( + keys_in_ptr, + keys_out_ptr, + values_in_ptr, + values_out_ptr, + n_segments, + start_offset_it, + end_offset_it, + build_cache, + test_key); + + // Create expected result + std::vector> key_value_pairs; + for (std::size_t i = 0; i < n_elems; ++i) + { + key_value_pairs.emplace_back(host_keys[i], host_values[i]); + } + + std::vector expected_keys(n_elems); + std::vector expected_values(n_elems); + + for (std::size_t i = 0; i < n_segments; ++i) + { + std::size_t segment_start = segments[i]; + std::size_t segment_end = segments[i + 1]; + + // Sort this segment by key + std::sort( + key_value_pairs.begin() + segment_start, key_value_pairs.begin() + segment_end, [](const auto& a, const auto& b) { + return a.first < b.first; + }); + + // Extract sorted keys and values + for (std::size_t j = segment_start; j < segment_end; ++j) + { + expected_keys[j] = key_value_pairs[j].first; + expected_values[j] = key_value_pairs[j].second; + } + } + + auto result_keys = std::vector(keys_out_ptr); + auto result_values = std::vector(values_out_ptr); + + REQUIRE(expected_keys == result_keys); + REQUIRE(expected_values == result_values); +} + +using SizeT = unsigned long long; + +struct variable_segment_offset_iterator_state_t +{ + SizeT linear_id; + const SizeT* offsets; +}; + +static std::tuple make_variable_segment_iterator_sources() +{ + static constexpr std::string_view it_state_src = R"XXX( +struct variable_segment_offset_iterator_state_t { + unsigned long long linear_id; + const unsigned long long* offsets; +}; +)XXX"; + + static constexpr std::string_view it_advance_src = R"XXX( +extern "C" __device__ void advance_variable_offset_it(variable_segment_offset_iterator_state_t* state, unsigned long long offset) +{ + state->linear_id += offset; +} +)XXX"; + + static constexpr std::string_view it_deref_src = R"XXX( +extern "C" __device__ unsigned long long dereference_variable_offset_it(variable_segment_offset_iterator_state_t* state) +{ + return state->offsets[state->linear_id]; +} +)XXX"; + + return std::make_tuple(std::string(it_state_src), std::string(it_advance_src), std::string(it_deref_src)); +} + +struct SegmentedSort_VariableSegments_Fixture_Tag; +C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][variable_segments]") +{ + using KeyT = std::int32_t; + using ValueT = float; + + const std::size_t n_segments = 20; + + // Create variable segment sizes + std::vector segment_sizes = {1, 5, 10, 20, 30, 15, 8, 3, 25, 12, 7, 18, 22, 4, 35, 9, 14, 6, 28, 11}; + REQUIRE(segment_sizes.size() == n_segments); + + std::size_t n_elems = std::accumulate(segment_sizes.begin(), segment_sizes.end(), 0ULL); + + std::vector host_keys = generate(n_elems); + // Generate float values by first generating ints and then transforming + std::vector host_values_int = generate(n_elems); + std::vector host_values(n_elems); + std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { + return static_cast(x); + }); + std::vector host_keys_out(n_elems); + std::vector host_values_out(n_elems); + + pointer_t keys_in_ptr(host_keys); + pointer_t keys_out_ptr(host_keys_out); + pointer_t values_in_ptr(host_values); + pointer_t values_out_ptr(host_values_out); + + // Create segment offset arrays + std::vector start_offsets(n_segments); + std::vector end_offsets(n_segments); + + SizeT current_offset = 0; + for (std::size_t i = 0; i < n_segments; ++i) + { + start_offsets[i] = current_offset; + current_offset += segment_sizes[i]; + end_offsets[i] = current_offset; + } + + pointer_t start_offsets_ptr(start_offsets); + pointer_t end_offsets_ptr(end_offsets); + + const auto& [offset_state_src, offset_advance_src, offset_deref_src] = make_variable_segment_iterator_sources(); + + iterator_t start_offset_it = + make_iterator( + {"variable_segment_offset_iterator_state_t", offset_state_src}, + {"advance_variable_offset_it", offset_advance_src}, + {"dereference_variable_offset_it", offset_deref_src}); + + start_offset_it.state.linear_id = 0; + start_offset_it.state.offsets = start_offsets_ptr.ptr; + + iterator_t end_offset_it = + make_iterator( + {"variable_segment_offset_iterator_state_t", ""}, + {"advance_variable_offset_it", ""}, + {"dereference_variable_offset_it", ""}); + + end_offset_it.state.linear_id = 0; + end_offset_it.state.offsets = end_offsets_ptr.ptr; + + auto& build_cache = get_cache(); + const auto& test_key = make_key(); + + segmented_sort( + keys_in_ptr, + keys_out_ptr, + values_in_ptr, + values_out_ptr, + n_segments, + start_offset_it, + end_offset_it, + build_cache, + test_key); + + // Create expected result + std::vector> key_value_pairs; + for (std::size_t i = 0; i < n_elems; ++i) + { + key_value_pairs.emplace_back(host_keys[i], host_values[i]); + } + + std::vector expected_keys(n_elems); + std::vector expected_values(n_elems); + + for (std::size_t i = 0; i < n_segments; ++i) + { + std::size_t segment_start = start_offsets[i]; + std::size_t segment_end = end_offsets[i]; + + // Sort this segment by key + std::sort( + key_value_pairs.begin() + segment_start, key_value_pairs.begin() + segment_end, [](const auto& a, const auto& b) { + return a.first < b.first; + }); + + // Extract sorted keys and values + for (std::size_t j = segment_start; j < segment_end; ++j) + { + expected_keys[j] = key_value_pairs[j].first; + expected_values[j] = key_value_pairs[j].second; + } + } + + auto result_keys = std::vector(keys_out_ptr); + auto result_values = std::vector(values_out_ptr); + + REQUIRE(expected_keys == result_keys); + REQUIRE(expected_values == result_values); +} From 8382203fa2f1c0d251166f4b324ab64881cd931a Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 16:20:52 +0000 Subject: [PATCH 007/100] Add error checks --- c/parallel/src/segmented_sort.cu | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 0c6617bfa35..dcffd7f1889 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -272,22 +272,36 @@ CUresult cccl_device_segmented_sort_build( { CUresult error = CUDA_SUCCESS; + if (keys_in_it.value_type.type != keys_out_it.value_type.type) + { + fflush(stderr); + printf("\nERROR in cccl_device_segmented_sort_build(): keys_in_it and keys_out_it must have the same type\n "); + fflush(stdout); + return CUDA_ERROR_UNKNOWN; + } + + if (values_in_it.value_type.type != values_out_it.value_type.type) + { + fflush(stderr); + printf("\nERROR in cccl_device_segmented_sort_build(): values_in_it and values_out_it must have the same type\n "); + fflush(stdout); + return CUDA_ERROR_UNKNOWN; + } + try { const char* name = "device_segmented_sort"; const int cc = cc_major * 10 + cc_minor; - // Get iterator specializations const auto [keys_in_iterator_name, keys_in_iterator_src] = get_specialization(template_id(), keys_in_it); const auto [keys_out_iterator_name, keys_out_iterator_src] = get_specialization( - template_id(), keys_out_it, keys_in_it.value_type); + template_id(), keys_out_it, keys_out_it.value_type); - // Determine if this is keys-only sorting - const bool keys_only = (values_in_it.state == nullptr || values_out_it.state == nullptr); + const bool keys_only = values_in_it.type == cccl_iterator_kind_t::CCCL_POINTER && values_in_it.state == nullptr; std::string values_in_iterator_name, values_in_iterator_src; std::string values_out_iterator_name, values_out_iterator_src; @@ -306,7 +320,6 @@ CUresult cccl_device_segmented_sort_build( } else { - // For keys-only sorting, use NullType for values values_in_iterator_name = "cub::NullType*"; values_out_iterator_name = "cub::NullType*"; values_in_iterator_src = ""; @@ -320,10 +333,8 @@ CUresult cccl_device_segmented_sort_build( const auto [end_offset_iterator_name, end_offset_iterator_src] = get_specialization(template_id(), end_offset_it); - // OffsetT is checked to match have 64-bit size const auto offset_t = cccl_type_enum_to_name(cccl_type_enum::CCCL_UINT64); - // Get type names for keys and values const std::string key_t = cccl_type_enum_to_name(keys_in_it.value_type.type); const std::string value_t = keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); From 233a73a97b5d04ceee44c5fe7437ffaaadc87b34 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 16:35:35 +0000 Subject: [PATCH 008/100] Make segment selector operator() device only (for c.parallel) and revert change to continuation kernel --- .../dispatch/dispatch_segmented_sort.cuh | 73 +++++++++++-------- 1 file changed, 42 insertions(+), 31 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index abebf3450e2..ffb0d390a86 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -319,7 +319,7 @@ struct DispatchSegmentedSort , d_offset_end(d_offset_end) {} - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const + _CCCL_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const { const OffsetT segment_size = d_offset_end[base_segment_offset + segment_id] - d_offset_begin[base_segment_offset + segment_id]; @@ -341,7 +341,7 @@ struct DispatchSegmentedSort , d_offset_end(d_offset_end) {} - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const + _CCCL_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const { const OffsetT segment_size = d_offset_end[base_segment_offset + segment_id] - d_offset_begin[base_segment_offset + segment_id]; @@ -769,35 +769,46 @@ private: #else // CUB_RDC_ENABLED -# define CUB_TEMP_DEVICE_CODE \ - error = \ - launcher_factory(1, 1, 0, stream) \ - .doit(kernel_source.SegmentedSortContinuationKernel(), \ - large_kernel, \ - small_kernel, \ - current_num_segments, \ - d_keys.Current(), \ - GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), \ - d_keys_double_buffer, \ - d_values.Current(), \ - GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), \ - d_values_double_buffer, \ - current_begin_offset, \ - current_end_offset, \ - group_sizes.get(), \ - large_and_medium_segments_indices.get(), \ - small_segments_indices.get()); \ - error = CubDebug(error); \ - \ - if (cudaSuccess != error) \ - { \ - return error; \ - } \ - \ - error = CubDebug(detail::DebugSyncStream(stream)); \ - if (cudaSuccess != error) \ - { \ - return error; \ +# define CUB_TEMP_DEVICE_CODE \ + error = \ + launcher_factory(1, 1, 0, stream) \ + .doit( \ + detail::segmented_sort::DeviceSegmentedSortContinuationKernel< \ + WrappedPolicyT, \ + LargeKernelT, \ + SmallKernelT, \ + KeyT, \ + ValueT, \ + BeginOffsetIteratorT, \ + EndOffsetIteratorT, \ + KernelLauncherFactory>, \ + large_kernel, \ + small_kernel, \ + current_num_segments, \ + d_keys.Current(), \ + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), \ + d_keys_double_buffer, \ + d_values.Current(), \ + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), \ + d_values_double_buffer, \ + current_begin_offset, \ + current_end_offset, \ + group_sizes.get(), \ + large_and_medium_segments_indices.get(), \ + small_segments_indices.get(), \ + launcher_factory, \ + wrapped_policy); \ + error = CubDebug(error); \ + \ + if (cudaSuccess != error) \ + { \ + return error; \ + } \ + \ + error = CubDebug(detail::DebugSyncStream(stream)); \ + if (cudaSuccess != error) \ + { \ + return error; \ } #endif // CUB_RDC_ENABLED From 2f579abc6380732b3f7c1cc27a84bf3c43a32f70 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 16:55:28 +0000 Subject: [PATCH 009/100] move three way partition kernels to separate header since they are used in dispatch segmented sort --- .../dispatch/dispatch_three_way_partition.cuh | 101 +----------------- .../dispatch/kernels/segmented_sort.cuh | 50 +++++++++ 2 files changed, 51 insertions(+), 100 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index 8cb7e02d091..00590054938 100644 --- a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -37,10 +37,9 @@ # pragma system_header #endif // no system header -#include #include +#include #include -#include #include #include @@ -128,104 +127,6 @@ public: } } }; - -/****************************************************************************** - * Kernel entry points - *****************************************************************************/ -template -__launch_bounds__(int(ChainedPolicyT::ActivePolicy::ThreeWayPartitionPolicy::BLOCK_THREADS)) - CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceThreeWayPartitionKernel( - InputIteratorT d_in, - FirstOutputIteratorT d_first_part_out, - SecondOutputIteratorT d_second_part_out, - UnselectedOutputIteratorT d_unselected_out, - NumSelectedIteratorT d_num_selected_out, - ScanTileStateT tile_status, - SelectFirstPartOp select_first_part_op, - SelectSecondPartOp select_second_part_op, - OffsetT num_items, - int num_tiles, - _CCCL_GRID_CONSTANT const StreamingContextT streaming_context) -{ - using AgentThreeWayPartitionPolicyT = typename ChainedPolicyT::ActivePolicy::ThreeWayPartitionPolicy; - - // Thread block type for selecting data from input tiles - using AgentThreeWayPartitionT = AgentThreeWayPartition< - AgentThreeWayPartitionPolicyT, - InputIteratorT, - FirstOutputIteratorT, - SecondOutputIteratorT, - UnselectedOutputIteratorT, - SelectFirstPartOp, - SelectSecondPartOp, - OffsetT, - StreamingContextT>; - - // Shared memory for AgentThreeWayPartition - __shared__ typename AgentThreeWayPartitionT::TempStorage temp_storage; - - // Process tiles - AgentThreeWayPartitionT( - temp_storage, - d_in, - d_first_part_out, - d_second_part_out, - d_unselected_out, - select_first_part_op, - select_second_part_op, - num_items, - streaming_context) - .ConsumeRange(num_tiles, tile_status, d_num_selected_out); -} - -/** - * @brief Initialization kernel for tile status initialization (multi-block) - * - * @tparam ScanTileStateT - * Tile status interface type - * - * @tparam NumSelectedIteratorT - * Output iterator type for recording the number of items selected - * - * @param[in] tile_state_1 - * Tile status interface - * - * @param[in] tile_state_2 - * Tile status interface - * - * @param[in] num_tiles - * Number of tiles - * - * @param[out] d_num_selected_out - * Pointer to the total number of items selected - * (i.e., length of @p d_selected_out) - */ -template -CUB_DETAIL_KERNEL_ATTRIBUTES void -DeviceThreeWayPartitionInitKernel(ScanTileStateT tile_state, int num_tiles, NumSelectedIteratorT d_num_selected_out) -{ - // Initialize tile status - tile_state.InitializeStatus(num_tiles); - - // Initialize d_num_selected_out - if (blockIdx.x == 0) - { - if (threadIdx.x < 2) - { - d_num_selected_out[threadIdx.x] = 0; - } - } -} } // namespace detail::three_way_partition /****************************************************************************** diff --git a/cub/cub/device/dispatch/kernels/segmented_sort.cuh b/cub/cub/device/dispatch/kernels/segmented_sort.cuh index 35436269875..2d1c2a1fa18 100644 --- a/cub/cub/device/dispatch/kernels/segmented_sort.cuh +++ b/cub/cub/device/dispatch/kernels/segmented_sort.cuh @@ -29,6 +29,56 @@ using local_segment_index_t = ::cuda::std::uint32_t; // Type used for total number of segments and to index within segments globally using global_segment_offset_t = ::cuda::std::int64_t; +template +struct LargeSegmentsSelectorT +{ + OffsetT value{}; + BeginOffsetIteratorT d_offset_begin{}; + EndOffsetIteratorT d_offset_end{}; + global_segment_offset_t base_segment_offset{}; + +#if !_CCCL_COMPILER(NVRTC) + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE + LargeSegmentsSelectorT(OffsetT value, BeginOffsetIteratorT d_offset_begin, EndOffsetIteratorT d_offset_end) + : value(value) + , d_offset_begin(d_offset_begin) + , d_offset_end(d_offset_end) + {} +#endif + + _CCCL_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const + { + const OffsetT segment_size = + d_offset_end[base_segment_offset + segment_id] - d_offset_begin[base_segment_offset + segment_id]; + return segment_size > value; + } +}; + +template +struct SmallSegmentsSelectorT +{ + OffsetT value{}; + BeginOffsetIteratorT d_offset_begin{}; + EndOffsetIteratorT d_offset_end{}; + global_segment_offset_t base_segment_offset{}; + +#if !_CCCL_COMPILER(NVRTC) + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE + SmallSegmentsSelectorT(OffsetT value, BeginOffsetIteratorT d_offset_begin, EndOffsetIteratorT d_offset_end) + : value(value) + , d_offset_begin(d_offset_begin) + , d_offset_end(d_offset_end) + {} +#endif + + _CCCL_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const + { + const OffsetT segment_size = + d_offset_end[base_segment_offset + segment_id] - d_offset_begin[base_segment_offset + segment_id]; + return segment_size < value; + } +}; + /** * @brief Fallback kernel, in case there's not enough segments to * take advantage of partitioning. From d76073707a2afd6444747817d20db01cc0d7efd5 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 17:48:29 +0000 Subject: [PATCH 010/100] Enable dynamic cub dispatch in three way partition --- .../dispatch/dispatch_three_way_partition.cuh | 218 ++++++++++-------- .../tuning/tuning_three_way_partition.cuh | 26 +++ 2 files changed, 143 insertions(+), 101 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index 00590054938..3be8a562c49 100644 --- a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -37,6 +37,7 @@ # pragma system_header #endif // no system header +#include #include #include #include @@ -45,7 +46,8 @@ #include -#include +#include +#include #include @@ -53,10 +55,53 @@ CUB_NAMESPACE_BEGIN namespace detail::three_way_partition { + +template +struct DeviceThreeWayPartitionKernelSource +{ + CUB_DEFINE_KERNEL_GETTER(ThreeWayPartitionInitKernel, + DeviceThreeWayPartitionInitKernel); + + CUB_DEFINE_KERNEL_GETTER( + ThreeWayPartitionKernel, + DeviceThreeWayPartitionKernel< + MaxPolicyT, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + per_partition_offset_t, + streaming_context_t>); + + CUB_RUNTIME_FUNCTION static constexpr size_t OffsetSize() + { + return sizeof(OffsetT); + } +}; + // Offset type used to instantiate the stream three-way-partition-kernel and agent to index the items within one // partition using per_partition_offset_t = ::cuda::std::int32_t; +using AccumPackHelperT = detail::three_way_partition::accumulator_pack_t; +using AccumPackT = typename AccumPackHelperT::pack_t; +using ScanTileStateT = cub::ScanTileState; + template class streaming_context_t { @@ -133,16 +178,31 @@ public: * Dispatch ******************************************************************************/ -template , detail::three_way_partition::per_partition_offset_t>> +template < + typename InputIteratorT, + typename FirstOutputIteratorT, + typename SecondOutputIteratorT, + typename UnselectedOutputIteratorT, + typename NumSelectedIteratorT, + typename SelectFirstPartOp, + typename SelectSecondPartOp, + typename OffsetT, + typename PolicyHub = detail::three_way_partition::policy_hub, + detail::three_way_partition::per_partition_offset_t>, + typename KernelSource = detail::three_way_partition::DeviceThreeWayPartitionKernelSource< + typename PolicyHub::MaxPolicy, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + detail::three_way_partition::ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + detail::three_way_partition::per_partition_offset_t, + detail::three_way_partition::streaming_context_t, + OffsetT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> struct DispatchThreeWayPartitionIf { /***************************************************************************** @@ -157,9 +217,7 @@ struct DispatchThreeWayPartitionIf using streaming_context_t = detail::three_way_partition::streaming_context_t; - using AccumPackHelperT = detail::three_way_partition::accumulator_pack_t; - using AccumPackT = typename AccumPackHelperT::pack_t; - using ScanTileStateT = cub::ScanTileState; + using ScanTileStateT = detail::three_way_partition::ScanTileStateT; static constexpr int INIT_KERNEL_THREADS = 256; @@ -174,31 +232,8 @@ struct DispatchThreeWayPartitionIf SelectSecondPartOp select_second_part_op; OffsetT num_items; cudaStream_t stream; - - CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchThreeWayPartitionIf( - void* d_temp_storage, - size_t& temp_storage_bytes, - InputIteratorT d_in, - FirstOutputIteratorT d_first_part_out, - SecondOutputIteratorT d_second_part_out, - UnselectedOutputIteratorT d_unselected_out, - NumSelectedIteratorT d_num_selected_out, - SelectFirstPartOp select_first_part_op, - SelectSecondPartOp select_second_part_op, - OffsetT num_items, - cudaStream_t stream) - : d_temp_storage(d_temp_storage) - , temp_storage_bytes(temp_storage_bytes) - , d_in(d_in) - , d_first_part_out(d_first_part_out) - , d_second_part_out(d_second_part_out) - , d_unselected_out(d_unselected_out) - , d_num_selected_out(d_num_selected_out) - , select_first_part_op(select_first_part_op) - , select_second_part_op(select_second_part_op) - , num_items(num_items) - , stream(stream) - {} + KernelSource kernel_source; + KernelLauncherFactory launcher_factory; /***************************************************************************** * Dispatch entrypoints @@ -206,13 +241,15 @@ struct DispatchThreeWayPartitionIf template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t - Invoke(ScanInitKernelPtrT three_way_partition_init_kernel, SelectIfKernelPtrT three_way_partition_kernel) + Invoke(ActivePolicyT policy, + ScanInitKernelPtrT three_way_partition_init_kernel, + SelectIfKernelPtrT three_way_partition_kernel) { cudaError error = cudaSuccess; - constexpr int block_threads = ActivePolicyT::ThreeWayPartitionPolicy::BLOCK_THREADS; - constexpr int items_per_thread = ActivePolicyT::ThreeWayPartitionPolicy::ITEMS_PER_THREAD; - constexpr int tile_size = block_threads * items_per_thread; + const int block_threads = policy.ThreeWayPartition().BlockThreads(); + const int items_per_thread = policy.ThreeWayPartition().ItemsPerThread(); + const int tile_size = block_threads * items_per_thread; // The maximum number of items for which we will ever invoke the kernel (i.e. largest partition size) auto const max_partition_size = static_cast( @@ -229,7 +266,7 @@ struct DispatchThreeWayPartitionIf constexpr ::cuda::std::size_t num_counters_per_pass = 3; constexpr ::cuda::std::size_t num_streaming_counters = 2 * num_counters_per_pass; ::cuda::std::size_t streaming_selection_storage_bytes = - (num_partitions > 1) ? num_streaming_counters * sizeof(OffsetT) : ::cuda::std::size_t{0}; + (num_partitions > 1) ? num_streaming_counters * kernel_source.OffsetSize() : ::cuda::std::size_t{0}; // Specify temporary storage allocation requirements size_t allocation_sizes[2] = {0ULL, streaming_selection_storage_bytes}; @@ -282,7 +319,7 @@ struct DispatchThreeWayPartitionIf } // Log three_way_partition_init_kernel configuration - int init_grid_size = _CUDA_VSTD::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); + const int init_grid_size = _CUDA_VSTD::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); #ifdef CUB_DEBUG_LOG _CubLog("Invoking three_way_partition_init_kernel<<<%d, %d, 0, %lld>>>()\n", @@ -292,7 +329,7 @@ struct DispatchThreeWayPartitionIf #endif // CUB_DEBUG_LOG // Invoke three_way_partition_init_kernel to initialize tile descriptors - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream) .doit(three_way_partition_init_kernel, tile_status, current_num_tiles, d_num_selected_out); // Check for failure to launch @@ -321,9 +358,10 @@ struct DispatchThreeWayPartitionIf { // Get SM occupancy for select_if_kernel int range_select_sm_occupancy; - error = CubDebug(MaxSmOccupancy(range_select_sm_occupancy, // out - three_way_partition_kernel, - block_threads)); + error = CubDebug(launcher_factory.MaxSmOccupancy( + range_select_sm_occupancy, // out + three_way_partition_kernel, + block_threads)); if (cudaSuccess != error) { return error; @@ -340,7 +378,7 @@ struct DispatchThreeWayPartitionIf #endif // CUB_DEBUG_LOG // Invoke select_if_kernel - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(current_num_tiles, block_threads, 0, stream) + launcher_factory(current_num_tiles, block_threads, 0, stream) .doit(three_way_partition_kernel, d_in, d_first_part_out, @@ -376,28 +414,16 @@ struct DispatchThreeWayPartitionIf } template - CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke() + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) { - using MaxPolicyT = typename PolicyHub::MaxPolicy; - return Invoke( - detail::three_way_partition::DeviceThreeWayPartitionInitKernel, - detail::three_way_partition::DeviceThreeWayPartitionKernel< - MaxPolicyT, - InputIteratorT, - FirstOutputIteratorT, - SecondOutputIteratorT, - UnselectedOutputIteratorT, - NumSelectedIteratorT, - ScanTileStateT, - SelectFirstPartOp, - SelectSecondPartOp, - per_partition_offset_t, - streaming_context_t>); + const auto wrapped_policy = detail::three_way_partition::MakeThreeWayPartitionPolicyWrapper(active_policy); + return Invoke(wrapped_policy, kernel_source.ThreeWayPartitionInitKernel(), kernel_source.ThreeWayPartitionKernel()); } /** * Internal dispatch routine */ + template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -409,44 +435,34 @@ struct DispatchThreeWayPartitionIf SelectFirstPartOp select_first_part_op, SelectSecondPartOp select_second_part_op, OffsetT num_items, - cudaStream_t stream) + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) { - using MaxPolicyT = typename PolicyHub::MaxPolicy; - - cudaError error = cudaSuccess; - - do + // Get PTX version + int ptx_version = 0; + if (cudaError error = CubDebug(launcher_factory.PtxVersion(ptx_version)); cudaSuccess != error) { - // Get PTX version - int ptx_version = 0; - error = CubDebug(cub::PtxVersion(ptx_version)); - if (cudaSuccess != error) - { - break; - } - - DispatchThreeWayPartitionIf dispatch( - d_temp_storage, - temp_storage_bytes, - d_in, - d_first_part_out, - d_second_part_out, - d_unselected_out, - d_num_selected_out, - select_first_part_op, - select_second_part_op, - num_items, - stream); - - // Dispatch - error = CubDebug(MaxPolicyT::Invoke(ptx_version, dispatch)); - if (cudaSuccess != error) - { - break; - } - } while (0); + return error; + } - return error; + DispatchThreeWayPartitionIf dispatch{ + d_temp_storage, + temp_storage_bytes, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + d_num_selected_out, + select_first_part_op, + select_second_part_op, + num_items, + stream, + kernel_source, + launcher_factory}; + + return CubDebug(max_policy.Invoke(ptx_version, dispatch)); } }; diff --git a/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh b/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh index cf2d3740f54..8bb69b9fbaa 100644 --- a/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh @@ -51,6 +51,32 @@ namespace detail { namespace three_way_partition { + +template +struct ThreeWayPartitionPolicyWrapper : PolicyT +{ + CUB_RUNTIME_FUNCTION ThreeWayPartitionPolicyWrapper(PolicyT base) + : PolicyT(base) + {} +}; + +template +struct ThreeWayPartitionPolicyWrapper> + : StaticPolicyT +{ + CUB_RUNTIME_FUNCTION ThreeWayPartitionPolicyWrapper(StaticPolicyT base) + : StaticPolicyT(base) + {} + + CUB_DEFINE_SUB_POLICY_GETTER(ThreeWayPartition) +}; + +template +CUB_RUNTIME_FUNCTION ThreeWayPartitionPolicyWrapper MakeThreeWayPartitionPolicyWrapper(PolicyT policy) +{ + return ThreeWayPartitionPolicyWrapper{policy}; +} + enum class input_size { _1, From 93d5fdab3feb638f3128204208d8199bd9e8a848 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 18:27:33 +0000 Subject: [PATCH 011/100] Move three way partition kernels to separate file --- .../dispatch/kernels/three_way_partition.cuh | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) create mode 100644 cub/cub/device/dispatch/kernels/three_way_partition.cuh diff --git a/cub/cub/device/dispatch/kernels/three_way_partition.cuh b/cub/cub/device/dispatch/kernels/three_way_partition.cuh new file mode 100644 index 00000000000..9ea919f7110 --- /dev/null +++ b/cub/cub/device/dispatch/kernels/three_way_partition.cuh @@ -0,0 +1,121 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#pragma once + +#include + +#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 + +CUB_NAMESPACE_BEGIN + +namespace detail::three_way_partition +{ +/****************************************************************************** + * Kernel entry points + *****************************************************************************/ +template +__launch_bounds__(int(ChainedPolicyT::ActivePolicy::ThreeWayPartitionPolicy::BLOCK_THREADS)) + CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceThreeWayPartitionKernel( + InputIteratorT d_in, + FirstOutputIteratorT d_first_part_out, + SecondOutputIteratorT d_second_part_out, + UnselectedOutputIteratorT d_unselected_out, + NumSelectedIteratorT d_num_selected_out, + ScanTileStateT tile_status, + SelectFirstPartOp select_first_part_op, + SelectSecondPartOp select_second_part_op, + OffsetT num_items, + int num_tiles, + _CCCL_GRID_CONSTANT const StreamingContextT streaming_context) +{ + using AgentThreeWayPartitionPolicyT = typename ChainedPolicyT::ActivePolicy::ThreeWayPartitionPolicy; + + // Thread block type for selecting data from input tiles + using AgentThreeWayPartitionT = AgentThreeWayPartition< + AgentThreeWayPartitionPolicyT, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + SelectFirstPartOp, + SelectSecondPartOp, + OffsetT, + StreamingContextT>; + + // Shared memory for AgentThreeWayPartition + __shared__ typename AgentThreeWayPartitionT::TempStorage temp_storage; + + // Process tiles + AgentThreeWayPartitionT( + temp_storage, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + select_first_part_op, + select_second_part_op, + num_items, + streaming_context) + .ConsumeRange(num_tiles, tile_status, d_num_selected_out); +} + +/** + * @brief Initialization kernel for tile status initialization (multi-block) + * + * @tparam ScanTileStateT + * Tile status interface type + * + * @tparam NumSelectedIteratorT + * Output iterator type for recording the number of items selected + * + * @param[in] tile_state_1 + * Tile status interface + * + * @param[in] tile_state_2 + * Tile status interface + * + * @param[in] num_tiles + * Number of tiles + * + * @param[out] d_num_selected_out + * Pointer to the total number of items selected + * (i.e., length of @p d_selected_out) + */ +template +CUB_DETAIL_KERNEL_ATTRIBUTES void +DeviceThreeWayPartitionInitKernel(ScanTileStateT tile_state, int num_tiles, NumSelectedIteratorT d_num_selected_out) +{ + // Initialize tile status + tile_state.InitializeStatus(num_tiles); + + // Initialize d_num_selected_out + if (blockIdx.x == 0) + { + if (threadIdx.x < 2) + { + d_num_selected_out[threadIdx.x] = 0; + } + } +} +} // namespace detail::three_way_partition + +CUB_NAMESPACE_END From 48b09969a8748fbb0fb669254a9ff3034fd9a58e Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 19:11:23 +0000 Subject: [PATCH 012/100] Call partition through dispatch and add template params for partition policy hub and kernel source --- cub/cub/device/device_partition.cuh | 4 +- .../dispatch/dispatch_segmented_sort.cuh | 125 +++++++++--------- 2 files changed, 67 insertions(+), 62 deletions(-) diff --git a/cub/cub/device/device_partition.cuh b/cub/cub/device/device_partition.cuh index 259b2d40d63..472bed4ef6e 100644 --- a/cub/cub/device/device_partition.cuh +++ b/cub/cub/device/device_partition.cuh @@ -384,7 +384,9 @@ private: typename EndOffsetIteratorT, typename PolicyHub, typename KernelSource, - typename KernelLauncherFactory> + typename KernelLauncherFactory, + typename PartitionPolicyHub, + typename PartitionKernelSource> friend class DispatchSegmentedSort; // Internal version without NVTX range diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index ffb0d390a86..71bc9b44787 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -282,22 +282,39 @@ struct DeviceSegmentedSortKernelSource }; } // namespace detail::segmented_sort -template , - typename KernelSource = detail::segmented_sort::DeviceSegmentedSortKernelSource< - typename PolicyHub::MaxPolicy, - Order, - KeyT, - ValueT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT>, - typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +template < + SortOrder Order, + typename KeyT, + typename ValueT, + typename OffsetT, + typename BeginOffsetIteratorT, + typename EndOffsetIteratorT, + typename PolicyHub = detail::segmented_sort::policy_hub, + typename KernelSource = detail::segmented_sort::DeviceSegmentedSortKernelSource< + typename PolicyHub::MaxPolicy, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT>, + typename PartitionPolicyHub = detail::three_way_partition::policy_hub< + cub::detail::it_value_t>, + detail::three_way_partition::per_partition_offset_t>, + typename PartitionKernelSource = detail::three_way_partition::DeviceThreeWayPartitionKernelSource< + typename PartitionPolicyHub::MaxPolicy, + THRUST_NS_QUALIFIER::counting_iterator, + cub::detail::segmented_sort::local_segment_index_t*, + cub::detail::segmented_sort::local_segment_index_t*, + THRUST_NS_QUALIFIER::reverse_iterator, + cub::detail::segmented_sort::local_segment_index_t*, + detail::three_way_partition::ScanTileStateT, + cub::detail::segmented_sort::LargeSegmentsSelectorT, + cub::detail::segmented_sort::SmallSegmentsSelectorT, + detail::three_way_partition::per_partition_offset_t, + detail::three_way_partition::streaming_context_t, + detail::choose_signed_offset::type>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> struct DispatchSegmentedSort { using local_segment_index_t = detail::segmented_sort::local_segment_index_t; @@ -305,49 +322,10 @@ struct DispatchSegmentedSort static constexpr int KEYS_ONLY = ::cuda::std::is_same_v; - struct LargeSegmentsSelectorT - { - OffsetT value{}; - BeginOffsetIteratorT d_offset_begin{}; - EndOffsetIteratorT d_offset_end{}; - global_segment_offset_t base_segment_offset{}; - - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE - LargeSegmentsSelectorT(OffsetT value, BeginOffsetIteratorT d_offset_begin, EndOffsetIteratorT d_offset_end) - : value(value) - , d_offset_begin(d_offset_begin) - , d_offset_end(d_offset_end) - {} - - _CCCL_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const - { - const OffsetT segment_size = - d_offset_end[base_segment_offset + segment_id] - d_offset_begin[base_segment_offset + segment_id]; - return segment_size > value; - } - }; - - struct SmallSegmentsSelectorT - { - OffsetT value{}; - BeginOffsetIteratorT d_offset_begin{}; - EndOffsetIteratorT d_offset_end{}; - global_segment_offset_t base_segment_offset{}; - - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE - SmallSegmentsSelectorT(OffsetT value, BeginOffsetIteratorT d_offset_begin, EndOffsetIteratorT d_offset_end) - : value(value) - , d_offset_begin(d_offset_begin) - , d_offset_end(d_offset_end) - {} - - _CCCL_DEVICE _CCCL_FORCEINLINE bool operator()(local_segment_index_t segment_id) const - { - const OffsetT segment_size = - d_offset_end[base_segment_offset + segment_id] - d_offset_begin[base_segment_offset + segment_id]; - return segment_size < value; - } - }; + using LargeSegmentsSelectorT = + cub::detail::segmented_sort::LargeSegmentsSelectorT; + using SmallSegmentsSelectorT = + cub::detail::segmented_sort::SmallSegmentsSelectorT; // Partition selects large and small groups. The middle group is not selected. static constexpr size_t num_selected_groups = 2; @@ -497,7 +475,32 @@ struct DispatchSegmentedSort auto medium_indices_iterator = THRUST_NS_QUALIFIER::make_reverse_iterator(large_and_medium_segments_indices.get()); - cub::DevicePartition::IfNoNVTX( + // We call partition through dispatch instead of device because c.parallel needs to be able to call the kernel. + // This approach propagates the type erasure to partition. + using ChooseOffsetT = detail::choose_signed_offset; + using PartitionOffsetT = typename ChooseOffsetT::type; + using DispatchThreeWayPartitionIfT = cub::DispatchThreeWayPartitionIf< + THRUST_NS_QUALIFIER::counting_iterator, + decltype(large_and_medium_segments_indices.get()), + decltype(small_segments_indices.get()), + decltype(medium_indices_iterator), + decltype(group_sizes.get()), + LargeSegmentsSelectorT, + SmallSegmentsSelectorT, + PartitionOffsetT, + PartitionPolicyHub, + PartitionKernelSource, + KernelLauncherFactory>; + + // Signed integer type for global offsets + // Check if the number of items exceeds the range covered by the selected signed offset type + cudaError_t error = ChooseOffsetT::is_exceeding_offset_type(num_items); + if (error) + { + return error; + } + + DispatchThreeWayPartitionIfT::Dispatch( nullptr, three_way_partition_temp_storage_bytes, THRUST_NS_QUALIFIER::counting_iterator(0), @@ -505,9 +508,9 @@ struct DispatchSegmentedSort small_segments_indices.get(), medium_indices_iterator, group_sizes.get(), - max_num_segments_per_invocation, large_segments_selector, small_segments_selector, + max_num_segments_per_invocation, stream); device_partition_temp_storage.grow(three_way_partition_temp_storage_bytes); From cce3899a34aeb0c4b36cf12c0ea64037417d00c0 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 20:12:13 +0000 Subject: [PATCH 013/100] Various compilation fixes --- c/parallel/include/cccl/c/segmented_sort.h | 7 +- c/parallel/src/segmented_sort.cu | 86 +++++++++++++++++++--- c/parallel/test/test_segmented_sort.cpp | 22 +++++- 3 files changed, 100 insertions(+), 15 deletions(-) diff --git a/c/parallel/include/cccl/c/segmented_sort.h b/c/parallel/include/cccl/c/segmented_sort.h index 29a9fbfae9b..a4efe792a14 100644 --- a/c/parallel/include/cccl/c/segmented_sort.h +++ b/c/parallel/include/cccl/c/segmented_sort.h @@ -28,10 +28,14 @@ typedef struct cccl_device_segmented_sort_build_result_t void* cubin; size_t cubin_size; CUlibrary library; + cccl_type_info offset_type; CUkernel segmented_sort_fallback_kernel; CUkernel segmented_sort_kernel_small; CUkernel segmented_sort_kernel_large; + CUkernel three_way_partition_init_kernel; + CUkernel three_way_partition_kernel; void* runtime_policy; + void* partition_runtime_policy; } cccl_device_segmented_sort_build_result_t; // TODO return a union of nvtx/cuda/nvrtc errors or a string? @@ -58,7 +62,8 @@ CCCL_C_API CUresult cccl_device_segmented_sort( cccl_iterator_t d_keys_out, cccl_iterator_t d_values_in, cccl_iterator_t d_values_out, - uint64_t num_segments, + int64_t num_items, + int64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, CUstream stream); diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index dcffd7f1889..ff2df1cd707 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -38,12 +38,12 @@ #include struct device_segmented_sort_policy; -using OffsetT = unsigned long long; -static_assert(std::is_same_v, OffsetT>, "OffsetT must be size_t"); +using OffsetT = long; +static_assert(std::is_same_v, OffsetT>, "OffsetT must be long"); -// check we can map OffsetT to ::cuda::std::uint64_t -static_assert(std::is_unsigned_v); -static_assert(sizeof(OffsetT) == sizeof(::cuda::std::uint64_t)); +// check we can map OffsetT to ::cuda::std::int64_t +static_assert(std::is_signed_v); +static_assert(sizeof(OffsetT) == sizeof(::cuda::std::int64_t)); namespace segmented_sort { @@ -246,6 +246,66 @@ struct segmented_sort_kernel_source return build.segmented_sort_kernel_large; } }; + +struct partition_kernel_source +{ + cccl_device_segmented_sort_build_result_t& build; + + CUkernel ThreeWayPartitionInitKernel() const + { + return build.three_way_partition_init_kernel; + } + CUkernel ThreeWayPartitionKernel() const + { + return build.three_way_partition_kernel; + } + + std::size_t OffsetSize() const + { + return build.offset_type.size; + } +}; + +struct segmented_sort_runtime_tuning_policy +{ + cub::detail::RuntimeRadixSortDownsweepAgentPolicy large_segment; + cub::detail::RuntimeSmallAndMediumSegmentedSortAgentPolicy small_and_medium_segment; + + auto LargeSegment() const + { + return large_segment; + } + auto SmallAndMediumSegmentedSort() const + { + return small_and_medium_segment; + } + + using MaxPolicy = segmented_sort_runtime_tuning_policy; + + template + cudaError_t Invoke(int, F& op) + { + return op.template Invoke(*this); + } +}; + +struct partition_runtime_tuning_policy +{ + cub::detail::RuntimeThreeWayPartitionAgentPolicy three_way_partition; + + auto ThreeWayPartition() const + { + return three_way_partition; + } + + using MaxPolicy = partition_runtime_tuning_policy; + + template + cudaError_t Invoke(int, F& op) + { + return op.template Invoke(*this); + } +}; } // namespace segmented_sort struct segmented_sort_keys_input_iterator_tag; @@ -554,7 +614,8 @@ CUresult cccl_device_segmented_sort( cccl_iterator_t d_keys_out, cccl_iterator_t d_values_in, cccl_iterator_t d_values_out, - uint64_t num_segments, + int64_t num_items, + int64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, CUstream stream) @@ -582,23 +643,28 @@ CUresult cccl_device_segmented_sort( OffsetT, // OffsetT indirect_iterator_t, // BeginOffsetIteratorT indirect_iterator_t, // EndOffsetIteratorT - cub::detail::segmented_sort::policy_hub, // PolicyHub + segmented_sort::segmented_sort_runtime_tuning_policy, // PolicyHub segmented_sort::segmented_sort_kernel_source, // KernelSource + segmented_sort::partition_runtime_tuning_policy, // PartitionPolicyHub + segmented_sort::partition_kernel_source, // PartitionKernelSource cub::detail::CudaDriverLauncherFactory>:: // KernelLaunchFactory Dispatch( d_temp_storage, *temp_storage_bytes, d_keys_double_buffer, d_values_double_buffer, - 0, // num_items - not used in segmented sort - static_cast(num_segments), + num_items, + num_segments, indirect_iterator_t{start_offset_in}, indirect_iterator_t{end_offset_in}, true, // is_overwrite_okay stream, /* kernel_source */ {build}, + /* partition_kernel_source */ {build}, /* launcher_factory */ cub::detail::CudaDriverLauncherFactory{cu_device, build.cc}, - /* policy */ *static_cast(build.runtime_policy)); + /* policy */ *reinterpret_cast(build.runtime_policy), + /* partition_policy */ + *reinterpret_cast(build.partition_runtime_policy)); error = static_cast(exec_status); } diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 16595bb5e3d..d52f8a36ee1 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -52,7 +52,8 @@ struct segmented_sort_build cccl_iterator_t keys_out, cccl_iterator_t values_in, cccl_iterator_t values_out, - uint64_t, + int64_t /*num_items*/, + int64_t /*num_segments*/, cccl_iterator_t start_offsets, cccl_iterator_t end_offsets, int cc_major, @@ -94,14 +95,15 @@ void segmented_sort( cccl_iterator_t keys_out, cccl_iterator_t values_in, cccl_iterator_t values_out, - uint64_t num_segments, + int64_t num_items, + int64_t num_segments, cccl_iterator_t start_offsets, cccl_iterator_t end_offsets, std::optional& cache, const std::optional& lookup_key) { AlgorithmExecute( - cache, lookup_key, keys_in, keys_out, values_in, values_out, num_segments, start_offsets, end_offsets); + cache, lookup_key, keys_in, keys_out, values_in, values_out, num_items, num_segments, start_offsets, end_offsets); } // ============== @@ -215,7 +217,16 @@ C2H_TEST_LIST("segmented_sort can sort keys-only with integral types", const auto& test_key = make_key(); segmented_sort( - keys_in_ptr, keys_out_ptr, values_in, values_out, n_segments, start_offset_it, end_offset_it, build_cache, test_key); + keys_in_ptr, + keys_out_ptr, + values_in, + values_out, + n_elems, + n_segments, + start_offset_it, + end_offset_it, + build_cache, + test_key); // Create expected result by sorting each segment std::vector expected_keys = host_keys; @@ -298,6 +309,7 @@ C2H_TEST_LIST("segmented_sort can sort key-value pairs with integral types", keys_out_ptr, values_in_ptr, values_out_ptr, + n_elems, n_segments, start_offset_it, end_offset_it, @@ -405,6 +417,7 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust keys_out_ptr, values_in_ptr, values_out_ptr, + n_elems, n_segments, start_offset_it, end_offset_it, @@ -553,6 +566,7 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va keys_out_ptr, values_in_ptr, values_out_ptr, + n_elems, n_segments, start_offset_it, end_offset_it, From b4c720d15bcc4aca2d5df18bff47e26cc4c409ba Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 20:12:52 +0000 Subject: [PATCH 014/100] Begin work to reuse CUB policies in c.parallel --- cub/cub/agent/agent_radix_sort_downsweep.cuh | 20 ++++++++++++++++ cub/cub/agent/agent_sub_warp_merge_sort.cuh | 12 ++++++++++ cub/cub/agent/agent_three_way_partition.cuh | 13 +++++++++++ .../dispatch/dispatch_segmented_sort.cuh | 23 ++++++++++++++----- 4 files changed, 62 insertions(+), 6 deletions(-) diff --git a/cub/cub/agent/agent_radix_sort_downsweep.cuh b/cub/cub/agent/agent_radix_sort_downsweep.cuh index f69dedc60d7..e069d82e990 100644 --- a/cub/cub/agent/agent_radix_sort_downsweep.cuh +++ b/cub/cub/agent/agent_radix_sort_downsweep.cuh @@ -119,6 +119,26 @@ struct AgentRadixSortDownsweepPolicy : ScalingType static constexpr BlockScanAlgorithm SCAN_ALGORITHM = _SCAN_ALGORITHM; }; +#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) +namespace detail +{ +// Only define this when needed. +// Because of overload woes, this depends on C++20 concepts. util_device.h checks that concepts are available when +// either runtime policies or PTX JSON information are enabled, so if they are, this is always valid. The generic +// version is always defined, and that's the only one needed for regular CUB operations. +// +// TODO: enable this unconditionally once concepts are always available +CUB_DETAIL_POLICY_WRAPPER_DEFINE( + RadixSortDownsweepAgentPolicy, + (GenericAgentPolicy), + (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), + (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), + (RANK_ALGORITHM, RankAlgorithm, cub::RadixRankAlgorithm), + (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm), + (RADIX_BITS, RadixBits, int) ) +} // namespace detail +#endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) + /****************************************************************************** * Thread block abstractions ******************************************************************************/ diff --git a/cub/cub/agent/agent_sub_warp_merge_sort.cuh b/cub/cub/agent/agent_sub_warp_merge_sort.cuh index 757c8151528..88aaa01489f 100644 --- a/cub/cub/agent/agent_sub_warp_merge_sort.cuh +++ b/cub/cub/agent/agent_sub_warp_merge_sort.cuh @@ -76,6 +76,18 @@ struct AgentSmallAndMediumSegmentedSortPolicy static constexpr int SEGMENTS_PER_SMALL_BLOCK = BLOCK_THREADS / SmallPolicyT::WARP_THREADS; }; +#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) +namespace detail +{ +CUB_DETAIL_POLICY_WRAPPER_DEFINE( + SmallAndMediumSegmentedSortAgentPolicy, + (GenericAgentPolicy), + (BLOCK_THREADS, BlockThreads, int), + (SEGMENTS_PER_MEDIUM_BLOCK, SegmentsPerMediumBlock, int), + (SEGMENTS_PER_SMALL_BLOCK, SegmentsPerSmallBlock, int) ) +} // namespace detail +#endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) + namespace detail { namespace sub_warp_merge_sort diff --git a/cub/cub/agent/agent_three_way_partition.cuh b/cub/cub/agent/agent_three_way_partition.cuh index 6ef3d1d6395..169eff4e636 100644 --- a/cub/cub/agent/agent_three_way_partition.cuh +++ b/cub/cub/agent/agent_three_way_partition.cuh @@ -73,9 +73,22 @@ struct AgentThreeWayPartitionPolicy }; }; +#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) namespace detail { +CUB_DETAIL_POLICY_WRAPPER_DEFINE( + ThreeWayPartitionAgentPolicy, + (GenericAgentPolicy), + (BLOCK_THREADS, BlockThreads, int), + (ITEMS_PER_THREAD, ItemsPerThread, int), + (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), + (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), + (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm)) +} // namespace detail +#endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) +namespace detail +{ namespace three_way_partition { diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 71bc9b44787..491965549c0 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -382,8 +382,12 @@ struct DispatchSegmentedSort KernelSource kernel_source; + PartitionKernelSource partition_kernel_source; + KernelLauncherFactory launcher_factory; + typename PartitionPolicyHub::MaxPolicy partition_max_policy; + template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT policy = {}) { @@ -511,7 +515,9 @@ struct DispatchSegmentedSort large_segments_selector, small_segments_selector, max_num_segments_per_invocation, - stream); + stream, + partition_kernel_source, + launcher_factory); device_partition_temp_storage.grow(three_way_partition_temp_storage_bytes); } @@ -638,7 +644,8 @@ struct DispatchSegmentedSort return error; } - template + template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -650,9 +657,11 @@ struct DispatchSegmentedSort EndOffsetIteratorT d_end_offsets, bool is_overwrite_okay, cudaStream_t stream, - KernelSource kernel_source = {}, - KernelLauncherFactory launcher_factory = {}, - MaxPolicyT max_policy = {}) + KernelSource kernel_source = {}, + PartitionKernelSource partition_kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}, + PartitionMaxPolicyT partition_max_policy = {}) { // Get PTX version int ptx_version = 0; @@ -674,7 +683,9 @@ struct DispatchSegmentedSort is_overwrite_okay, stream, kernel_source, - launcher_factory}; + partition_kernel_source, + launcher_factory, + partition_max_policy}; // Dispatch to chained policy return CubDebug(max_policy.Invoke(ptx_version, dispatch)); From ed7d6aa687fd69bde573faf10740db750f763ff9 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 21:32:45 +0000 Subject: [PATCH 015/100] Add encoded policy method --- cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index 3056ddf3bc3..b06c88584bd 100644 --- a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -140,6 +140,15 @@ struct SegmentedSortPolicyWrapper() = LargeSegment().EncodedPolicy(), + key<"SmallAndMediumSegmentedSortPolicy">() = SmallAndMediumSegmentedSort().EncodedPolicy()>(); + } +#endif }; template From c210f84c4410b43b1d72b56dd986cb78cdc693f3 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 21:47:10 +0000 Subject: [PATCH 016/100] Fix other call to device partition --- .../dispatch/dispatch_segmented_sort.cuh | 34 +++++++++++++++++-- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 491965549c0..707fe740900 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -757,7 +757,32 @@ private: auto medium_indices_iterator = THRUST_NS_QUALIFIER::make_reverse_iterator(large_and_medium_segments_indices.get() + current_num_segments); - error = CubDebug(cub::DevicePartition::IfNoNVTX( + // We call partition through dispatch instead of device because c.parallel needs to be able to call the kernel. + // This approach propagates the type erasure to partition. + using ChooseOffsetT = detail::choose_signed_offset; + using PartitionOffsetT = typename ChooseOffsetT::type; + using DispatchThreeWayPartitionIfT = cub::DispatchThreeWayPartitionIf< + THRUST_NS_QUALIFIER::counting_iterator, + decltype(large_and_medium_segments_indices.get()), + decltype(small_segments_indices.get()), + decltype(medium_indices_iterator), + decltype(group_sizes.get()), + LargeSegmentsSelectorT, + SmallSegmentsSelectorT, + PartitionOffsetT, + PartitionPolicyHub, + PartitionKernelSource, + KernelLauncherFactory>; + + // Signed integer type for global offsets + // Check if the number of items exceeds the range covered by the selected signed offset type + cudaError_t error = ChooseOffsetT::is_exceeding_offset_type(num_items); + if (error) + { + return error; + } + + DispatchThreeWayPartitionIfT::Dispatch( device_partition_temp_storage.get(), three_way_partition_temp_storage_bytes, THRUST_NS_QUALIFIER::counting_iterator(0), @@ -765,10 +790,13 @@ private: small_segments_indices.get(), medium_indices_iterator, group_sizes.get(), - current_num_segments, large_segments_selector, small_segments_selector, - stream)); + current_num_segments, + stream, + partition_kernel_source, + launcher_factory); + if (cudaSuccess != error) { return error; From e4f7748e20b67566fb4c15fed222079df26c37ad Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 23:36:49 +0000 Subject: [PATCH 017/100] Fix issue with tuning policies --- c/parallel/src/segmented_sort.cu | 72 ++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index ff2df1cd707..c3f07aba579 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -270,6 +270,9 @@ struct segmented_sort_runtime_tuning_policy { cub::detail::RuntimeRadixSortDownsweepAgentPolicy large_segment; cub::detail::RuntimeSmallAndMediumSegmentedSortAgentPolicy small_and_medium_segment; + cub::detail::RuntimeSubWarpMergeSortAgentPolicy small_segment; + cub::detail::RuntimeSubWarpMergeSortAgentPolicy medium_segment; + int partitioning_threshold; auto LargeSegment() const { @@ -280,6 +283,75 @@ struct segmented_sort_runtime_tuning_policy return small_and_medium_segment; } + void CheckLoadModifierIsNotLDG() const + { + if (large_segment.LoadModifier() == cub::CacheLoadModifier::LOAD_LDG) + { + throw std::runtime_error("The memory consistency model does not apply to texture accesses"); + } + } + + void CheckLoadAlgorithmIsNotStriped() const + { + if (large_segment.LoadAlgorithm() == cub::BLOCK_LOAD_STRIPED + || medium_segment.LoadAlgorithm() == cub::WARP_LOAD_STRIPED + || small_segment.LoadAlgorithm() == cub::WARP_LOAD_STRIPED) + { + throw std::runtime_error("Striped load will make this algorithm unstable"); + } + } + + void CheckStoreAlgorithmIsNotStriped() const + { + if (medium_segment.StoreAlgorithm() == cub::WARP_STORE_STRIPED + || small_segment.StoreAlgorithm() == cub::WARP_STORE_STRIPED) + { + throw std::runtime_error("Striped stores will produce unsorted results"); + } + } + + int PartitioningThreshold() const + { + return partitioning_threshold; + } + + int LargeSegmentRadixBits() const + { + return large_segment.RadixBits(); + } + + int SegmentsPerSmallBlock() const + { + return small_and_medium_segment.SegmentsPerSmallBlock(); + } + + int SegmentsPerMediumBlock() const + { + return small_and_medium_segment.SegmentsPerMediumBlock(); + } + + int SmallPolicyItemsPerTile() const + { + return small_segment.ItemsPerTile(); + } + + int MediumPolicyItemsPerTile() const + { + return medium_segment.ItemsPerTile(); + } + + template + int BlockThreads(PolicyT policy) const + { + return policy.BlockThreads(); + } + + template + int ItemsPerThread(PolicyT policy) const + { + return policy.ItemsPerThread(); + } + using MaxPolicy = segmented_sort_runtime_tuning_policy; template From c725a18fa26d8a597f136ffe4ab931e8be4a1a6b Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 14 Aug 2025 23:46:50 +0000 Subject: [PATCH 018/100] Add missing runtime policies and use advance iterators function instead of addition --- cub/cub/agent/agent_radix_sort_downsweep.cuh | 5 ++--- cub/cub/agent/agent_sub_warp_merge_sort.cuh | 10 ++++++++++ cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 9 +++++++-- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/cub/cub/agent/agent_radix_sort_downsweep.cuh b/cub/cub/agent/agent_radix_sort_downsweep.cuh index e069d82e990..026abbe4bbd 100644 --- a/cub/cub/agent/agent_radix_sort_downsweep.cuh +++ b/cub/cub/agent/agent_radix_sort_downsweep.cuh @@ -131,11 +131,10 @@ namespace detail CUB_DETAIL_POLICY_WRAPPER_DEFINE( RadixSortDownsweepAgentPolicy, (GenericAgentPolicy), - (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), + (RADIX_BITS, RadixBits, int) (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), (RANK_ALGORITHM, RankAlgorithm, cub::RadixRankAlgorithm), - (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm), - (RADIX_BITS, RadixBits, int) ) + (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm)) } // namespace detail #endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) diff --git a/cub/cub/agent/agent_sub_warp_merge_sort.cuh b/cub/cub/agent/agent_sub_warp_merge_sort.cuh index 88aaa01489f..a7b1c1b4686 100644 --- a/cub/cub/agent/agent_sub_warp_merge_sort.cuh +++ b/cub/cub/agent/agent_sub_warp_merge_sort.cuh @@ -79,6 +79,16 @@ struct AgentSmallAndMediumSegmentedSortPolicy #if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) namespace detail { +CUB_DETAIL_POLICY_WRAPPER_DEFINE( + SubWarpMergeSortAgentPolicy, + (GenericAgentPolicy), + (WARP_THREADS, WarpThreads, int), + (ITEMS_PER_THREAD, ItemsPerThread, int), + (ITEMS_PER_TILE, ItemsPerTile, int), + (LOAD_ALGORITHM, LoadAlgorithm, cub::WarpLoadAlgorithm), + (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), + (STORE_ALGORITHM, StoreAlgorithm, cub::WarpStoreAlgorithm)) + CUB_DETAIL_POLICY_WRAPPER_DEFINE( SmallAndMediumSegmentedSortAgentPolicy, (GenericAgentPolicy), diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 707fe740900..62871395a2c 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -40,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -751,8 +752,12 @@ private: large_segments_selector.base_segment_offset = current_seg_offset; small_segments_selector.base_segment_offset = current_seg_offset; - [[maybe_unused]] auto current_begin_offset = d_begin_offsets + current_seg_offset; - [[maybe_unused]] auto current_end_offset = d_end_offsets + current_seg_offset; + + BeginOffsetIteratorT current_begin_offset = d_begin_offsets; + EndOffsetIteratorT current_end_offset = d_end_offsets; + + detail::advance_iterators_inplace_if_supported(current_begin_offset, current_seg_offset); + detail::advance_iterators_inplace_if_supported(current_end_offset, current_seg_offset); auto medium_indices_iterator = THRUST_NS_QUALIFIER::make_reverse_iterator(large_and_medium_segments_indices.get() + current_num_segments); From 96194dce2e9b0ef9f04e3f5fdbae4fc79a1d334b Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 18 Aug 2025 22:51:42 +0000 Subject: [PATCH 019/100] Begin fixing runtime policy in segmented sort --- c/parallel/src/segmented_sort.cu | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index c3f07aba579..d0b3b400d9e 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -112,18 +112,6 @@ struct segmented_sort_runtime_policy } }; -// Function to create runtime policy from JSON -segmented_sort_runtime_policy from_json(const nlohmann::json& j) -{ - return segmented_sort_runtime_policy{ - .partitioning_threshold = j["PartitioningThreshold"].get(), - .large_segment_radix_bits = j["LargeSegmentRadixBits"].get(), - .segments_per_small_block = j["SegmentsPerSmallBlock"].get(), - .segments_per_medium_block = j["SegmentsPerMediumBlock"].get(), - .small_policy_items_per_tile = j["SmallPolicyItemsPerTile"].get(), - .medium_policy_items_per_tile = j["MediumPolicyItemsPerTile"].get()}; -} - std::string get_device_segmented_sort_fallback_kernel_name( std::string_view /* key_iterator_t */, std::string_view /* value_iterator_t */, @@ -516,7 +504,9 @@ struct __align__({1}) storage_t {{ nlohmann::json runtime_policy = get_policy(policy_wrapper_expr, ptx_query_tu_src, ptx_args); - auto segmented_sort_policy = segmented_sort::from_json(runtime_policy); + using cub::detail::RuntimeSegmentedSortAgentPolicy; + auto [segmented_sort_policy, segmented_sort_policy_str] = + RuntimeSegmentedSortAgentPolicy::from_json(runtime_policy, "SegmentedSortPolicy"); // Extract sub-policy information if available std::string small_and_medium_policy_str; From eb3c174348290d9dbe4b9464edbbe72a587411d7 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 18 Aug 2025 22:53:09 +0000 Subject: [PATCH 020/100] Fix missing comma --- cub/cub/agent/agent_radix_sort_downsweep.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cub/cub/agent/agent_radix_sort_downsweep.cuh b/cub/cub/agent/agent_radix_sort_downsweep.cuh index 026abbe4bbd..a7440e7f850 100644 --- a/cub/cub/agent/agent_radix_sort_downsweep.cuh +++ b/cub/cub/agent/agent_radix_sort_downsweep.cuh @@ -131,7 +131,8 @@ namespace detail CUB_DETAIL_POLICY_WRAPPER_DEFINE( RadixSortDownsweepAgentPolicy, (GenericAgentPolicy), - (RADIX_BITS, RadixBits, int) (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), + (RADIX_BITS, RadixBits, int), + (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), (RANK_ALGORITHM, RankAlgorithm, cub::RadixRankAlgorithm), (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm)) From 0b4ba625072b7adcd2d59e05360a012f5e56ee31 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 18 Aug 2025 23:27:52 +0000 Subject: [PATCH 021/100] Split SmallAndMediumSegmentedSortPolicy into separate small and medium policies --- cub/benchmarks/bench/segmented_sort/keys.cu | 17 +- cub/cub/agent/agent_sub_warp_merge_sort.cuh | 32 +-- .../dispatch/dispatch_segmented_sort.cuh | 32 +-- .../dispatch/kernels/segmented_sort.cuh | 17 +- .../dispatch/tuning/tuning_segmented_sort.cuh | 196 ++++++++++-------- cub/test/catch2_segmented_sort_helper.cuh | 11 +- 6 files changed, 147 insertions(+), 158 deletions(-) diff --git a/cub/benchmarks/bench/segmented_sort/keys.cu b/cub/benchmarks/bench/segmented_sort/keys.cu index c8011db8915..d455ce4531b 100644 --- a/cub/benchmarks/bench/segmented_sort/keys.cu +++ b/cub/benchmarks/bench/segmented_sort/keys.cu @@ -128,19 +128,10 @@ struct device_seg_sort_policy_hub static constexpr int ITEMS_PER_SMALL_THREAD = TUNE_S_ITEMS; static constexpr int ITEMS_PER_MEDIUM_THREAD = TUNE_M_ITEMS; - using SmallAndMediumSegmentedSortPolicyT = cub::AgentSmallAndMediumSegmentedSortPolicy< - - BLOCK_THREADS, - - // Small policy - cub:: - AgentSubWarpMergeSortPolicy, - - // Medium policy - cub::AgentSubWarpMergeSortPolicy>; + using SmallSegmentPolicy = cub:: + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = cub:: + AgentSubWarpMergeSortPolicy; }; using MaxPolicy = Policy500; diff --git a/cub/cub/agent/agent_sub_warp_merge_sort.cuh b/cub/cub/agent/agent_sub_warp_merge_sort.cuh index a7b1c1b4686..063b985424e 100644 --- a/cub/cub/agent/agent_sub_warp_merge_sort.cuh +++ b/cub/cub/agent/agent_sub_warp_merge_sort.cuh @@ -48,53 +48,39 @@ CUB_NAMESPACE_BEGIN -template struct AgentSubWarpMergeSortPolicy { - static constexpr int WARP_THREADS = WARP_THREADS_ARG; - static constexpr int ITEMS_PER_THREAD = ITEMS_PER_THREAD_ARG; - static constexpr int ITEMS_PER_TILE = WARP_THREADS * ITEMS_PER_THREAD; + static constexpr int BLOCK_THREADS = BLOCK_THREADS_ARG; + static constexpr int WARP_THREADS = WARP_THREADS_ARG; + static constexpr int ITEMS_PER_THREAD = ITEMS_PER_THREAD_ARG; + static constexpr int ITEMS_PER_TILE = WARP_THREADS * ITEMS_PER_THREAD; + static constexpr int SEGMENTS_PER_BLOCK = BLOCK_THREADS / WARP_THREADS; static constexpr cub::WarpLoadAlgorithm LOAD_ALGORITHM = LOAD_ALGORITHM_ARG; static constexpr cub::CacheLoadModifier LOAD_MODIFIER = LOAD_MODIFIER_ARG; static constexpr cub::WarpStoreAlgorithm STORE_ALGORITHM = STORE_ALGORITHM_ARG; }; -template -struct AgentSmallAndMediumSegmentedSortPolicy -{ - static constexpr int BLOCK_THREADS = BLOCK_THREADS_ARG; - using SmallPolicyT = SmallPolicy; - using MediumPolicyT = MediumPolicy; - - static constexpr int SEGMENTS_PER_MEDIUM_BLOCK = BLOCK_THREADS / MediumPolicyT::WARP_THREADS; - - static constexpr int SEGMENTS_PER_SMALL_BLOCK = BLOCK_THREADS / SmallPolicyT::WARP_THREADS; -}; - #if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) namespace detail { CUB_DETAIL_POLICY_WRAPPER_DEFINE( SubWarpMergeSortAgentPolicy, (GenericAgentPolicy), + (BLOCK_THREADS, BlockThreads, int), (WARP_THREADS, WarpThreads, int), (ITEMS_PER_THREAD, ItemsPerThread, int), (ITEMS_PER_TILE, ItemsPerTile, int), + (SEGMENTS_PER_BLOCK, SegmentsPerBlock, int), (LOAD_ALGORITHM, LoadAlgorithm, cub::WarpLoadAlgorithm), (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), (STORE_ALGORITHM, StoreAlgorithm, cub::WarpStoreAlgorithm)) - -CUB_DETAIL_POLICY_WRAPPER_DEFINE( - SmallAndMediumSegmentedSortAgentPolicy, - (GenericAgentPolicy), - (BLOCK_THREADS, BlockThreads, int), - (SEGMENTS_PER_MEDIUM_BLOCK, SegmentsPerMediumBlock, int), - (SEGMENTS_PER_SMALL_BLOCK, SegmentsPerSmallBlock, int) ) } // namespace detail #endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 62871395a2c..baaf419e138 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -111,11 +111,11 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont _CubLog("Invoking " "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", static_cast(blocks_in_grid), - wrapped_policy.BlockThreads(wrapped_policy.LargeSegment()), + wrapped_policy.LargeSegment().BlockThreads(), (long long) stream); #endif // CUB_DEBUG_LOG - launcher_factory(blocks_in_grid, wrapped_policy.BlockThreads(wrapped_policy.LargeSegment()), 0, stream) + launcher_factory(blocks_in_grid, wrapped_policy.LargeSegment().BlockThreads(), 0, stream) .doit(large_kernel, large_and_medium_segments_indices, d_current_keys, @@ -159,14 +159,11 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t DeviceSegmentedSortCont _CubLog("Invoking " "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", static_cast(small_and_medium_blocks_in_grid), - wrapped_policy.BlockThreads(wrapped_policy.SmallAndMediumSegmentedSort()), + wrapped_policy.SmallSegment().BlockThreads(), (long long) stream); #endif // CUB_DEBUG_LOG - launcher_factory(small_and_medium_blocks_in_grid, - wrapped_policy.BlockThreads(wrapped_policy.SmallAndMediumSegmentedSort()), - 0, - stream) + launcher_factory(small_and_medium_blocks_in_grid, wrapped_policy.SmallSegment().BlockThreads(), 0, stream) .doit(small_kernel, small_segments, medium_segments, @@ -394,30 +391,16 @@ struct DispatchSegmentedSort { auto wrapped_policy = detail::segmented_sort::MakeSegmentedSortPolicyWrapper(policy); - // using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - // using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - wrapped_policy.CheckLoadModifierIsNotLDG(); - // static_assert(LargeSegmentPolicyT::LOAD_MODIFIER != CacheLoadModifier::LOAD_LDG, - // "The memory consistency model does not apply to texture accesses"); - if constexpr (!KEYS_ONLY) { wrapped_policy.CheckLoadAlgorithmIsNotStriped(); } - // static_assert(KEYS_ONLY || LargeSegmentPolicyT::LOAD_ALGORITHM != BLOCK_LOAD_STRIPED - // || SmallAndMediumPolicyT::MediumPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED - // || SmallAndMediumPolicyT::SmallPolicyT::LOAD_ALGORITHM != WARP_LOAD_STRIPED, - // "Striped load will make this algorithm unstable"); wrapped_policy.CheckStoreAlgorithmIsNotStriped(); - // static_assert(SmallAndMediumPolicyT::MediumPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED - // || SmallAndMediumPolicyT::SmallPolicyT::STORE_ALGORITHM != WARP_STORE_STRIPED, - // "Striped stores will produce unsorted results"); const int radix_bits = wrapped_policy.LargeSegmentRadixBits(); - // constexpr int radix_bits = LargeSegmentPolicyT::RADIX_BITS; cudaError error = cudaSuccess; @@ -920,9 +903,8 @@ private: { cudaError_t error = cudaSuccess; - const auto blocks_in_grid = static_cast(num_segments); - constexpr auto threads_in_block = - static_cast(wrapped_policy.BlockThreads(wrapped_policy.LargeSegment())); + const auto blocks_in_grid = static_cast(num_segments); + constexpr auto threads_in_block = static_cast(wrapped_policy.LargeSegment().BlockThreads()); // Log kernel configuration #ifdef CUB_DEBUG_LOG @@ -931,7 +913,7 @@ private: blocks_in_grid, threads_in_block, (long long) stream, - wrapped_policy.ItemsPerThread(wrapped_policy.LargeSegment()), + wrapped_policy.LargeSegment().ItemsPerThread(), wrapped_policy.LargeSegmentRadixBits()); #endif // CUB_DEBUG_LOG diff --git a/cub/cub/device/dispatch/kernels/segmented_sort.cuh b/cub/cub/device/dispatch/kernels/segmented_sort.cuh index 2d1c2a1fa18..6d965580637 100644 --- a/cub/cub/device/dispatch/kernels/segmented_sort.cuh +++ b/cub/cub/device/dispatch/kernels/segmented_sort.cuh @@ -139,7 +139,7 @@ __launch_bounds__(ChainedPolicyT::ActivePolicy::LargeSegmentPolicy::BLOCK_THREAD { using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - using MediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT::MediumPolicyT; + using MediumPolicyT = typename ActivePolicyT::MediumSegmentPolicy; const auto segment_id = static_cast(blockIdx.x); OffsetT segment_begin = d_begin_offsets[segment_id]; @@ -303,7 +303,7 @@ template -__launch_bounds__(ChainedPolicyT::ActivePolicy::SmallAndMediumSegmentedSortPolicyT::BLOCK_THREADS) +__launch_bounds__(ChainedPolicyT::ActivePolicy::SmallSegmentPolicy::BLOCK_THREADS) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceSegmentedSortKernelSmall( local_segment_index_t small_segments, local_segment_index_t medium_segments, @@ -322,10 +322,9 @@ __launch_bounds__(ChainedPolicyT::ActivePolicy::SmallAndMediumSegmentedSortPolic const local_segment_index_t tid = threadIdx.x; const local_segment_index_t bid = blockIdx.x; - using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; - using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - using MediumPolicyT = typename SmallAndMediumPolicyT::MediumPolicyT; - using SmallPolicyT = typename SmallAndMediumPolicyT::SmallPolicyT; + using ActivePolicyT = typename ChainedPolicyT::ActivePolicy; + using SmallPolicyT = typename ActivePolicyT::SmallSegmentPolicy; + using MediumPolicyT = typename ActivePolicyT::MediumSegmentPolicy; constexpr auto threads_per_medium_segment = static_cast(MediumPolicyT::WARP_THREADS); constexpr auto threads_per_small_segment = static_cast(SmallPolicyT::WARP_THREADS); @@ -336,11 +335,9 @@ __launch_bounds__(ChainedPolicyT::ActivePolicy::SmallAndMediumSegmentedSortPolic using SmallAgentWarpMergeSortT = sub_warp_merge_sort::AgentSubWarpSort; - constexpr auto segments_per_medium_block = - static_cast(SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK); + constexpr auto segments_per_medium_block = static_cast(MediumPolicyT::SEGMENTS_PER_BLOCK); - constexpr auto segments_per_small_block = - static_cast(SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK); + constexpr auto segments_per_small_block = static_cast(SmallPolicyT::SEGMENTS_PER_BLOCK); __shared__ union { diff --git a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index b06c88584bd..5bac6740507 100644 --- a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -59,8 +59,8 @@ struct SegmentedSortPolicyWrapper : PolicyT template struct SegmentedSortPolicyWrapper> - : StaticPolicyT + typename StaticPolicyT::SmallSegmentPolicy, + typename StaticPolicyT::MediumSegmentPolicy>> : StaticPolicyT { CUB_RUNTIME_FUNCTION SegmentedSortPolicyWrapper(StaticPolicyT base) : StaticPolicyT(base) @@ -71,9 +71,14 @@ struct SegmentedSortPolicyWrapper - CUB_RUNTIME_FUNCTION static constexpr int BlockThreads(PolicyT /*policy*/) - { - return PolicyT::BLOCK_THREADS; - } - - template - CUB_RUNTIME_FUNCTION static constexpr int ItemsPerThread(PolicyT /*policy*/) - { - return PolicyT::ITEMS_PER_THREAD; + return StaticPolicyT::MediumSegmentPolicy::ITEMS_PER_TILE; } #if defined(CUB_ENABLE_POLICY_PTX_JSON) _CCCL_DEVICE static constexpr auto EncodedPolicy() { using namespace ptx_json; - return object() = LargeSegment().EncodedPolicy(), - key<"SmallAndMediumSegmentedSortPolicy">() = SmallAndMediumSegmentedSort().EncodedPolicy()>(); + return object() = LargeSegment().EncodedPolicy(), + key<"SmallSegmentPolicy">() = SmallSegment().EncodedPolicy(), + key<"MediumSegmentPolicy">() = MediumSegment().EncodedPolicy()>(); } #endif }; @@ -181,12 +173,19 @@ struct policy_hub static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(7); static constexpr int ITEMS_PER_MEDIUM_THREAD = Nominal4BItemsToItems(7); - using SmallAndMediumSegmentedSortPolicyT = AgentSmallAndMediumSegmentedSortPolicy< - BLOCK_THREADS, - // Small policy - AgentSubWarpMergeSortPolicy<4 /* Threads per segment */, ITEMS_PER_SMALL_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>, - // Medium policy - AgentSubWarpMergeSortPolicy<32 /* Threads per segment */, ITEMS_PER_MEDIUM_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>>; + + using SmallSegmentPolicy = + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = + AgentSubWarpMergeSortPolicy; }; struct Policy600 : ChainedPolicy<600, Policy600, Policy500> @@ -207,12 +206,19 @@ struct policy_hub static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(9); static constexpr int ITEMS_PER_MEDIUM_THREAD = Nominal4BItemsToItems(9); - using SmallAndMediumSegmentedSortPolicyT = AgentSmallAndMediumSegmentedSortPolicy< - BLOCK_THREADS, - // Small policy - AgentSubWarpMergeSortPolicy<4 /* Threads per segment */, ITEMS_PER_SMALL_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>, - // Medium policy - AgentSubWarpMergeSortPolicy<32 /* Threads per segment */, ITEMS_PER_MEDIUM_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>>; + + using SmallSegmentPolicy = + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = + AgentSubWarpMergeSortPolicy; }; struct Policy610 : ChainedPolicy<610, Policy610, Policy600> @@ -233,12 +239,19 @@ struct policy_hub static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(9); static constexpr int ITEMS_PER_MEDIUM_THREAD = Nominal4BItemsToItems(9); - using SmallAndMediumSegmentedSortPolicyT = AgentSmallAndMediumSegmentedSortPolicy< - BLOCK_THREADS, - // Small policy - AgentSubWarpMergeSortPolicy<4 /* Threads per segment */, ITEMS_PER_SMALL_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>, - // Medium policy - AgentSubWarpMergeSortPolicy<32 /* Threads per segment */, ITEMS_PER_MEDIUM_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>>; + + using SmallSegmentPolicy = + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = + AgentSubWarpMergeSortPolicy; }; struct Policy620 : ChainedPolicy<620, Policy620, Policy610> @@ -259,12 +272,19 @@ struct policy_hub static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(9); static constexpr int ITEMS_PER_MEDIUM_THREAD = Nominal4BItemsToItems(9); - using SmallAndMediumSegmentedSortPolicyT = AgentSmallAndMediumSegmentedSortPolicy< - BLOCK_THREADS, - // Small policy - AgentSubWarpMergeSortPolicy<4 /* Threads per segment */, ITEMS_PER_SMALL_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>, - // Medium policy - AgentSubWarpMergeSortPolicy<32 /* Threads per segment */, ITEMS_PER_MEDIUM_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>>; + + using SmallSegmentPolicy = + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = + AgentSubWarpMergeSortPolicy; }; struct Policy700 : ChainedPolicy<700, Policy700, Policy620> @@ -285,15 +305,19 @@ struct policy_hub static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(7); static constexpr int ITEMS_PER_MEDIUM_THREAD = Nominal4BItemsToItems(KEYS_ONLY ? 11 : 7); - using SmallAndMediumSegmentedSortPolicyT = AgentSmallAndMediumSegmentedSortPolicy< - BLOCK_THREADS, - // Small policy - AgentSubWarpMergeSortPolicy, - // Medium policy - AgentSubWarpMergeSortPolicy<32 /* Threads per segment */, ITEMS_PER_MEDIUM_THREAD, WARP_LOAD_DIRECT, LOAD_DEFAULT>>; + + using SmallSegmentPolicy = + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = + AgentSubWarpMergeSortPolicy; }; struct Policy800 : ChainedPolicy<800, Policy800, Policy700> @@ -312,15 +336,19 @@ struct policy_hub static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(9); static constexpr int ITEMS_PER_MEDIUM_THREAD = Nominal4BItemsToItems(KEYS_ONLY ? 7 : 11); - using SmallAndMediumSegmentedSortPolicyT = AgentSmallAndMediumSegmentedSortPolicy< - BLOCK_THREADS, - // Small policy - AgentSubWarpMergeSortPolicy, - // Medium policy - AgentSubWarpMergeSortPolicy<32 /* Threads per segment */, ITEMS_PER_MEDIUM_THREAD, WARP_LOAD_TRANSPOSE, LOAD_DEFAULT>>; + + using SmallSegmentPolicy = + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = + AgentSubWarpMergeSortPolicy; }; struct Policy860 : ChainedPolicy<860, Policy860, Policy800> @@ -340,15 +368,19 @@ struct policy_hub static constexpr bool LARGE_ITEMS = sizeof(DominantT) > 4; static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(LARGE_ITEMS ? 7 : 9); static constexpr int ITEMS_PER_MEDIUM_THREAD = Nominal4BItemsToItems(LARGE_ITEMS ? 9 : 7); - using SmallAndMediumSegmentedSortPolicyT = AgentSmallAndMediumSegmentedSortPolicy< - BLOCK_THREADS, - // Small policy - AgentSubWarpMergeSortPolicy, - // Medium policy - AgentSubWarpMergeSortPolicy<16 /* Threads per segment */, ITEMS_PER_MEDIUM_THREAD, WARP_LOAD_TRANSPOSE, LOAD_LDG>>; + + using SmallSegmentPolicy = + AgentSubWarpMergeSortPolicy; + using MediumSegmentPolicy = + AgentSubWarpMergeSortPolicy; }; using MaxPolicy = Policy860; diff --git a/cub/test/catch2_segmented_sort_helper.cuh b/cub/test/catch2_segmented_sort_helper.cuh index f0e16d6b596..2acb0f78d61 100644 --- a/cub/test/catch2_segmented_sort_helper.cuh +++ b/cub/test/catch2_segmented_sort_helper.cuh @@ -1572,12 +1572,13 @@ struct generate_edge_case_offsets_dispatch { NV_IF_TARGET( NV_IS_HOST, - (using SmallAndMediumPolicyT = typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT; - using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; + (using SmallPolicyT = typename ActivePolicyT::SmallSegmentPolicy; + using MediumPolicyT = typename ActivePolicyT::MediumSegmentPolicy; + using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy; - small_segment_max_segment_size = SmallAndMediumPolicyT::SmallPolicyT::ITEMS_PER_TILE; - items_per_small_segment = SmallAndMediumPolicyT::SmallPolicyT::ITEMS_PER_THREAD; - medium_segment_max_segment_size = SmallAndMediumPolicyT::MediumPolicyT::ITEMS_PER_TILE; + small_segment_max_segment_size = SmallPolicyT::ITEMS_PER_TILE; + items_per_small_segment = SmallPolicyT::ITEMS_PER_THREAD; + medium_segment_max_segment_size = MediumPolicyT::ITEMS_PER_TILE; single_thread_segment_size = items_per_small_segment; large_cached_segment_max_segment_size = LargeSegmentPolicyT::BLOCK_THREADS * LargeSegmentPolicyT::ITEMS_PER_THREAD; // From 57764c687ae61f1c4fdecf544abf9904a0a7a655 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 19 Aug 2025 00:53:10 +0000 Subject: [PATCH 022/100] Continue fixing runtime segmented sort policies --- c/parallel/src/segmented_sort.cu | 159 ++++++------------------------- 1 file changed, 29 insertions(+), 130 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index d0b3b400d9e..0c4fd2f61e8 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -41,77 +41,12 @@ struct device_segmented_sort_policy; using OffsetT = long; static_assert(std::is_same_v, OffsetT>, "OffsetT must be long"); -// check we can map OffsetT to ::cuda::std::int64_t +// check we can map OffsetT to cuda::std::int64_t static_assert(std::is_signed_v); -static_assert(sizeof(OffsetT) == sizeof(::cuda::std::int64_t)); +static_assert(sizeof(OffsetT) == sizeof(cuda::std::int64_t)); namespace segmented_sort { - -// Runtime policy structure for segmented sort -struct segmented_sort_runtime_policy -{ - int partitioning_threshold; - int large_segment_radix_bits; - int segments_per_small_block; - int segments_per_medium_block; - int small_policy_items_per_tile; - int medium_policy_items_per_tile; - - // Required methods for SegmentedSortPolicyWrapper - constexpr int PartitioningThreshold() const - { - return partitioning_threshold; - } - constexpr int LargeSegmentRadixBits() const - { - return large_segment_radix_bits; - } - constexpr int SegmentsPerSmallBlock() const - { - return segments_per_small_block; - } - constexpr int SegmentsPerMediumBlock() const - { - return segments_per_medium_block; - } - constexpr int SmallPolicyItemsPerTile() const - { - return small_policy_items_per_tile; - } - constexpr int MediumPolicyItemsPerTile() const - { - return medium_policy_items_per_tile; - } - - // Additional methods expected by SegmentedSortPolicyWrapper - constexpr void CheckLoadModifierIsNotLDG() const {} // No-op validation - constexpr void CheckLoadAlgorithmIsNotStriped() const {} // No-op validation - constexpr void CheckStoreAlgorithmIsNotStriped() const {} // No-op validation - - // Policy accessor methods - constexpr int BlockThreads(int /* large_segment_policy */) const - { - return 256; - } // Default block size - constexpr int LargeSegment() const - { - return 0; - } // Return index for large segment policy - constexpr auto SmallAndMediumSegmentedSort() const - { - return *this; - } // Return policy for small/medium segments - - using MaxPolicy = segmented_sort_runtime_policy; - - template - cudaError_t Invoke(int, F& op) - { - return op.template Invoke(*this); - } -}; - std::string get_device_segmented_sort_fallback_kernel_name( std::string_view /* key_iterator_t */, std::string_view /* value_iterator_t */, @@ -257,7 +192,6 @@ struct partition_kernel_source struct segmented_sort_runtime_tuning_policy { cub::detail::RuntimeRadixSortDownsweepAgentPolicy large_segment; - cub::detail::RuntimeSmallAndMediumSegmentedSortAgentPolicy small_and_medium_segment; cub::detail::RuntimeSubWarpMergeSortAgentPolicy small_segment; cub::detail::RuntimeSubWarpMergeSortAgentPolicy medium_segment; int partitioning_threshold; @@ -266,9 +200,15 @@ struct segmented_sort_runtime_tuning_policy { return large_segment; } - auto SmallAndMediumSegmentedSort() const + + auto SmallSegment() const + { + return small_segment; + } + + auto MediumSegment() const { - return small_and_medium_segment; + return medium_segment; } void CheckLoadModifierIsNotLDG() const @@ -310,12 +250,12 @@ struct segmented_sort_runtime_tuning_policy int SegmentsPerSmallBlock() const { - return small_and_medium_segment.SegmentsPerSmallBlock(); + return small_segment.SegmentsPerBlock(); } int SegmentsPerMediumBlock() const { - return small_and_medium_segment.SegmentsPerMediumBlock(); + return medium_segment.SegmentsPerBlock(); } int SmallPolicyItemsPerTile() const @@ -328,18 +268,6 @@ struct segmented_sort_runtime_tuning_policy return medium_segment.ItemsPerTile(); } - template - int BlockThreads(PolicyT policy) const - { - return policy.BlockThreads(); - } - - template - int ItemsPerThread(PolicyT policy) const - { - return policy.ItemsPerThread(); - } - using MaxPolicy = segmented_sort_runtime_tuning_policy; template @@ -504,50 +432,16 @@ struct __align__({1}) storage_t {{ nlohmann::json runtime_policy = get_policy(policy_wrapper_expr, ptx_query_tu_src, ptx_args); - using cub::detail::RuntimeSegmentedSortAgentPolicy; - auto [segmented_sort_policy, segmented_sort_policy_str] = - RuntimeSegmentedSortAgentPolicy::from_json(runtime_policy, "SegmentedSortPolicy"); + using cub::detail::RuntimeRadixSortDownsweepAgentPolicy; + auto [large_segment_policy, large_segment_policy_str] = + RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "LargeSegmentPolicy"); - // Extract sub-policy information if available - std::string small_and_medium_policy_str; - if (runtime_policy.contains("SmallAndMediumSegmentedSort")) - { - auto sub_policy = runtime_policy["SmallAndMediumSegmentedSort"]; - auto block_threads = sub_policy["BlockThreads"].get(); - auto segments_per_medium = sub_policy["SegmentsPerMediumBlock"].get(); - auto segments_per_small = sub_policy["SegmentsPerSmallBlock"].get(); - - small_and_medium_policy_str = std::format( - R"XXX( - // Small and Medium Segment Policy - static constexpr int SMALL_MEDIUM_BLOCK_THREADS = {0}; - static constexpr int SMALL_MEDIUM_SEGMENTS_PER_MEDIUM_BLOCK = {1}; - static constexpr int SMALL_MEDIUM_SEGMENTS_PER_SMALL_BLOCK = {2};)XXX", - block_threads, - segments_per_medium, - segments_per_small); - } + using cub::detail::RuntimeSubWarpMergeSortAgentPolicy; + auto [small_segment_policy, small_segment_policy_str] = + RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "SmallSegmentPolicy"); - // Build the policy structure manually - const std::string segmented_sort_policy_str = std::format( - R"XXX( - static constexpr int PARTITIONING_THRESHOLD = {0}; - static constexpr int LARGE_SEGMENT_RADIX_BITS = {1}; - static constexpr int SEGMENTS_PER_SMALL_BLOCK = {2}; - static constexpr int SEGMENTS_PER_MEDIUM_BLOCK = {3}; - static constexpr int SMALL_POLICY_ITEMS_PER_TILE = {4}; - static constexpr int MEDIUM_POLICY_ITEMS_PER_TILE = {5};{6} - using MaxPolicy = cub::detail::segmented_sort::policy_hub<{7}, {8}>::MaxPolicy; -)XXX", - segmented_sort_policy.partitioning_threshold, // 0 - segmented_sort_policy.large_segment_radix_bits, // 1 - segmented_sort_policy.segments_per_small_block, // 2 - segmented_sort_policy.segments_per_medium_block, // 3 - segmented_sort_policy.small_policy_items_per_tile, // 4 - segmented_sort_policy.medium_policy_items_per_tile, // 5 - small_and_medium_policy_str, // 6 - key_t, // 7 - value_t); // 8 + auto [medium_segment_policy, medium_segment_policy_str] = + RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "MediumSegmentPolicy"); // agent_policy_t is to specify parameters like policy_hub does in dispatch_segmented_sort.cuh constexpr std::string_view program_preamble_template = R"XXX( @@ -557,6 +451,8 @@ struct __align__({1}) storage_t {{ struct device_segmented_sort_policy {{ struct ActivePolicy {{ {2} + {3} + {4} }}; }}; )XXX"; @@ -565,7 +461,9 @@ struct device_segmented_sort_policy {{ program_preamble_template, jit_template_header_contents, // 0 dependent_definitions_src, // 1 - segmented_sort_policy_str); // 2 + large_segment_policy_str, // 2 + small_segment_policy_str, // 3 + medium_segment_policy_str); // 4 std::string segmented_sort_fallback_kernel_name = segmented_sort::get_device_segmented_sort_fallback_kernel_name( keys_in_iterator_name, @@ -655,7 +553,8 @@ struct device_segmented_sort_policy {{ build_ptr->cubin = (void*) result.data.release(); build_ptr->cubin_size = result.size; // Use the runtime policy extracted via from_json - build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_policy{segmented_sort_policy}; + build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_tuning_policy{ + large_segment_policy, small_segment_policy, medium_segment_policy}; } catch (const std::exception& exc) { @@ -724,7 +623,7 @@ CUresult cccl_device_segmented_sort( /* kernel_source */ {build}, /* partition_kernel_source */ {build}, /* launcher_factory */ cub::detail::CudaDriverLauncherFactory{cu_device, build.cc}, - /* policy */ *reinterpret_cast(build.runtime_policy), + /* policy */ *reinterpret_cast(build.runtime_policy), /* partition_policy */ *reinterpret_cast(build.partition_runtime_policy)); @@ -760,7 +659,7 @@ CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_res std::unique_ptr cubin(reinterpret_cast(build_ptr->cubin)); // Clean up the runtime policy - delete static_cast(build_ptr->runtime_policy); + delete static_cast(build_ptr->runtime_policy); check(cuLibraryUnload(build_ptr->library)); } catch (const std::exception& exc) From 01cc9d1a040b8a1df3618f04b6b3ec0217b96d4e Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 3 Sep 2025 00:24:53 +0000 Subject: [PATCH 023/100] Fix compilation errors --- c/parallel/src/segmented_sort.cu | 4 ++- c/parallel/test/test_segmented_sort.cpp | 34 +++++++++++-------------- 2 files changed, 18 insertions(+), 20 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 0c4fd2f61e8..edad1e0c833 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -443,6 +443,8 @@ struct __align__({1}) storage_t {{ auto [medium_segment_policy, medium_segment_policy_str] = RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "MediumSegmentPolicy"); + auto partitioning_threshold = static_cast(runtime_policy["PartitioningThreshold"].get()); + // agent_policy_t is to specify parameters like policy_hub does in dispatch_segmented_sort.cuh constexpr std::string_view program_preamble_template = R"XXX( #include @@ -554,7 +556,7 @@ struct device_segmented_sort_policy {{ build_ptr->cubin_size = result.size; // Use the runtime policy extracted via from_json build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_tuning_policy{ - large_segment_policy, small_segment_policy, medium_segment_policy}; + large_segment_policy, small_segment_policy, medium_segment_policy, partitioning_threshold}; } catch (const std::exception& exc) { diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index d52f8a36ee1..5a91ce1747e 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -25,6 +25,9 @@ #include #include +using key_types = c2h::type_list; +using item_t = float; + using BuildResultT = cccl_device_segmented_sort_build_result_t; struct segmented_sort_cleanup @@ -149,13 +152,9 @@ extern "C" __device__ {2} {0}({1}* state) } struct SegmentedSort_KeysOnly_Fixture_Tag; -C2H_TEST_LIST("segmented_sort can sort keys-only with integral types", - "[segmented_sort][keys_only]", - std::int32_t, - std::int64_t, - std::uint32_t, - std::uint64_t) +C2H_TEST("segmented_sort can sort keys-only with integral types", "[segmented_sort][keys_only]", key_types) { + using KeyT = c2h::get<0, TestType>; // generate choices for n_segments: 0, 13 and 2 random samples from [50, 200) const std::size_t n_segments = GENERATE(0, 13, take(2, random(50, 200))); // generate choices for segment size: 1, 20 and random samples @@ -163,21 +162,18 @@ C2H_TEST_LIST("segmented_sort can sort keys-only with integral types", const std::size_t n_elems = n_segments * segment_size; - std::vector host_keys = generate(n_elems); - std::vector host_keys_out(n_elems); + std::vector host_keys_int = generate(n_elems); + std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); + std::vector host_keys_out(n_elems); REQUIRE(host_keys.size() == n_elems); REQUIRE(host_keys_out.size() == n_elems); - pointer_t keys_in_ptr(host_keys); // copy from host to device - pointer_t keys_out_ptr(host_keys_out); // copy from host to device + pointer_t keys_in_ptr(host_keys); // copy from host to device + pointer_t keys_out_ptr(host_keys_out); // copy from host to device - // Create null value iterators for keys-only sorting - // For keys-only sorting, we create dummy iterators that won't be used - auto dummy_values_it = make_constant_iterator(std::string{"TestType"}); - dummy_values_it.state.value = TestType{}; - cccl_iterator_t values_in = dummy_values_it; - cccl_iterator_t values_out = dummy_values_it; + pointer_t values_in; + pointer_t values_out; using SizeT = unsigned long long; static constexpr std::string_view index_ty_name = "unsigned long long"; @@ -214,7 +210,7 @@ C2H_TEST_LIST("segmented_sort can sort keys-only with integral types", end_offset_it.state.segment_size = segment_size; auto& build_cache = get_cache(); - const auto& test_key = make_key(); + const auto& test_key = make_key(); segmented_sort( keys_in_ptr, @@ -229,7 +225,7 @@ C2H_TEST_LIST("segmented_sort can sort keys-only with integral types", test_key); // Create expected result by sorting each segment - std::vector expected_keys = host_keys; + std::vector expected_keys = host_keys; for (std::size_t i = 0; i < n_segments; ++i) { std::size_t segment_start = i * segment_size; @@ -237,7 +233,7 @@ C2H_TEST_LIST("segmented_sort can sort keys-only with integral types", std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end); } - REQUIRE(expected_keys == std::vector(keys_out_ptr)); + REQUIRE(expected_keys == std::vector(keys_out_ptr)); } struct SegmentedSort_KeyValuePairs_Fixture_Tag; From fa60f52ed24056870f9091b82f7f56c23cc6cfc2 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 3 Sep 2025 00:26:01 +0000 Subject: [PATCH 024/100] Change policies to make them work from c.parallel --- cub/cub/agent/agent_radix_sort_downsweep.cuh | 2 ++ cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 4 ++-- cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh | 7 ++++--- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/cub/cub/agent/agent_radix_sort_downsweep.cuh b/cub/cub/agent/agent_radix_sort_downsweep.cuh index a7440e7f850..e3e7499b928 100644 --- a/cub/cub/agent/agent_radix_sort_downsweep.cuh +++ b/cub/cub/agent/agent_radix_sort_downsweep.cuh @@ -131,6 +131,8 @@ namespace detail CUB_DETAIL_POLICY_WRAPPER_DEFINE( RadixSortDownsweepAgentPolicy, (GenericAgentPolicy), + (BLOCK_THREADS, BlockThreads, int), + (ITEMS_PER_THREAD, ItemsPerThread, int), (RADIX_BITS, RadixBits, int), (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index baaf419e138..d41383b6925 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -903,8 +903,8 @@ private: { cudaError_t error = cudaSuccess; - const auto blocks_in_grid = static_cast(num_segments); - constexpr auto threads_in_block = static_cast(wrapped_policy.LargeSegment().BlockThreads()); + const auto blocks_in_grid = static_cast(num_segments); + const auto threads_in_block = static_cast(wrapped_policy.LargeSegment().BlockThreads()); // Log kernel configuration #ifdef CUB_DEBUG_LOG diff --git a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index 5bac6740507..becb8d09cbc 100644 --- a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -136,9 +136,10 @@ struct SegmentedSortPolicyWrapper() = LargeSegment().EncodedPolicy(), - key<"SmallSegmentPolicy">() = SmallSegment().EncodedPolicy(), - key<"MediumSegmentPolicy">() = MediumSegment().EncodedPolicy()>(); + return object() = LargeSegment().EncodedPolicy(), + key<"SmallSegmentPolicy">() = SmallSegment().EncodedPolicy(), + key<"MediumSegmentPolicy">() = MediumSegment().EncodedPolicy(), + key<"PartitioningThreshold">() = value()>(); } #endif }; From bb8eb8fc33cadbdf109f27c8178e6e7e12e5f01b Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 3 Sep 2025 00:26:20 +0000 Subject: [PATCH 025/100] Add one more level for preprocessor for each --- libcudacxx/include/cuda/std/__cccl/preprocessor.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libcudacxx/include/cuda/std/__cccl/preprocessor.h b/libcudacxx/include/cuda/std/__cccl/preprocessor.h index e769b0a25d6..88cca9af23c 100644 --- a/libcudacxx/include/cuda/std/__cccl/preprocessor.h +++ b/libcudacxx/include/cuda/std/__cccl/preprocessor.h @@ -99,6 +99,8 @@ #define _CCCL_PP_FOR_EACH_7(_Mp, _1, _2, _3, _4, _5, _6, _7) _Mp(_1) _Mp(_2) _Mp(_3) _Mp(_4) _Mp(_5) _Mp(_6) _Mp(_7) #define _CCCL_PP_FOR_EACH_8(_Mp, _1, _2, _3, _4, _5, _6, _7, _8) \ _Mp(_1) _Mp(_2) _Mp(_3) _Mp(_4) _Mp(_5) _Mp(_6) _Mp(_7) _Mp(_8) +#define _CCCL_PP_FOR_EACH_9(_Mp, _1, _2, _3, _4, _5, _6, _7, _8, _9) \ + _Mp(_1) _Mp(_2) _Mp(_3) _Mp(_4) _Mp(_5) _Mp(_6) _Mp(_7) _Mp(_8) _Mp(_9) #define _CCCL_PP_PROBE_EMPTY_PROBE__CCCL_PP_PROBE_EMPTY _CCCL_PP_PROBE(~) From 750c3fc8f998f20a6693800290815eba1f0257e8 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 3 Sep 2025 17:25:25 +0000 Subject: [PATCH 026/100] Account for double buffers and sort order and use test tuple params similar to radix sort tests --- c/parallel/include/cccl/c/segmented_sort.h | 5 + c/parallel/src/segmented_sort.cu | 97 ++++-- c/parallel/test/test_segmented_sort.cpp | 370 +++++++++++++++------ 3 files changed, 357 insertions(+), 115 deletions(-) diff --git a/c/parallel/include/cccl/c/segmented_sort.h b/c/parallel/include/cccl/c/segmented_sort.h index a4efe792a14..00ca9a8821d 100644 --- a/c/parallel/include/cccl/c/segmented_sort.h +++ b/c/parallel/include/cccl/c/segmented_sort.h @@ -15,6 +15,7 @@ #endif // !CCCL_C_EXPERIMENTAL #include +#include #include #include @@ -36,11 +37,13 @@ typedef struct cccl_device_segmented_sort_build_result_t CUkernel three_way_partition_kernel; void* runtime_policy; void* partition_runtime_policy; + cccl_sort_order_t order; } cccl_device_segmented_sort_build_result_t; // TODO return a union of nvtx/cuda/nvrtc errors or a string? CCCL_C_API CUresult cccl_device_segmented_sort_build( cccl_device_segmented_sort_build_result_t* build, + cccl_sort_order_t sort_order, cccl_iterator_t d_keys_in, cccl_iterator_t d_keys_out, cccl_iterator_t d_values_in, @@ -66,6 +69,8 @@ CCCL_C_API CUresult cccl_device_segmented_sort( int64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, + bool is_overwrite_okay, + int* selector, CUstream stream); CCCL_C_API CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_result_t* bld_ptr); diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index edad1e0c833..c29a9acea07 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -53,7 +53,8 @@ std::string get_device_segmented_sort_fallback_kernel_name( std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t, std::string_view key_t, - std::string_view value_t) + std::string_view value_t, + cccl_sort_order_t sort_order) { std::string chained_policy_t; check(nvrtcGetTypeName(&chained_policy_t)); @@ -72,8 +73,8 @@ std::string get_device_segmented_sort_fallback_kernel_name( DeviceSegmentedSortFallbackKernel(...); */ return std::format( - "cub::detail::segmented_sort::DeviceSegmentedSortFallbackKernel", + "cub::detail::segmented_sort::DeviceSegmentedSortFallbackKernel<{0}, {1}, {2}, {3}, {4}, {5}, {6}>", + (sort_order == CCCL_ASCENDING) ? "cub::SortOrder::Ascending" : "cub::SortOrder::Descending", chained_policy_t, // 0 key_t, // 1 value_t, // 2 @@ -88,7 +89,8 @@ std::string get_device_segmented_sort_kernel_small_name( std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t, std::string_view key_t, - std::string_view value_t) + std::string_view value_t, + cccl_sort_order_t sort_order) { std::string chained_policy_t; check(nvrtcGetTypeName(&chained_policy_t)); @@ -107,8 +109,8 @@ std::string get_device_segmented_sort_kernel_small_name( DeviceSegmentedSortKernelSmall(...); */ return std::format( - "cub::detail::segmented_sort::DeviceSegmentedSortKernelSmall", + "cub::detail::segmented_sort::DeviceSegmentedSortKernelSmall<{0}, {1}, {2}, {3}, {4}, {5}, {6}>", + (sort_order == CCCL_ASCENDING) ? "cub::SortOrder::Ascending" : "cub::SortOrder::Descending", chained_policy_t, // 0 key_t, // 1 value_t, // 2 @@ -123,7 +125,8 @@ std::string get_device_segmented_sort_kernel_large_name( std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t, std::string_view key_t, - std::string_view value_t) + std::string_view value_t, + cccl_sort_order_t sort_order) { std::string chained_policy_t; check(nvrtcGetTypeName(&chained_policy_t)); @@ -142,8 +145,8 @@ std::string get_device_segmented_sort_kernel_large_name( DeviceSegmentedSortKernelLarge(...); */ return std::format( - "cub::detail::segmented_sort::DeviceSegmentedSortKernelLarge", + "cub::detail::segmented_sort::DeviceSegmentedSortKernelLarge<{0}, {1}, {2}, {3}, {4}, {5}, {6}>", + (sort_order == CCCL_ASCENDING) ? "cub::SortOrder::Ascending" : "cub::SortOrder::Descending", chained_policy_t, // 0 key_t, // 1 value_t, // 2 @@ -305,6 +308,7 @@ struct segmented_sort_end_offset_iterator_tag; CUresult cccl_device_segmented_sort_build( cccl_device_segmented_sort_build_result_t* build_ptr, + cccl_sort_order_t sort_order, cccl_iterator_t keys_in_it, cccl_iterator_t keys_out_it, cccl_iterator_t values_in_it, @@ -473,7 +477,8 @@ struct device_segmented_sort_policy {{ start_offset_iterator_name, end_offset_iterator_name, key_t, - value_t); + value_t, + sort_order); std::string segmented_sort_kernel_small_name = segmented_sort::get_device_segmented_sort_kernel_small_name( keys_in_iterator_name, @@ -481,7 +486,8 @@ struct device_segmented_sort_policy {{ start_offset_iterator_name, end_offset_iterator_name, key_t, - value_t); + value_t, + sort_order); std::string segmented_sort_kernel_large_name = segmented_sort::get_device_segmented_sort_kernel_large_name( keys_in_iterator_name, @@ -489,7 +495,8 @@ struct device_segmented_sort_policy {{ start_offset_iterator_name, end_offset_iterator_name, key_t, - value_t); + value_t, + sort_order); std::string segmented_sort_fallback_kernel_lowered_name; std::string segmented_sort_kernel_small_lowered_name; @@ -557,6 +564,7 @@ struct device_segmented_sort_policy {{ // Use the runtime policy extracted via from_json build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_tuning_policy{ large_segment_policy, small_segment_policy, medium_segment_policy, partitioning_threshold}; + build_ptr->order = sort_order; } catch (const std::exception& exc) { @@ -569,7 +577,8 @@ struct device_segmented_sort_policy {{ return error; } -CUresult cccl_device_segmented_sort( +template +CUresult cccl_device_segmented_sort_impl( cccl_device_segmented_sort_build_result_t build, void* d_temp_storage, size_t* temp_storage_bytes, @@ -581,6 +590,8 @@ CUresult cccl_device_segmented_sort( int64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, + bool is_overwrite_okay, + int* selector, CUstream stream) { bool pushed = false; @@ -594,13 +605,18 @@ CUresult cccl_device_segmented_sort( // Create DoubleBuffer structures for keys and values // CUB will handle keys-only vs key-value sorting internally - auto d_keys_double_buffer = cub::DoubleBuffer( - static_cast(d_keys_in.state), static_cast(d_keys_out.state)); - auto d_values_double_buffer = cub::DoubleBuffer( - static_cast(d_values_in.state), static_cast(d_values_out.state)); + indirect_arg_t key_arg_in{d_keys_in}; + indirect_arg_t key_arg_out{d_keys_out}; + cub::DoubleBuffer d_keys_double_buffer( + *static_cast(&key_arg_in), *static_cast(&key_arg_out)); + + indirect_arg_t val_arg_in{d_values_in}; + indirect_arg_t val_arg_out{d_values_out}; + cub::DoubleBuffer d_values_double_buffer( + *static_cast(&val_arg_in), *static_cast(&val_arg_out)); auto exec_status = cub::DispatchSegmentedSort< - cub::SortOrder::Ascending, + Order, indirect_arg_t, // KeyT indirect_arg_t, // ValueT OffsetT, // OffsetT @@ -620,7 +636,7 @@ CUresult cccl_device_segmented_sort( num_segments, indirect_iterator_t{start_offset_in}, indirect_iterator_t{end_offset_in}, - true, // is_overwrite_okay + is_overwrite_okay, stream, /* kernel_source */ {build}, /* partition_kernel_source */ {build}, @@ -629,6 +645,11 @@ CUresult cccl_device_segmented_sort( /* partition_policy */ *reinterpret_cast(build.partition_runtime_policy)); + if (selector != nullptr) + { + *selector = d_keys_double_buffer.selector; + } + error = static_cast(exec_status); } catch (const std::exception& exc) @@ -648,6 +669,44 @@ CUresult cccl_device_segmented_sort( return error; } +CUresult cccl_device_segmented_sort( + cccl_device_segmented_sort_build_result_t build, + void* d_temp_storage, + size_t* temp_storage_bytes, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_keys_out, + cccl_iterator_t d_values_in, + cccl_iterator_t d_values_out, + int64_t num_items, + int64_t num_segments, + cccl_iterator_t start_offset_in, + cccl_iterator_t end_offset_in, + bool is_overwrite_okay, + int* selector, + CUstream stream) +{ + auto segmented_sort_impl = + (build.order == CCCL_ASCENDING) + ? cccl_device_segmented_sort_impl + : cccl_device_segmented_sort_impl; + + return segmented_sort_impl( + build, + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + start_offset_in, + end_offset_in, + is_overwrite_okay, + selector, + stream); +} + CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_result_t* build_ptr) { try diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 5a91ce1747e..99f13b58903 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -41,6 +41,31 @@ struct segmented_sort_cleanup using segmented_sort_deleter = BuildResultDeleter; using segmented_sort_build_cache_t = build_cache_t>; +template +struct TestParameters +{ + using KeyT = KeyTy; + static constexpr bool m_descending = descending; + static constexpr bool m_overwrite_okay = overwrite_okay; + + constexpr TestParameters() {} + + bool is_descending() const + { + return m_descending; + } + bool is_overwrite_okay() const + { + return m_overwrite_okay; + } +}; + +using test_params_tuple = + c2h::type_list, false, false>, + TestParameters, true, false>, + TestParameters, false, true>, + TestParameters, true, true>>; + template auto& get_cache() { @@ -51,6 +76,7 @@ struct segmented_sort_build { CUresult operator()( BuildResultT* build_ptr, + cccl_sort_order_t sort_order, cccl_iterator_t keys_in, cccl_iterator_t keys_out, cccl_iterator_t values_in, @@ -59,6 +85,8 @@ struct segmented_sort_build int64_t /*num_segments*/, cccl_iterator_t start_offsets, cccl_iterator_t end_offsets, + bool /*is_overwrite_okay*/, + int* /*selector*/, int cc_major, int cc_minor, const char* cub_path, @@ -68,6 +96,7 @@ struct segmented_sort_build { return cccl_device_segmented_sort_build( build_ptr, + sort_order, keys_in, keys_out, values_in, @@ -85,15 +114,41 @@ struct segmented_sort_build struct segmented_sort_run { - template - CUresult operator()(Ts... args) const noexcept + template + CUresult operator()( + BuildResultT build, + void* temp_storage, + size_t* temp_storage_bytes, + cccl_sort_order_t, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_keys_out, + cccl_iterator_t d_values_in, + cccl_iterator_t d_values_out, + int64_t num_items, + int64_t num_segments, + cccl_iterator_t start_offsets, + cccl_iterator_t end_offsets, + Rest... rest) const noexcept { - return cccl_device_segmented_sort(args...); + return cccl_device_segmented_sort( + build, + temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + start_offsets, + end_offsets, + rest...); } }; template void segmented_sort( + cccl_sort_order_t sort_order, cccl_iterator_t keys_in, cccl_iterator_t keys_out, cccl_iterator_t values_in, @@ -102,11 +157,25 @@ void segmented_sort( int64_t num_segments, cccl_iterator_t start_offsets, cccl_iterator_t end_offsets, + bool is_overwrite_okay, + int* selector, std::optional& cache, const std::optional& lookup_key) { AlgorithmExecute( - cache, lookup_key, keys_in, keys_out, values_in, values_out, num_items, num_segments, start_offsets, end_offsets); + cache, + lookup_key, + sort_order, + keys_in, + keys_out, + values_in, + values_out, + num_items, + num_segments, + start_offsets, + end_offsets, + is_overwrite_okay, + selector); } // ============== @@ -152,9 +221,16 @@ extern "C" __device__ {2} {0}({1}* state) } struct SegmentedSort_KeysOnly_Fixture_Tag; -C2H_TEST("segmented_sort can sort keys-only with integral types", "[segmented_sort][keys_only]", key_types) +C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", test_params_tuple) { - using KeyT = c2h::get<0, TestType>; + using T = c2h::get<0, TestType>; + using key_t = typename T::KeyT; + constexpr auto this_test_params = T(); + const bool is_descending = this_test_params.is_descending(); + const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); + int selector = -1; + // generate choices for n_segments: 0, 13 and 2 random samples from [50, 200) const std::size_t n_segments = GENERATE(0, 13, take(2, random(50, 200))); // generate choices for segment size: 1, 20 and random samples @@ -163,14 +239,14 @@ C2H_TEST("segmented_sort can sort keys-only with integral types", "[segmented_so const std::size_t n_elems = n_segments * segment_size; std::vector host_keys_int = generate(n_elems); - std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); - std::vector host_keys_out(n_elems); + std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); + std::vector host_keys_out(n_elems); REQUIRE(host_keys.size() == n_elems); REQUIRE(host_keys_out.size() == n_elems); - pointer_t keys_in_ptr(host_keys); // copy from host to device - pointer_t keys_out_ptr(host_keys_out); // copy from host to device + pointer_t keys_in_ptr(host_keys); + pointer_t keys_out_ptr(host_keys_out); pointer_t values_in; pointer_t values_out; @@ -209,10 +285,15 @@ C2H_TEST("segmented_sort can sort keys-only with integral types", "[segmented_so end_offset_it.state.linear_id = 1; end_offset_it.state.segment_size = segment_size; - auto& build_cache = get_cache(); - const auto& test_key = make_key(); + auto& build_cache = get_cache(); + const std::string& key_string = KeyBuilder::join( + {KeyBuilder::bool_as_key(is_descending), + KeyBuilder::type_as_key(), + KeyBuilder::bool_as_key(is_overwrite_okay)}); + const auto& test_key = std::make_optional(key_string); segmented_sort( + order, keys_in_ptr, keys_out_ptr, values_in, @@ -221,29 +302,43 @@ C2H_TEST("segmented_sort can sort keys-only with integral types", "[segmented_so n_segments, start_offset_it, end_offset_it, + is_overwrite_okay, + &selector, build_cache, test_key); // Create expected result by sorting each segment - std::vector expected_keys = host_keys; + std::vector expected_keys = host_keys; for (std::size_t i = 0; i < n_segments; ++i) { std::size_t segment_start = i * segment_size; std::size_t segment_end = segment_start + segment_size; - std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end); + if (is_descending) + { + std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end, std::greater()); + } + else + { + std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end); + } } - REQUIRE(expected_keys == std::vector(keys_out_ptr)); + auto& output_keys = (is_overwrite_okay && selector == 0) ? keys_in_ptr : keys_out_ptr; + REQUIRE(expected_keys == std::vector(output_keys)); } struct SegmentedSort_KeyValuePairs_Fixture_Tag; -C2H_TEST_LIST("segmented_sort can sort key-value pairs with integral types", - "[segmented_sort][key_value]", - std::int32_t, - std::int64_t, - std::uint32_t, - std::uint64_t) +C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value]", test_params_tuple) { + using T = c2h::get<0, TestType>; + using key_t = typename T::KeyT; + + constexpr auto this_test_params = T(); + const bool is_descending = this_test_params.is_descending(); + const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); + int selector = -1; + // generate choices for n_segments: 0, 10 and random samples const std::size_t n_segments = GENERATE(0, 10, take(2, random(30, 100))); // generate choices for segment size @@ -251,18 +346,27 @@ C2H_TEST_LIST("segmented_sort can sort key-value pairs with integral types", const std::size_t n_elems = n_segments * segment_size; - std::vector host_keys = generate(n_elems); - std::vector host_values = generate(n_elems); - std::vector host_keys_out(n_elems); - std::vector host_values_out(n_elems); + std::vector host_keys_int = generate(n_elems); + std::vector host_keys(n_elems); + std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int x) { + return static_cast(x); + }); + std::vector host_values_int = generate(n_elems); + std::vector host_values(n_elems); + std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { + return static_cast(x); + }); + + std::vector host_keys_out(n_elems); + std::vector host_values_out(n_elems); REQUIRE(host_keys.size() == n_elems); REQUIRE(host_values.size() == n_elems); - pointer_t keys_in_ptr(host_keys); - pointer_t keys_out_ptr(host_keys_out); - pointer_t values_in_ptr(host_values); - pointer_t values_out_ptr(host_values_out); + pointer_t keys_in_ptr(host_keys); + pointer_t keys_out_ptr(host_keys_out); + pointer_t values_in_ptr(host_values); + pointer_t values_out_ptr(host_values_out); using SizeT = unsigned long long; static constexpr std::string_view index_ty_name = "unsigned long long"; @@ -297,10 +401,16 @@ C2H_TEST_LIST("segmented_sort can sort key-value pairs with integral types", end_offset_it.state.linear_id = 1; end_offset_it.state.segment_size = segment_size; - auto& build_cache = get_cache(); - const auto& test_key = make_key(); + auto& build_cache = get_cache(); + const std::string& key_string = KeyBuilder::join( + {KeyBuilder::bool_as_key(is_descending), + KeyBuilder::type_as_key(), + KeyBuilder::type_as_key(), + KeyBuilder::bool_as_key(is_overwrite_okay)}); + const auto& test_key = std::make_optional(key_string); segmented_sort( + order, keys_in_ptr, keys_out_ptr, values_in_ptr, @@ -309,29 +419,42 @@ C2H_TEST_LIST("segmented_sort can sort key-value pairs with integral types", n_segments, start_offset_it, end_offset_it, + is_overwrite_okay, + &selector, build_cache, test_key); // Create expected result by sorting each segment with key-value pairs - std::vector> key_value_pairs; + std::vector> key_value_pairs; for (std::size_t i = 0; i < n_elems; ++i) { key_value_pairs.emplace_back(host_keys[i], host_values[i]); } - std::vector expected_keys(n_elems); - std::vector expected_values(n_elems); + std::vector expected_keys(n_elems); + std::vector expected_values(n_elems); for (std::size_t i = 0; i < n_segments; ++i) { std::size_t segment_start = i * segment_size; std::size_t segment_end = segment_start + segment_size; - // Sort this segment by key - std::sort( - key_value_pairs.begin() + segment_start, key_value_pairs.begin() + segment_end, [](const auto& a, const auto& b) { - return a.first < b.first; - }); + if (is_descending) + { + std::sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return b.first < a.first; + }); + } + else + { + std::sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return a.first < b.first; + }); + } // Extract sorted keys and values for (std::size_t j = segment_start; j < segment_end; ++j) @@ -341,8 +464,10 @@ C2H_TEST_LIST("segmented_sort can sort key-value pairs with integral types", } } - REQUIRE(expected_keys == std::vector(keys_out_ptr)); - REQUIRE(expected_values == std::vector(values_out_ptr)); + auto& output_keys = (is_overwrite_okay && selector == 0) ? keys_in_ptr : keys_out_ptr; + auto& output_vals = (is_overwrite_okay && selector == 0) ? values_in_ptr : values_out_ptr; + REQUIRE(expected_keys == std::vector(output_keys)); + REQUIRE(expected_values == std::vector(output_vals)); } struct custom_pair @@ -362,17 +487,23 @@ struct custom_pair }; struct SegmentedSort_CustomTypes_Fixture_Tag; -C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][custom_types]") +C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][custom_types]", test_params_tuple) { - using KeyT = custom_pair; - using ValueT = float; + using T = c2h::get<0, TestType>; + using key_t = custom_pair; + + constexpr auto this_test_params = T(); + const bool is_descending = this_test_params.is_descending(); + const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); + int selector = -1; const std::size_t n_segments = 25; const std::size_t segment_size = 20; const std::size_t n_elems = n_segments * segment_size; // Generate custom key data - std::vector host_keys(n_elems); + std::vector host_keys(n_elems); for (std::size_t i = 0; i < n_elems; ++i) { host_keys[i] = custom_pair{static_cast(i % 1000), static_cast(i % 100)}; @@ -380,19 +511,19 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust // Generate float values by first generating ints and then transforming std::vector host_values_int = generate(n_elems); - std::vector host_values(n_elems); + std::vector host_values(n_elems); std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { - return static_cast(x); + return static_cast(x); }); - std::vector host_keys_out(n_elems); - std::vector host_values_out(n_elems); + std::vector host_keys_out(n_elems); + std::vector host_values_out(n_elems); - pointer_t keys_in_ptr(host_keys); - pointer_t keys_out_ptr(host_keys_out); - pointer_t values_in_ptr(host_values); - pointer_t values_out_ptr(host_values_out); + pointer_t keys_in_ptr(host_keys); + pointer_t keys_out_ptr(host_keys_out); + pointer_t values_in_ptr(host_values); + pointer_t values_out_ptr(host_values_out); - using SizeT = ::cuda::std::size_t; + using SizeT = cuda::std::size_t; std::vector segments(n_segments + 1); for (std::size_t i = 0; i <= n_segments; ++i) { @@ -405,10 +536,16 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust auto end_offset_it = start_offset_it; end_offset_it.state = offset_ptr.ptr + 1; - auto& build_cache = get_cache(); - const auto& test_key = make_key(); + auto& build_cache = get_cache(); + const std::string& key_string = KeyBuilder::join( + {KeyBuilder::bool_as_key(is_descending), + KeyBuilder::type_as_key(), + KeyBuilder::type_as_key(), + KeyBuilder::bool_as_key(is_overwrite_okay)}); + const auto& test_key = std::make_optional(key_string); segmented_sort( + order, keys_in_ptr, keys_out_ptr, values_in_ptr, @@ -417,29 +554,42 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust n_segments, start_offset_it, end_offset_it, + is_overwrite_okay, + &selector, build_cache, test_key); // Create expected result - std::vector> key_value_pairs; + std::vector> key_value_pairs; for (std::size_t i = 0; i < n_elems; ++i) { key_value_pairs.emplace_back(host_keys[i], host_values[i]); } - std::vector expected_keys(n_elems); - std::vector expected_values(n_elems); + std::vector expected_keys(n_elems); + std::vector expected_values(n_elems); for (std::size_t i = 0; i < n_segments; ++i) { std::size_t segment_start = segments[i]; std::size_t segment_end = segments[i + 1]; - // Sort this segment by key - std::sort( - key_value_pairs.begin() + segment_start, key_value_pairs.begin() + segment_end, [](const auto& a, const auto& b) { - return a.first < b.first; - }); + if (is_descending) + { + std::sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return b.first < a.first; + }); + } + else + { + std::sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return a.first < b.first; + }); + } // Extract sorted keys and values for (std::size_t j = segment_start; j < segment_end; ++j) @@ -449,11 +599,10 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust } } - auto result_keys = std::vector(keys_out_ptr); - auto result_values = std::vector(values_out_ptr); - - REQUIRE(expected_keys == result_keys); - REQUIRE(expected_values == result_values); + auto& output_keys = (is_overwrite_okay && selector == 0) ? keys_in_ptr : keys_out_ptr; + auto& output_vals = (is_overwrite_okay && selector == 0) ? values_in_ptr : values_out_ptr; + REQUIRE(expected_keys == std::vector(output_keys)); + REQUIRE(expected_values == std::vector(output_vals)); } using SizeT = unsigned long long; @@ -491,10 +640,16 @@ extern "C" __device__ unsigned long long dereference_variable_offset_it(variable } struct SegmentedSort_VariableSegments_Fixture_Tag; -C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][variable_segments]") +C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][variable_segments]", test_params_tuple) { - using KeyT = std::int32_t; - using ValueT = float; + using T = c2h::get<0, TestType>; + using key_t = std::int32_t; + + constexpr auto this_test_params = T(); + const bool is_descending = this_test_params.is_descending(); + const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); + int selector = -1; const std::size_t n_segments = 20; @@ -504,20 +659,25 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va std::size_t n_elems = std::accumulate(segment_sizes.begin(), segment_sizes.end(), 0ULL); - std::vector host_keys = generate(n_elems); + std::vector host_keys_int = generate(n_elems); + std::vector host_keys(n_elems); + std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int x) { + return static_cast(x); + }); + // Generate float values by first generating ints and then transforming std::vector host_values_int = generate(n_elems); - std::vector host_values(n_elems); + std::vector host_values(n_elems); std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { - return static_cast(x); + return static_cast(x); }); - std::vector host_keys_out(n_elems); - std::vector host_values_out(n_elems); + std::vector host_keys_out(n_elems); + std::vector host_values_out(n_elems); - pointer_t keys_in_ptr(host_keys); - pointer_t keys_out_ptr(host_keys_out); - pointer_t values_in_ptr(host_values); - pointer_t values_out_ptr(host_values_out); + pointer_t keys_in_ptr(host_keys); + pointer_t keys_out_ptr(host_keys_out); + pointer_t values_in_ptr(host_values); + pointer_t values_out_ptr(host_values_out); // Create segment offset arrays std::vector start_offsets(n_segments); @@ -554,10 +714,16 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va end_offset_it.state.linear_id = 0; end_offset_it.state.offsets = end_offsets_ptr.ptr; - auto& build_cache = get_cache(); - const auto& test_key = make_key(); + auto& build_cache = get_cache(); + const std::string& key_string = KeyBuilder::join( + {KeyBuilder::bool_as_key(is_descending), + KeyBuilder::type_as_key(), + KeyBuilder::type_as_key(), + KeyBuilder::bool_as_key(is_overwrite_okay)}); + const auto& test_key = std::make_optional(key_string); segmented_sort( + order, keys_in_ptr, keys_out_ptr, values_in_ptr, @@ -566,29 +732,42 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va n_segments, start_offset_it, end_offset_it, + is_overwrite_okay, + &selector, build_cache, test_key); // Create expected result - std::vector> key_value_pairs; + std::vector> key_value_pairs; for (std::size_t i = 0; i < n_elems; ++i) { key_value_pairs.emplace_back(host_keys[i], host_values[i]); } - std::vector expected_keys(n_elems); - std::vector expected_values(n_elems); + std::vector expected_keys(n_elems); + std::vector expected_values(n_elems); for (std::size_t i = 0; i < n_segments; ++i) { std::size_t segment_start = start_offsets[i]; std::size_t segment_end = end_offsets[i]; - // Sort this segment by key - std::sort( - key_value_pairs.begin() + segment_start, key_value_pairs.begin() + segment_end, [](const auto& a, const auto& b) { - return a.first < b.first; - }); + if (is_descending) + { + std::sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return b.first < a.first; + }); + } + else + { + std::sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return a.first < b.first; + }); + } // Extract sorted keys and values for (std::size_t j = segment_start; j < segment_end; ++j) @@ -598,9 +777,8 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va } } - auto result_keys = std::vector(keys_out_ptr); - auto result_values = std::vector(values_out_ptr); - - REQUIRE(expected_keys == result_keys); - REQUIRE(expected_values == result_values); + auto& output_keys = (is_overwrite_okay && selector == 0) ? keys_in_ptr : keys_out_ptr; + auto& output_vals = (is_overwrite_okay && selector == 0) ? values_in_ptr : values_out_ptr; + REQUIRE(expected_keys == std::vector(output_keys)); + REQUIRE(expected_values == std::vector(output_vals)); } From 5fddefe603fb01bb9bd61997d8c74764a81cf38a Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 3 Sep 2025 17:25:51 +0000 Subject: [PATCH 027/100] Add missing include --- cub/cub/agent/agent_radix_sort_downsweep.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/cub/cub/agent/agent_radix_sort_downsweep.cuh b/cub/cub/agent/agent_radix_sort_downsweep.cuh index e3e7499b928..5e1749633b0 100644 --- a/cub/cub/agent/agent_radix_sort_downsweep.cuh +++ b/cub/cub/agent/agent_radix_sort_downsweep.cuh @@ -51,6 +51,7 @@ #include #include #include +#include #include #include From bfb47e5a6094f066abb602a554778997aeed8e83 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 4 Sep 2025 18:53:22 +0000 Subject: [PATCH 028/100] Pass partition max policy to three way partition dispatch --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index d41383b6925..40566abe7d8 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -501,7 +501,8 @@ struct DispatchSegmentedSort max_num_segments_per_invocation, stream, partition_kernel_source, - launcher_factory); + launcher_factory, + partition_max_policy); device_partition_temp_storage.grow(three_way_partition_temp_storage_bytes); } @@ -783,7 +784,8 @@ private: current_num_segments, stream, partition_kernel_source, - launcher_factory); + launcher_factory, + partition_max_policy); if (cudaSuccess != error) { From 389f9b38ee25515bb7aabde28cfc1f250ec20340 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 4 Sep 2025 18:57:21 +0000 Subject: [PATCH 029/100] Move streaming_context_t and other typedefs to three way partition kernels file --- .../dispatch/dispatch_three_way_partition.cuh | 79 ------------------ .../dispatch/kernels/three_way_partition.cuh | 80 +++++++++++++++++++ 2 files changed, 80 insertions(+), 79 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index 3be8a562c49..6164237bc97 100644 --- a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -93,85 +93,6 @@ struct DeviceThreeWayPartitionKernelSource return sizeof(OffsetT); } }; - -// Offset type used to instantiate the stream three-way-partition-kernel and agent to index the items within one -// partition -using per_partition_offset_t = ::cuda::std::int32_t; - -using AccumPackHelperT = detail::three_way_partition::accumulator_pack_t; -using AccumPackT = typename AccumPackHelperT::pack_t; -using ScanTileStateT = cub::ScanTileState; - -template -class streaming_context_t -{ -private: - bool first_partition = true; - bool last_partition = false; - TotalNumItemsT total_previous_num_items{}; - - // We use a double-buffer for keeping track of the number of previously selected items - TotalNumItemsT* d_num_selected_in = nullptr; - TotalNumItemsT* d_num_selected_out = nullptr; - -public: - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE - streaming_context_t(TotalNumItemsT* d_num_selected_in, TotalNumItemsT* d_num_selected_out, bool is_last_partition) - : last_partition(is_last_partition) - , d_num_selected_in(d_num_selected_in) - , d_num_selected_out(d_num_selected_out) - {} - - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(TotalNumItemsT num_items, bool next_partition_is_the_last) - { - ::cuda::std::swap(d_num_selected_in, d_num_selected_out); - first_partition = false; - last_partition = next_partition_is_the_last; - total_previous_num_items += num_items; - }; - - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT input_offset() const - { - return first_partition ? TotalNumItemsT{0} : total_previous_num_items; - }; - - _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_first() const - { - return first_partition ? TotalNumItemsT{0} : d_num_selected_in[0]; - }; - - _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_second() const - { - return first_partition ? TotalNumItemsT{0} : d_num_selected_in[1]; - }; - - _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_rejected() const - { - return first_partition ? TotalNumItemsT{0} : d_num_selected_in[2]; - ; - }; - - template - _CCCL_DEVICE _CCCL_FORCEINLINE void update_num_selected( - NumSelectedIteratorT user_num_selected_out_it, - TotalNumItemsT num_selected_first, - TotalNumItemsT num_selected_second, - TotalNumItemsT num_items_in_partition) const - { - if (last_partition) - { - user_num_selected_out_it[0] = num_previously_selected_first() + num_selected_first; - user_num_selected_out_it[1] = num_previously_selected_second() + num_selected_second; - } - else - { - d_num_selected_out[0] = num_previously_selected_first() + num_selected_first; - d_num_selected_out[1] = num_previously_selected_second() + num_selected_second; - d_num_selected_out[2] = - num_previously_rejected() + (num_items_in_partition - num_selected_second - num_selected_first); - } - } -}; } // namespace detail::three_way_partition /****************************************************************************** diff --git a/cub/cub/device/dispatch/kernels/three_way_partition.cuh b/cub/cub/device/dispatch/kernels/three_way_partition.cuh index 9ea919f7110..5335b0e09d8 100644 --- a/cub/cub/device/dispatch/kernels/three_way_partition.cuh +++ b/cub/cub/device/dispatch/kernels/three_way_partition.cuh @@ -19,6 +19,86 @@ CUB_NAMESPACE_BEGIN namespace detail::three_way_partition { + +// Offset type used to instantiate the stream three-way-partition-kernel and agent to index the items within one +// partition +using per_partition_offset_t = ::cuda::std::int32_t; + +using AccumPackHelperT = detail::three_way_partition::accumulator_pack_t; +using AccumPackT = typename AccumPackHelperT::pack_t; +using ScanTileStateT = cub::ScanTileState; + +template +class streaming_context_t +{ +private: + bool first_partition = true; + bool last_partition = false; + TotalNumItemsT total_previous_num_items{}; + + // We use a double-buffer for keeping track of the number of previously selected items + TotalNumItemsT* d_num_selected_in = nullptr; + TotalNumItemsT* d_num_selected_out = nullptr; + +public: + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE + streaming_context_t(TotalNumItemsT* d_num_selected_in, TotalNumItemsT* d_num_selected_out, bool is_last_partition) + : last_partition(is_last_partition) + , d_num_selected_in(d_num_selected_in) + , d_num_selected_out(d_num_selected_out) + {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(TotalNumItemsT num_items, bool next_partition_is_the_last) + { + ::cuda::std::swap(d_num_selected_in, d_num_selected_out); + first_partition = false; + last_partition = next_partition_is_the_last; + total_previous_num_items += num_items; + }; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT input_offset() const + { + return first_partition ? TotalNumItemsT{0} : total_previous_num_items; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_first() const + { + return first_partition ? TotalNumItemsT{0} : d_num_selected_in[0]; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_second() const + { + return first_partition ? TotalNumItemsT{0} : d_num_selected_in[1]; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_rejected() const + { + return first_partition ? TotalNumItemsT{0} : d_num_selected_in[2]; + ; + }; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void update_num_selected( + NumSelectedIteratorT user_num_selected_out_it, + TotalNumItemsT num_selected_first, + TotalNumItemsT num_selected_second, + TotalNumItemsT num_items_in_partition) const + { + if (last_partition) + { + user_num_selected_out_it[0] = num_previously_selected_first() + num_selected_first; + user_num_selected_out_it[1] = num_previously_selected_second() + num_selected_second; + } + else + { + d_num_selected_out[0] = num_previously_selected_first() + num_selected_first; + d_num_selected_out[1] = num_previously_selected_second() + num_selected_second; + d_num_selected_out[2] = + num_previously_rejected() + (num_items_in_partition - num_selected_second - num_selected_first); + } + } +}; + /****************************************************************************** * Kernel entry points *****************************************************************************/ From 5842b63cd8854a788d7bf23a6215653442eb887c Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 4 Sep 2025 18:58:47 +0000 Subject: [PATCH 030/100] Encode delay contructor info into ptx-json. This requires defining new types that hold this info and adding support for unsigned int to ptx-json --- cub/cub/agent/single_pass_scan_operators.cuh | 30 +++++++++++++++++++ cub/cub/detail/ptx-json/value.h | 5 ++-- .../tuning/tuning_three_way_partition.cuh | 11 +++++++ 3 files changed, 44 insertions(+), 2 deletions(-) diff --git a/cub/cub/agent/single_pass_scan_operators.cuh b/cub/cub/agent/single_pass_scan_operators.cuh index 9f937ddee62..9e093608b4a 100644 --- a/cub/cub/agent/single_pass_scan_operators.cuh +++ b/cub/cub/agent/single_pass_scan_operators.cuh @@ -497,6 +497,36 @@ using default_reduce_by_key_delay_constructor_t = reduce_by_key_delay_constructor_t<350, 450>, default_delay_constructor_t>>; +#if defined(CUB_ENABLE_POLICY_PTX_JSON) +# include + +// ptx-json encoders for delay constructor types. Unlike the other agent policy +// member variables, this is defined as a type alias so we can't use the +// CUB_DETAIL_POLICY_WRAPPER_DEFINE macro to embed it with ptx-json. To work +// around this, we define the ptx-json encoders here. These can then be used in +// the policy wrapper's EncodedPolicy member function to explicitly encode the +// delay constructor. + +template +struct delay_constructor_json; + +template +struct delay_constructor_json> +{ + using type = + ptx_json::object() = ptx_json::value(), + ptx_json::key<"delay">() = ptx_json::value(), + ptx_json::key<"l2_write_latency">() = ptx_json::value()>; +}; + +template +struct delay_constructor_json> +{ + using type = ptx_json::object() = ptx_json::value(), + ptx_json::key<"l2_write_latency">() = ptx_json::value()>; +}; +#endif // CUB_ENABLE_POLICY_PTX_JSON + /** * @brief Alias template for a ScanTileState specialized for a given value type, `T`, and memory order `Order`. * diff --git a/cub/cub/detail/ptx-json/value.h b/cub/cub/detail/ptx-json/value.h index 1cc8537c8ee..89153d00954 100644 --- a/cub/cub/detail/ptx-json/value.h +++ b/cub/cub/detail/ptx-json/value.h @@ -70,8 +70,9 @@ struct value } }; -template -struct value +// Integral constants (matches both signed and unsigned integrals) +template +struct value, void>> { __forceinline__ __device__ static void emit() { diff --git a/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh b/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh index 8bb69b9fbaa..a1ab2b68451 100644 --- a/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh @@ -69,6 +69,17 @@ struct ThreeWayPartitionPolicyWrapper() = ThreeWayPartition().EncodedPolicy(), + key<"ThreeWayPartitionPolicyDelayConstructor">() = + typename detail::delay_constructor_json::type()>(); + } +#endif }; template From ce9ee5e07b8a955b393f6c3b8547275f426c88a9 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 4 Sep 2025 22:27:30 +0000 Subject: [PATCH 031/100] Get key size from kernel source instead of sizeof directly --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 40566abe7d8..787920a7951 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -277,6 +277,11 @@ struct DeviceSegmentedSortKernelSource CUB_DEFINE_KERNEL_GETTER( SegmentedSortKernelLarge, DeviceSegmentedSortKernelLarge); + + CUB_RUNTIME_FUNCTION static constexpr size_t KeySize() + { + return sizeof(KeyT); + } }; } // namespace detail::segmented_sort @@ -680,7 +685,7 @@ private: CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE int GetNumPasses(int radix_bits) { constexpr int byte_size = 8; - constexpr int num_bits = sizeof(KeyT) * byte_size; + const int num_bits = kernel_source.KeySize() * byte_size; const int num_passes = ::cuda::ceil_div(num_bits, radix_bits); return num_passes; } From 11a8007232a9f205a4e3efa98f0414e8dd5d1b9a Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 4 Sep 2025 22:30:50 +0000 Subject: [PATCH 032/100] Move segmented test utils to common header --- c/parallel/test/test_segmented_reduce.cpp | 58 +----------------- c/parallel/test/test_util.h | 71 +++++++++++++++++++++++ 2 files changed, 73 insertions(+), 56 deletions(-) diff --git a/c/parallel/test/test_segmented_reduce.cpp b/c/parallel/test/test_segmented_reduce.cpp index b15860d4e2e..a3e22f8d5a9 100644 --- a/c/parallel/test/test_segmented_reduce.cpp +++ b/c/parallel/test/test_segmented_reduce.cpp @@ -99,44 +99,6 @@ void segmented_reduce( // Test section // ============== -static std::tuple make_step_counting_iterator_sources( - std::string_view index_ty_name, - std::string_view state_name, - std::string_view advance_fn_name, - std::string_view dereference_fn_name) -{ - static constexpr std::string_view it_state_src_tmpl = R"XXX( -struct {0} {{ - {1} linear_id; - {1} row_size; -}}; -)XXX"; - - const std::string it_state_def_src = std::format(it_state_src_tmpl, state_name, index_ty_name); - - static constexpr std::string_view it_def_src_tmpl = R"XXX( -extern "C" __device__ void {0}({1}* state, {2} offset) -{{ - state->linear_id += offset; -}} -)XXX"; - - const std::string it_advance_fn_def_src = - std::format(it_def_src_tmpl, /*0*/ advance_fn_name, state_name, index_ty_name); - - static constexpr std::string_view it_deref_src_tmpl = R"XXX( -extern "C" __device__ {2} {0}({1}* state) -{{ - return (state->linear_id) * (state->row_size); -}} -)XXX"; - - const std::string it_deref_fn_def_src = - std::format(it_deref_src_tmpl, dereference_fn_name, state_name, index_ty_name); - - return std::make_tuple(it_state_def_src, it_advance_fn_def_src, it_deref_fn_def_src); -} - struct SegmentedReduce_SumOverRows_Fixture_Tag; C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type", "[segmented_reduce]", @@ -619,22 +581,6 @@ struct host_check_functor_state DataT* m_ptr; }; -template -void host_advance_transform_it_state(void* state, cccl_increment_t offset) -{ - auto st = reinterpret_cast(state); - using IndexT = decltype(st->base_it_state.value); - - if constexpr (std::is_signed_v) - { - st->base_it_state.value += offset.signed_offset; - } - else - { - st->base_it_state.value += offset.unsigned_offset; - } -} - namespace validate { @@ -833,8 +779,8 @@ extern "C" __device__ void {0}(const void *x1_p, const void *x2_p, void *out_p) auto cccl_end_offsets_it = static_cast(end_offsets_it); // set host_advance functions - cccl_start_offsets_it.host_advance = &host_advance_transform_it_state; - cccl_end_offsets_it.host_advance = &host_advance_transform_it_state; + cccl_start_offsets_it.host_advance = &host_advance_base_value; + cccl_end_offsets_it.host_advance = &host_advance_base_value; value_t h_init{DataT{0}}; diff --git a/c/parallel/test/test_util.h b/c/parallel/test/test_util.h index 7a9650b07f9..bdae495ad9c 100644 --- a/c/parallel/test/test_util.h +++ b/c/parallel/test/test_util.h @@ -1014,6 +1014,77 @@ inline std::tuple make_reverse_iterator_s return std::make_tuple(iterator_state_src, advance_fn_src, dereference_fn_src); } +// Common iterator helpers shared by segmented tests +inline std::tuple make_step_counting_iterator_sources( + std::string_view index_ty_name, + std::string_view state_name, + std::string_view advance_fn_name, + std::string_view dereference_fn_name) +{ + static constexpr std::string_view it_state_src_tmpl = R"XXX( +struct {0} {{ + {1} linear_id; + {1} row_size; +}}; +)XXX"; + + const std::string it_state_def_src = std::format(it_state_src_tmpl, state_name, index_ty_name); + + static constexpr std::string_view it_def_src_tmpl = R"XXX( +extern "C" __device__ void {0}({1}* state, {2} offset) +{{ + state->linear_id += offset; +}} +)XXX"; + + const std::string it_advance_fn_def_src = + std::format(it_def_src_tmpl, /*0*/ advance_fn_name, state_name, index_ty_name); + + static constexpr std::string_view it_deref_src_tmpl = R"XXX( +extern "C" __device__ {2} {0}({1}* state) +{{ + return (state->linear_id) * (state->row_size); +}} +)XXX"; + + const std::string it_deref_fn_def_src = + std::format(it_deref_src_tmpl, dereference_fn_name, state_name, index_ty_name); + + return std::make_tuple(it_state_def_src, it_advance_fn_def_src, it_deref_fn_def_src); +} + +// Host-side advance function for iterator states that have a `linear_id` member +template +inline void host_advance_linear_id(void* state, cccl_increment_t offset) +{ + auto* st = reinterpret_cast(state); + using Index = decltype(st->linear_id); + if constexpr (std::is_signed_v) + { + st->linear_id += offset.signed_offset; + } + else + { + st->linear_id += offset.unsigned_offset; + } +} + +// Host-side advance for iterator states that contain a nested `base_it_state.value` +template +inline void host_advance_base_value(void* state, cccl_increment_t offset) +{ + auto st = reinterpret_cast(state); + using IndexT = decltype(st->base_it_state.value); + if constexpr (std::is_signed_v) + { + st->base_it_state.value += offset.signed_offset; + } + else + { + st->base_it_state.value += offset.unsigned_offset; + } +} + template iterator_t> make_reverse_iterator( iterator_kind kind, std::string_view value_type, std::string_view prefix = "", std::string_view transform = "") From 5bae49431c14117c030ce78996e870b75562c1cd Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 4 Sep 2025 22:32:16 +0000 Subject: [PATCH 033/100] Add three way partition kernels and policy and implement KeySize --- c/parallel/include/cccl/c/segmented_sort.h | 1 + c/parallel/src/segmented_sort.cu | 197 +++++++++++++++++---- 2 files changed, 164 insertions(+), 34 deletions(-) diff --git a/c/parallel/include/cccl/c/segmented_sort.h b/c/parallel/include/cccl/c/segmented_sort.h index 00ca9a8821d..1950f481455 100644 --- a/c/parallel/include/cccl/c/segmented_sort.h +++ b/c/parallel/include/cccl/c/segmented_sort.h @@ -29,6 +29,7 @@ typedef struct cccl_device_segmented_sort_build_result_t void* cubin; size_t cubin_size; CUlibrary library; + cccl_type_info key_type; cccl_type_info offset_type; CUkernel segmented_sort_fallback_kernel; CUkernel segmented_sort_kernel_small; diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index c29a9acea07..14fb00dbb66 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -38,6 +38,7 @@ #include struct device_segmented_sort_policy; +struct device_three_way_partition_policy; using OffsetT = long; static_assert(std::is_same_v, OffsetT>, "OffsetT must be long"); @@ -48,8 +49,6 @@ static_assert(sizeof(OffsetT) == sizeof(cuda::std::int64_t)); namespace segmented_sort { std::string get_device_segmented_sort_fallback_kernel_name( - std::string_view /* key_iterator_t */, - std::string_view /* value_iterator_t */, std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t, std::string_view key_t, @@ -84,8 +83,6 @@ std::string get_device_segmented_sort_fallback_kernel_name( } std::string get_device_segmented_sort_kernel_small_name( - std::string_view /* key_iterator_t */, - std::string_view /* value_iterator_t */, std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t, std::string_view key_t, @@ -120,8 +117,6 @@ std::string get_device_segmented_sort_kernel_small_name( } std::string get_device_segmented_sort_kernel_large_name( - std::string_view /* key_iterator_t */, - std::string_view /* value_iterator_t */, std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t, std::string_view key_t, @@ -171,8 +166,73 @@ struct segmented_sort_kernel_source { return build.segmented_sort_kernel_large; } + + std::size_t KeySize() const + { + return build.key_type.size; + } }; +std::string get_three_way_partition_init_kernel_name() +{ + constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT"; + + constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + + return std::format("cub::detail::three_way_partition::DeviceThreeWayPartitionInitKernel<{0}, {1}>", + scan_tile_state_t, // 0 + num_selected_it_t); // 1 +} + +std::string +get_three_way_partition_kernel_name(std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t) +{ + std::string chained_policy_t; + check(nvrtcGetTypeName(&chained_policy_t)); + + constexpr std::string_view input_it_t = + "thrust::counting_iterator"; + constexpr std::string_view first_out_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + constexpr std::string_view second_out_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + constexpr std::string_view unselected_out_it_t = + "thrust::reverse_iterator"; + constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT"; + std::string offset_t; + check(nvrtcGetTypeName(&offset_t)); + + std::string select_first_part_op_t = std::format( + "cub::detail::segmented_sort::LargeSegmentsSelectorT<{0}, {1}, {2}>", + offset_t, // 0 + start_offset_iterator_t, // 1 + end_offset_iterator_t); // 2 + + std::string select_second_part_op_t = std::format( + "cub::detail::segmented_sort::SmallSegmentsSelectorT<{0}, {1}, {2}>", + offset_t, // 0 + start_offset_iterator_t, // 1 + end_offset_iterator_t); // 2 + + constexpr std::string_view per_partition_offset_t = "cub::detail::three_way_partition::per_partition_offset_t"; + constexpr std::string_view streaming_context_t = + "cub::detail::three_way_partition::streaming_context_t"; + + return std::format( + "cub::detail::three_way_partition::DeviceThreeWayPartitionKernel<{0}, {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}, " + "{10}>", + chained_policy_t, // 0 (ChainedPolicyT) + input_it_t, // 1 (InputIteratorT) + first_out_it_t, // 2 (FirstOutputIteratorT) + second_out_it_t, // 3 (SecondOutputIteratorT) + unselected_out_it_t, // 4 (UnselectedOutputIteratorT) + num_selected_it_t, // 5 (NumSelectedIteratorT) + scan_tile_state_t, // 6 (ScanTileStateT) + select_first_part_op_t, // 7 (SelectFirstPartOp) + select_second_part_op_t, // 8 (SelectSecondPartOp) + per_partition_offset_t, // 9 (OffsetT) + streaming_context_t); // 10 (StreamingContextT) +} + struct partition_kernel_source { cccl_device_segmented_sort_build_result_t& build; @@ -297,6 +357,42 @@ struct partition_runtime_tuning_policy return op.template Invoke(*this); } }; + +std::string get_three_way_partition_policy_delay_constructor(const nlohmann::json& partition_policy) +{ + auto dc_json = partition_policy["ThreeWayPartitionPolicyDelayConstructor"]; // optional; not used further + auto delay_constructor_type = dc_json["type"].get(); + + if (delay_constructor_type == "fixed_delay_constructor_t") + { + auto delay = dc_json["delay"].get(); + auto l2_write_latency = dc_json["l2_write_latency"].get(); + return std::format("cub::detail::fixed_delay_constructor_t<{}, {}>", delay, l2_write_latency); + } + else if (delay_constructor_type == "no_delay_constructor_t") + { + auto l2_write_latency = dc_json["l2_write_latency"].get(); + return std::format("cub::detail::no_delay_constructor_t<{}>", l2_write_latency); + } + throw std::runtime_error("Invalid delay constructor type: " + delay_constructor_type); +} + +std::string inject_delay_constructor_into_three_way_policy( + const std::string& three_way_partition_policy_str, const std::string& delay_constructor_type) +{ + // Insert before the final closing of the struct (right before the sequence "};") + const std::string needle = "};"; + const auto pos = three_way_partition_policy_str.rfind(needle); + if (pos == std::string::npos) + { + return three_way_partition_policy_str; // unexpected; return as-is + } + const std::string insertion = + std::format("\n struct detail {{ using delay_constructor_t = {}; }}; \n", delay_constructor_type); + std::string out = three_way_partition_policy_str; + out.insert(pos, insertion); + return out; +} } // namespace segmented_sort struct segmented_sort_keys_input_iterator_tag; @@ -385,7 +481,7 @@ CUresult cccl_device_segmented_sort_build( const auto [end_offset_iterator_name, end_offset_iterator_src] = get_specialization(template_id(), end_offset_it); - const auto offset_t = cccl_type_enum_to_name(cccl_type_enum::CCCL_UINT64); + const auto offset_t = cccl_type_enum_to_name(cccl_type_enum::CCCL_INT64); const std::string key_t = cccl_type_enum_to_name(keys_in_it.value_type.type); const std::string value_t = keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); @@ -427,6 +523,7 @@ struct __align__({1}) storage_t {{ static constexpr std::string_view ptx_query_tu_src_tmpl = R"XXXX( #include +#include {0} {1} )XXXX"; @@ -447,11 +544,35 @@ struct __align__({1}) storage_t {{ auto [medium_segment_policy, medium_segment_policy_str] = RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "MediumSegmentPolicy"); - auto partitioning_threshold = static_cast(runtime_policy["PartitioningThreshold"].get()); + auto partitioning_threshold = runtime_policy["PartitioningThreshold"].get(); + + static constexpr std::string_view partition_policy_wrapper_expr_tmpl = + R"XXXX(cub::detail::three_way_partition::MakeThreeWayPartitionPolicyWrapper(cub::detail::three_way_partition::policy_hub<{0}, {1}>::MaxPolicy::ActivePolicy{{}}))XXXX"; + const auto partition_policy_wrapper_expr = std::format( + partition_policy_wrapper_expr_tmpl, + "::cuda::std::uint32_t", // This is local_segment_index_t defined in segmented_sort.cuh + "::cuda::std::int32_t"); // This is per_partition_offset_t defined in segmented_sort.cuh + + nlohmann::json partition_policy = get_policy(partition_policy_wrapper_expr, ptx_query_tu_src, ptx_args); + + using cub::detail::RuntimeThreeWayPartitionAgentPolicy; + auto [three_way_partition_policy, three_way_partition_policy_str] = + RuntimeThreeWayPartitionAgentPolicy::from_json(partition_policy, "ThreeWayPartitionPolicy"); + + const std::string three_way_partition_policy_delay_constructor = + segmented_sort::get_three_way_partition_policy_delay_constructor(partition_policy); + + // Inject delay constructor alias into the ThreeWayPartitionPolicy struct string + const std::string injected_three_way_partition_policy_str = + segmented_sort::inject_delay_constructor_into_three_way_policy( + three_way_partition_policy_str, three_way_partition_policy_delay_constructor); - // agent_policy_t is to specify parameters like policy_hub does in dispatch_segmented_sort.cuh constexpr std::string_view program_preamble_template = R"XXX( #include +#include +#include // used in three_way_partition kernel +#include // used in three_way_partition kernel +#include // used in three_way_partition kernel {0} {1} struct device_segmented_sort_policy {{ @@ -461,6 +582,11 @@ struct device_segmented_sort_policy {{ {4} }}; }}; +struct device_three_way_partition_policy {{ + struct ActivePolicy {{ + {5} + }}; +}}; )XXX"; std::string final_src = std::format( @@ -469,38 +595,28 @@ struct device_segmented_sort_policy {{ dependent_definitions_src, // 1 large_segment_policy_str, // 2 small_segment_policy_str, // 3 - medium_segment_policy_str); // 4 + medium_segment_policy_str, // 4 + injected_three_way_partition_policy_str); // 5 std::string segmented_sort_fallback_kernel_name = segmented_sort::get_device_segmented_sort_fallback_kernel_name( - keys_in_iterator_name, - values_in_iterator_name, - start_offset_iterator_name, - end_offset_iterator_name, - key_t, - value_t, - sort_order); + start_offset_iterator_name, end_offset_iterator_name, key_t, value_t, sort_order); std::string segmented_sort_kernel_small_name = segmented_sort::get_device_segmented_sort_kernel_small_name( - keys_in_iterator_name, - values_in_iterator_name, - start_offset_iterator_name, - end_offset_iterator_name, - key_t, - value_t, - sort_order); + start_offset_iterator_name, end_offset_iterator_name, key_t, value_t, sort_order); std::string segmented_sort_kernel_large_name = segmented_sort::get_device_segmented_sort_kernel_large_name( - keys_in_iterator_name, - values_in_iterator_name, - start_offset_iterator_name, - end_offset_iterator_name, - key_t, - value_t, - sort_order); + start_offset_iterator_name, end_offset_iterator_name, key_t, value_t, sort_order); + + std::string three_way_partition_init_kernel_name = segmented_sort::get_three_way_partition_init_kernel_name(); + + std::string three_way_partition_kernel_name = + segmented_sort::get_three_way_partition_kernel_name(start_offset_iterator_name, end_offset_iterator_name); std::string segmented_sort_fallback_kernel_lowered_name; std::string segmented_sort_kernel_small_lowered_name; std::string segmented_sort_kernel_large_lowered_name; + std::string three_way_partition_init_kernel_lowered_name; + std::string three_way_partition_kernel_lowered_name; const std::string arch = std::format("-arch=sm_{0}{1}", cc_major, cc_minor); @@ -540,10 +656,14 @@ struct device_segmented_sort_policy {{ ->add_expression({segmented_sort_fallback_kernel_name}) ->add_expression({segmented_sort_kernel_small_name}) ->add_expression({segmented_sort_kernel_large_name}) + ->add_expression({three_way_partition_init_kernel_name}) + ->add_expression({three_way_partition_kernel_name}) ->compile_program({args, num_args}) ->get_name({segmented_sort_fallback_kernel_name, segmented_sort_fallback_kernel_lowered_name}) ->get_name({segmented_sort_kernel_small_name, segmented_sort_kernel_small_lowered_name}) ->get_name({segmented_sort_kernel_large_name, segmented_sort_kernel_large_lowered_name}) + ->get_name({three_way_partition_init_kernel_name, three_way_partition_init_kernel_lowered_name}) + ->get_name({three_way_partition_kernel_name, three_way_partition_kernel_lowered_name}) ->link_program() ->add_link_list(ltoir_list) ->finalize_program(); @@ -557,13 +677,22 @@ struct device_segmented_sort_policy {{ &build_ptr->segmented_sort_kernel_small, build_ptr->library, segmented_sort_kernel_small_lowered_name.c_str())); check(cuLibraryGetKernel( &build_ptr->segmented_sort_kernel_large, build_ptr->library, segmented_sort_kernel_large_lowered_name.c_str())); + check(cuLibraryGetKernel(&build_ptr->three_way_partition_init_kernel, + build_ptr->library, + three_way_partition_init_kernel_lowered_name.c_str())); + check(cuLibraryGetKernel( + &build_ptr->three_way_partition_kernel, build_ptr->library, three_way_partition_kernel_lowered_name.c_str())); - build_ptr->cc = cc; - build_ptr->cubin = (void*) result.data.release(); - build_ptr->cubin_size = result.size; + build_ptr->cc = cc; + build_ptr->cubin = (void*) result.data.release(); + build_ptr->cubin_size = result.size; + build_ptr->key_type = keys_in_it.value_type; + build_ptr->offset_type = cccl_type_info{sizeof(OffsetT), alignof(OffsetT), cccl_type_enum::CCCL_INT64}; // Use the runtime policy extracted via from_json build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_tuning_policy{ large_segment_policy, small_segment_policy, medium_segment_policy, partitioning_threshold}; + build_ptr->partition_runtime_policy = + new segmented_sort::partition_runtime_tuning_policy{three_way_partition_policy}; build_ptr->order = sort_order; } catch (const std::exception& exc) From c75d306ccb00c772c4de36756fdebae93c1df612 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 4 Sep 2025 22:33:28 +0000 Subject: [PATCH 034/100] Use common utils from test_util --- c/parallel/test/test_segmented_sort.cpp | 54 +++++-------------------- 1 file changed, 10 insertions(+), 44 deletions(-) diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 99f13b58903..0702692f246 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -182,44 +182,6 @@ void segmented_sort( // Test section // ============== -static std::tuple make_step_counting_iterator_sources( - std::string_view index_ty_name, - std::string_view state_name, - std::string_view advance_fn_name, - std::string_view dereference_fn_name) -{ - static constexpr std::string_view it_state_src_tmpl = R"XXX( -struct {0} {{ - {1} linear_id; - {1} row_size; -}}; -)XXX"; - - const std::string it_state_def_src = std::format(it_state_src_tmpl, state_name, index_ty_name); - - static constexpr std::string_view it_def_src_tmpl = R"XXX( -extern "C" __device__ void {0}({1}* state, {2} offset) -{{ - state->linear_id += offset; -}} -)XXX"; - - const std::string it_advance_fn_def_src = - std::format(it_def_src_tmpl, /*0*/ advance_fn_name, state_name, index_ty_name); - - static constexpr std::string_view it_deref_src_tmpl = R"XXX( -extern "C" __device__ {2} {0}({1}* state) -{{ - return (state->linear_id) * (state->row_size); -}} -)XXX"; - - const std::string it_deref_fn_def_src = - std::format(it_deref_src_tmpl, dereference_fn_name, state_name, index_ty_name); - - return std::make_tuple(it_state_def_src, it_advance_fn_def_src, it_deref_fn_def_src); -} - struct SegmentedSort_KeysOnly_Fixture_Tag; C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", test_params_tuple) { @@ -231,10 +193,8 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); int selector = -1; - // generate choices for n_segments: 0, 13 and 2 random samples from [50, 200) - const std::size_t n_segments = GENERATE(0, 13, take(2, random(50, 200))); - // generate choices for segment size: 1, 20 and random samples - const std::size_t segment_size = GENERATE(1, 20, take(2, random(10, 100))); + const std::size_t n_segments = GENERATE(0, 13, take(2, random(1 << 10, 1 << 12))); + const std::size_t segment_size = GENERATE(1, 12, take(2, random(1 << 10, 1 << 12))); const std::size_t n_elems = n_segments * segment_size; @@ -285,6 +245,12 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes end_offset_it.state.linear_id = 1; end_offset_it.state.segment_size = segment_size; + // Provide host-advance callbacks for offset iterators + auto start_offsets_cccl = static_cast(start_offset_it); + auto end_offsets_cccl = static_cast(end_offset_it); + start_offsets_cccl.host_advance = &host_advance_linear_id; + end_offsets_cccl.host_advance = &host_advance_linear_id; + auto& build_cache = get_cache(); const std::string& key_string = KeyBuilder::join( {KeyBuilder::bool_as_key(is_descending), @@ -300,8 +266,8 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes values_out, n_elems, n_segments, - start_offset_it, - end_offset_it, + start_offsets_cccl, + end_offsets_cccl, is_overwrite_okay, &selector, build_cache, From 04a4a6de649af584d945b9e7dd9a4cd7a9fdd5bf Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 5 Sep 2025 17:04:41 +0000 Subject: [PATCH 035/100] Delete partition tuning policy --- c/parallel/src/segmented_sort.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 14fb00dbb66..89201f9a5ea 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -360,7 +360,7 @@ struct partition_runtime_tuning_policy std::string get_three_way_partition_policy_delay_constructor(const nlohmann::json& partition_policy) { - auto dc_json = partition_policy["ThreeWayPartitionPolicyDelayConstructor"]; // optional; not used further + auto dc_json = partition_policy["ThreeWayPartitionPolicyDelayConstructor"]; auto delay_constructor_type = dc_json["type"].get(); if (delay_constructor_type == "fixed_delay_constructor_t") @@ -850,6 +850,7 @@ CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_res // Clean up the runtime policy delete static_cast(build_ptr->runtime_policy); + delete static_cast(build_ptr->partition_runtime_policy); check(cuLibraryUnload(build_ptr->library)); } catch (const std::exception& exc) From 686ea77b3291c9e6a0645a0674b924e04587bdf5 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sat, 6 Sep 2025 13:32:00 +0000 Subject: [PATCH 036/100] Make parameter tuple member functions constexpr --- c/parallel/test/test_segmented_sort.cpp | 85 +++++++++++++------------ 1 file changed, 46 insertions(+), 39 deletions(-) diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 0702692f246..e3e53c3aca4 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -50,11 +50,11 @@ struct TestParameters constexpr TestParameters() {} - bool is_descending() const + constexpr bool is_descending() const { return m_descending; } - bool is_overwrite_okay() const + constexpr bool is_overwrite_okay() const { return m_overwrite_okay; } @@ -185,13 +185,13 @@ void segmented_sort( struct SegmentedSort_KeysOnly_Fixture_Tag; C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", test_params_tuple) { - using T = c2h::get<0, TestType>; - using key_t = typename T::KeyT; - constexpr auto this_test_params = T(); - const bool is_descending = this_test_params.is_descending(); - const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; - const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - int selector = -1; + using T = c2h::get<0, TestType>; + using key_t = typename T::KeyT; + + constexpr auto this_test_params = T(); + constexpr bool is_descending = this_test_params.is_descending(); + constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); const std::size_t n_segments = GENERATE(0, 13, take(2, random(1 << 10, 1 << 12))); const std::size_t segment_size = GENERATE(1, 12, take(2, random(1 << 10, 1 << 12))); @@ -211,8 +211,9 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes pointer_t values_in; pointer_t values_out; - using SizeT = unsigned long long; - static constexpr std::string_view index_ty_name = "unsigned long long"; + // Always use signed long long for offset iterator since negative advances are possible + using SizeT = signed long long; + static constexpr std::string_view index_ty_name = "signed long long"; struct segment_offset_iterator_state_t { @@ -258,6 +259,8 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes KeyBuilder::bool_as_key(is_overwrite_okay)}); const auto& test_key = std::make_optional(key_string); + int selector = -1; + segmented_sort( order, keys_in_ptr, @@ -299,11 +302,10 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] using T = c2h::get<0, TestType>; using key_t = typename T::KeyT; - constexpr auto this_test_params = T(); - const bool is_descending = this_test_params.is_descending(); - const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; - const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - int selector = -1; + constexpr auto this_test_params = T(); + constexpr bool is_descending = this_test_params.is_descending(); + constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); // generate choices for n_segments: 0, 10 and random samples const std::size_t n_segments = GENERATE(0, 10, take(2, random(30, 100))); @@ -313,15 +315,9 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] const std::size_t n_elems = n_segments * segment_size; std::vector host_keys_int = generate(n_elems); - std::vector host_keys(n_elems); - std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int x) { - return static_cast(x); - }); + std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); std::vector host_values_int = generate(n_elems); - std::vector host_values(n_elems); - std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { - return static_cast(x); - }); + std::vector host_values(host_values_int.begin(), host_values_int.end()); std::vector host_keys_out(n_elems); std::vector host_values_out(n_elems); @@ -331,11 +327,12 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] pointer_t keys_in_ptr(host_keys); pointer_t keys_out_ptr(host_keys_out); + pointer_t values_in_ptr(host_values); pointer_t values_out_ptr(host_values_out); - using SizeT = unsigned long long; - static constexpr std::string_view index_ty_name = "unsigned long long"; + using SizeT = signed long long; + static constexpr std::string_view index_ty_name = "signed long long"; struct segment_offset_iterator_state_t { @@ -367,6 +364,12 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] end_offset_it.state.linear_id = 1; end_offset_it.state.segment_size = segment_size; + // Provide host-advance callbacks for offset iterators + auto start_offsets_cccl = static_cast(start_offset_it); + auto end_offsets_cccl = static_cast(end_offset_it); + start_offsets_cccl.host_advance = &host_advance_linear_id; + end_offsets_cccl.host_advance = &host_advance_linear_id; + auto& build_cache = get_cache(); const std::string& key_string = KeyBuilder::join( {KeyBuilder::bool_as_key(is_descending), @@ -375,6 +378,8 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] KeyBuilder::bool_as_key(is_overwrite_okay)}); const auto& test_key = std::make_optional(key_string); + int selector = -1; + segmented_sort( order, keys_in_ptr, @@ -383,15 +388,15 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] values_out_ptr, n_elems, n_segments, - start_offset_it, - end_offset_it, + start_offsets_cccl, + end_offsets_cccl, is_overwrite_okay, &selector, build_cache, test_key); // Create expected result by sorting each segment with key-value pairs - std::vector> key_value_pairs; + std::vector> key_value_pairs(n_elems); for (std::size_t i = 0; i < n_elems; ++i) { key_value_pairs.emplace_back(host_keys[i], host_values[i]); @@ -458,11 +463,10 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust using T = c2h::get<0, TestType>; using key_t = custom_pair; - constexpr auto this_test_params = T(); - const bool is_descending = this_test_params.is_descending(); - const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; - const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - int selector = -1; + constexpr auto this_test_params = T(); + constexpr bool is_descending = this_test_params.is_descending(); + constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); const std::size_t n_segments = 25; const std::size_t segment_size = 20; @@ -510,6 +514,8 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust KeyBuilder::bool_as_key(is_overwrite_okay)}); const auto& test_key = std::make_optional(key_string); + int selector = -1; + segmented_sort( order, keys_in_ptr, @@ -611,11 +617,10 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va using T = c2h::get<0, TestType>; using key_t = std::int32_t; - constexpr auto this_test_params = T(); - const bool is_descending = this_test_params.is_descending(); - const auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; - const bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - int selector = -1; + constexpr auto this_test_params = T(); + constexpr bool is_descending = this_test_params.is_descending(); + constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; + constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); const std::size_t n_segments = 20; @@ -688,6 +693,8 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va KeyBuilder::bool_as_key(is_overwrite_okay)}); const auto& test_key = std::make_optional(key_string); + int selector = -1; + segmented_sort( order, keys_in_ptr, From be17ec0e4c13f42571638fa70694737dab129f40 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sat, 6 Sep 2025 19:13:10 +0000 Subject: [PATCH 037/100] Rename row_size to segment_size and fix error in key value pair correctness check --- c/parallel/test/test_segmented_reduce.cpp | 32 +++++++++++------------ c/parallel/test/test_segmented_sort.cpp | 27 +++++++++---------- c/parallel/test/test_util.h | 4 +-- 3 files changed, 31 insertions(+), 32 deletions(-) diff --git a/c/parallel/test/test_segmented_reduce.cpp b/c/parallel/test/test_segmented_reduce.cpp index a3e22f8d5a9..c65260676f1 100644 --- a/c/parallel/test/test_segmented_reduce.cpp +++ b/c/parallel/test/test_segmented_reduce.cpp @@ -112,8 +112,8 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type", // generate 4 choices for number of columns const std::size_t n_cols = GENERATE(0, 12, take(2, random(1 << 10, 1 << 12))); - const std::size_t n_elems = n_rows * n_cols; - const std::size_t row_size = n_cols; + const std::size_t n_elems = n_rows * n_cols; + const std::size_t segment_size = n_cols; const std::vector host_input = generate(n_elems); std::vector host_output(n_rows, 0); @@ -130,7 +130,7 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type", struct row_offset_iterator_state_t { SizeT linear_id; - SizeT row_size; + SizeT segment_size; }; static constexpr std::string_view offset_iterator_state_name = "row_offset_iterator_state_t"; @@ -146,16 +146,16 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type", {advance_offset_method_name, offset_iterator_advance_src}, {deref_offset_method_name, offset_iterator_deref_src}); - start_offset_it.state.linear_id = 0; - start_offset_it.state.row_size = row_size; + start_offset_it.state.linear_id = 0; + start_offset_it.state.segment_size = segment_size; // a copy of offset iterator, so no need to define advance/dereference bodies, // just reused those defined above iterator_t end_offset_it = make_iterator( {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); - end_offset_it.state.linear_id = 1; - end_offset_it.state.row_size = row_size; + end_offset_it.state.linear_id = 1; + end_offset_it.state.segment_size = segment_size; operation_t op = make_operation("op", get_reduce_op(get_type_info().type)); value_t init{0}; @@ -170,7 +170,7 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type", for (std::size_t i = 0; i < n_rows; ++i) { - std::size_t row_offset = i * row_size; + std::size_t row_offset = i * segment_size; host_output_it[i] = std::reduce(host_input_it + row_offset, host_input_it + (row_offset + n_cols)); } REQUIRE(host_output == std::vector(output_ptr)); @@ -190,8 +190,8 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type " // generate 4 choices for number of columns const std::size_t n_cols = GENERATE(0, 12, take(2, random(1 << 10, 1 << 12))); - const std::size_t n_elems = n_rows * n_cols; - const std::size_t row_size = n_cols; + const std::size_t n_elems = n_rows * n_cols; + const std::size_t segment_size = n_cols; const std::vector host_input = generate(n_elems); std::vector host_output(n_rows, 0); @@ -208,7 +208,7 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type " struct row_offset_iterator_state_t { SizeT linear_id; - SizeT row_size; + SizeT segment_size; }; static constexpr std::string_view offset_iterator_state_name = "row_offset_iterator_state_t"; @@ -224,16 +224,16 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type " {advance_offset_method_name, offset_iterator_advance_src}, {deref_offset_method_name, offset_iterator_deref_src}); - start_offset_it.state.linear_id = 0; - start_offset_it.state.row_size = row_size; + start_offset_it.state.linear_id = 0; + start_offset_it.state.segment_size = segment_size; // a copy of offset iterator, so no need to define advance/dereference bodies, // just reused those defined above iterator_t end_offset_it = make_iterator( {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); - end_offset_it.state.linear_id = 1; - end_offset_it.state.row_size = row_size; + end_offset_it.state.linear_id = 1; + end_offset_it.state.segment_size = segment_size; cccl_op_t op = make_well_known_binary_operation(); value_t init{0}; @@ -248,7 +248,7 @@ C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type " for (std::size_t i = 0; i < n_rows; ++i) { - std::size_t row_offset = i * row_size; + std::size_t row_offset = i * segment_size; host_output_it[i] = std::reduce(host_input_it + row_offset, host_input_it + (row_offset + n_cols)); } REQUIRE(host_output == std::vector(output_ptr)); diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index e3e53c3aca4..289f0406390 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -307,9 +307,7 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - // generate choices for n_segments: 0, 10 and random samples - const std::size_t n_segments = GENERATE(0, 10, take(2, random(30, 100))); - // generate choices for segment size + const std::size_t n_segments = GENERATE(0, 10, take(2, random(30, 100))); const std::size_t segment_size = GENERATE(1, 15, take(2, random(5, 50))); const std::size_t n_elems = n_segments * segment_size; @@ -396,7 +394,8 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] test_key); // Create expected result by sorting each segment with key-value pairs - std::vector> key_value_pairs(n_elems); + std::vector> key_value_pairs; + key_value_pairs.reserve(n_elems); for (std::size_t i = 0; i < n_elems; ++i) { key_value_pairs.emplace_back(host_keys[i], host_values[i]); @@ -412,19 +411,19 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] if (is_descending) { - std::sort(key_value_pairs.begin() + segment_start, - key_value_pairs.begin() + segment_end, - [](const auto& a, const auto& b) { - return b.first < a.first; - }); + std::stable_sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return b.first < a.first; + }); } else { - std::sort(key_value_pairs.begin() + segment_start, - key_value_pairs.begin() + segment_end, - [](const auto& a, const auto& b) { - return a.first < b.first; - }); + std::stable_sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return a.first < b.first; + }); } // Extract sorted keys and values diff --git a/c/parallel/test/test_util.h b/c/parallel/test/test_util.h index bdae495ad9c..357f28c42b3 100644 --- a/c/parallel/test/test_util.h +++ b/c/parallel/test/test_util.h @@ -1024,7 +1024,7 @@ inline std::tuple make_step_counting_iter static constexpr std::string_view it_state_src_tmpl = R"XXX( struct {0} {{ {1} linear_id; - {1} row_size; + {1} segment_size; }}; )XXX"; @@ -1043,7 +1043,7 @@ extern "C" __device__ void {0}({1}* state, {2} offset) static constexpr std::string_view it_deref_src_tmpl = R"XXX( extern "C" __device__ {2} {0}({1}* state) {{ - return (state->linear_id) * (state->row_size); + return (state->linear_id) * (state->segment_size); }} )XXX"; From c8f3d5fa4782df32da6555c166f17849b252f0f9 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 10 Sep 2025 18:36:50 +0000 Subject: [PATCH 038/100] Allow passing in custom types as items and pass segment selectors through kernel source --- c/parallel/src/segmented_sort.cu | 68 ++++++++++++++++++++++++++------ 1 file changed, 56 insertions(+), 12 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 89201f9a5ea..a325e40fced 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -171,6 +171,23 @@ struct segmented_sort_kernel_source { return build.key_type.size; } + + using LargeSegmentsSelectorT = cub::detail::segmented_sort::LargeSegmentsSelectorT; + using SmallSegmentsSelectorT = cub::detail::segmented_sort::SmallSegmentsSelectorT; + + auto LargeSegmentsSelector( + OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const + { + return LargeSegmentsSelectorT( + offset, *reinterpret_cast(begin_offset_iterator.ptr), *reinterpret_cast(end_offset_iterator.ptr)); + } + + auto SmallSegmentsSelector( + OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const + { + return SmallSegmentsSelectorT( + offset, *reinterpret_cast(begin_offset_iterator.ptr), *reinterpret_cast(end_offset_iterator.ptr)); + } }; std::string get_three_way_partition_init_kernel_name() @@ -465,6 +482,27 @@ CUresult cccl_device_segmented_sort_build( template_id(), values_out_it, values_in_it.value_type); values_out_iterator_name = vo_name; values_out_iterator_src = vo_src; + + // For STORAGE values, ensure pointer types in iterator names/sources use items_storage_t* + if (values_in_it.value_type.type == cccl_type_enum::CCCL_STORAGE) + { + auto replace_all = [](std::string& s, const std::string& from, const std::string& to) { + if (from.empty()) + { + return; + } + size_t pos = 0; + while ((pos = s.find(from, pos)) != std::string::npos) + { + s.replace(pos, from.length(), to); + pos += to.length(); + } + }; + replace_all(values_in_iterator_src, "storage_t", "items_storage_t"); + replace_all(values_out_iterator_src, "storage_t", "items_storage_t"); + replace_all(values_in_iterator_name, "storage_t", "items_storage_t"); + replace_all(values_out_iterator_name, "storage_t", "items_storage_t"); + } } else { @@ -483,29 +521,35 @@ CUresult cccl_device_segmented_sort_build( const auto offset_t = cccl_type_enum_to_name(cccl_type_enum::CCCL_INT64); - const std::string key_t = cccl_type_enum_to_name(keys_in_it.value_type.type); - const std::string value_t = keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); + const std::string key_t = cccl_type_enum_to_name(keys_in_it.value_type.type); + const std::string value_t = + keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); const std::string dependent_definitions_src = std::format( R"XXX( struct __align__({1}) storage_t {{ char data[{0}]; }}; -{2} -{3} +struct __align__({3}) items_storage_t {{ + char data[{2}]; +}}; {4} {5} {6} {7} +{8} +{9} )XXX", keys_in_it.value_type.size, // 0 keys_in_it.value_type.alignment, // 1 - keys_in_iterator_src, // 2 - keys_out_iterator_src, // 3 - values_in_iterator_src, // 4 - values_out_iterator_src, // 5 - start_offset_iterator_src, // 6 - end_offset_iterator_src); // 7 + values_in_it.value_type.size, // 2 + values_in_it.value_type.alignment, // 3 + keys_in_iterator_src, // 4 + keys_out_iterator_src, // 5 + values_in_iterator_src, // 6 + values_out_iterator_src, // 7 + start_offset_iterator_src, // 8 + end_offset_iterator_src); // 9 // Runtime parameter tuning const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); @@ -763,8 +807,8 @@ CUresult cccl_device_segmented_sort_impl( d_values_double_buffer, num_items, num_segments, - indirect_iterator_t{start_offset_in}, - indirect_iterator_t{end_offset_in}, + start_offset_in, + end_offset_in, is_overwrite_okay, stream, /* kernel_source */ {build}, From d2b22bbcb1644994730818257ab0a3352c1bc06d Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 10 Sep 2025 18:37:28 +0000 Subject: [PATCH 039/100] Expand testing of segmented sort --- c/parallel/test/test_segmented_sort.cpp | 476 ++++++++++++++++-------- 1 file changed, 320 insertions(+), 156 deletions(-) diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 289f0406390..1e5a9d98b11 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -11,10 +11,8 @@ #include #include #include -#include #include // std::optional #include -#include #include #include @@ -211,46 +209,60 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes pointer_t values_in; pointer_t values_out; - // Always use signed long long for offset iterator since negative advances are possible - using SizeT = signed long long; - static constexpr std::string_view index_ty_name = "signed long long"; + // // Always use signed long long for offset iterator since negative advances are possible + // using SizeT = unsigned long long; + // static constexpr std::string_view index_ty_name = "unsigned long long"; - struct segment_offset_iterator_state_t - { - SizeT linear_id; - SizeT segment_size; - }; + // struct segment_offset_iterator_state_t + // { + // SizeT linear_id; + // SizeT segment_size; + // }; + + // static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; + // static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; + // static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; - static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; - static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; - static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; + // const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = + // make_step_counting_iterator_sources( + // index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); - const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = - make_step_counting_iterator_sources( - index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); + // iterator_t start_offset_it = + // make_iterator( + // {offset_iterator_state_name, offset_iterator_state_src}, + // {advance_offset_method_name, offset_iterator_advance_src}, + // {deref_offset_method_name, offset_iterator_deref_src}); - iterator_t start_offset_it = - make_iterator( - {offset_iterator_state_name, offset_iterator_state_src}, - {advance_offset_method_name, offset_iterator_advance_src}, - {deref_offset_method_name, offset_iterator_deref_src}); + // start_offset_it.state.linear_id = 0; + // start_offset_it.state.segment_size = segment_size; - start_offset_it.state.linear_id = 0; - start_offset_it.state.segment_size = segment_size; + // // Create end offset iterator (points to one past start) + // iterator_t end_offset_it = + // make_iterator( + // {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); - // Create end offset iterator (points to one past start) - iterator_t end_offset_it = - make_iterator( - {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); + // end_offset_it.state.linear_id = 1; + // end_offset_it.state.segment_size = segment_size; - end_offset_it.state.linear_id = 1; - end_offset_it.state.segment_size = segment_size; + // // Provide host-advance callbacks for offset iterators + // auto start_offsets_cccl = static_cast(start_offset_it); + // auto end_offsets_cccl = static_cast(end_offset_it); + // start_offsets_cccl.host_advance = &host_advance_linear_id; + // end_offsets_cccl.host_advance = &host_advance_linear_id; - // Provide host-advance callbacks for offset iterators - auto start_offsets_cccl = static_cast(start_offset_it); - auto end_offsets_cccl = static_cast(end_offset_it); - start_offsets_cccl.host_advance = &host_advance_linear_id; - end_offsets_cccl.host_advance = &host_advance_linear_id; + // Provide device arrays of start/end offsets instead of custom iterators + using SizeT = unsigned long long; + + std::vector start_offsets(n_segments); + std::vector end_offsets(n_segments); + for (std::size_t i = 0; i < n_segments; ++i) + { + start_offsets[i] = static_cast(i * segment_size); + end_offsets[i] = static_cast((i + 1) * segment_size); + } + + pointer_t start_offsets_ptr(start_offsets); + pointer_t end_offsets_ptr(end_offsets); auto& build_cache = get_cache(); const std::string& key_string = KeyBuilder::join( @@ -269,8 +281,10 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes values_out, n_elems, n_segments, - start_offsets_cccl, - end_offsets_cccl, + // start_offsets_cccl, + // end_offsets_cccl, + start_offsets_ptr, + end_offsets_ptr, is_overwrite_okay, &selector, build_cache, @@ -307,8 +321,8 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - const std::size_t n_segments = GENERATE(0, 10, take(2, random(30, 100))); - const std::size_t segment_size = GENERATE(1, 15, take(2, random(5, 50))); + const std::size_t n_segments = GENERATE(0, 13, take(2, random(1 << 10, 1 << 12))); + const std::size_t segment_size = GENERATE(1, 12, take(2, random(1 << 10, 1 << 12))); const std::size_t n_elems = n_segments * segment_size; @@ -329,51 +343,67 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] pointer_t values_in_ptr(host_values); pointer_t values_out_ptr(host_values_out); - using SizeT = signed long long; - static constexpr std::string_view index_ty_name = "signed long long"; + // using SizeT = signed long long; + // static constexpr std::string_view index_ty_name = "signed long long"; - struct segment_offset_iterator_state_t - { - SizeT linear_id; - SizeT segment_size; - }; + // struct segment_offset_iterator_state_t + // { + // SizeT linear_id; + // SizeT segment_size; + // }; + + // static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; + // static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; + // static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; - static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; - static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; - static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; + // const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = + // make_step_counting_iterator_sources( + // index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); - const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = - make_step_counting_iterator_sources( - index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); + // iterator_t start_offset_it = + // make_iterator( + // {offset_iterator_state_name, offset_iterator_state_src}, + // {advance_offset_method_name, offset_iterator_advance_src}, + // {deref_offset_method_name, offset_iterator_deref_src}); - iterator_t start_offset_it = - make_iterator( - {offset_iterator_state_name, offset_iterator_state_src}, - {advance_offset_method_name, offset_iterator_advance_src}, - {deref_offset_method_name, offset_iterator_deref_src}); + // start_offset_it.state.linear_id = 0; + // start_offset_it.state.segment_size = segment_size; - start_offset_it.state.linear_id = 0; - start_offset_it.state.segment_size = segment_size; + // iterator_t end_offset_it = + // make_iterator( + // {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); - iterator_t end_offset_it = - make_iterator( - {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); + // end_offset_it.state.linear_id = 1; + // end_offset_it.state.segment_size = segment_size; - end_offset_it.state.linear_id = 1; - end_offset_it.state.segment_size = segment_size; + // // Provide host-advance callbacks for offset iterators + // auto start_offsets_cccl = static_cast(start_offset_it); + // auto end_offsets_cccl = static_cast(end_offset_it); + // start_offsets_cccl.host_advance = &host_advance_linear_id; + // end_offsets_cccl.host_advance = &host_advance_linear_id; - // Provide host-advance callbacks for offset iterators - auto start_offsets_cccl = static_cast(start_offset_it); - auto end_offsets_cccl = static_cast(end_offset_it); - start_offsets_cccl.host_advance = &host_advance_linear_id; - end_offsets_cccl.host_advance = &host_advance_linear_id; + // Provide device arrays of start/end offsets instead of custom iterators + using SizeT = unsigned long long; + + std::vector start_offsets(n_segments); + std::vector end_offsets(n_segments); + for (std::size_t i = 0; i < n_segments; ++i) + { + start_offsets[i] = static_cast(i * segment_size); + end_offsets[i] = static_cast((i + 1) * segment_size); + } + + pointer_t start_offsets_ptr(start_offsets); + pointer_t end_offsets_ptr(end_offsets); auto& build_cache = get_cache(); const std::string& key_string = KeyBuilder::join( {KeyBuilder::bool_as_key(is_descending), KeyBuilder::type_as_key(), KeyBuilder::type_as_key(), - KeyBuilder::bool_as_key(is_overwrite_okay)}); + KeyBuilder::bool_as_key(is_overwrite_okay), + KeyBuilder::bool_as_key(n_elems == 0)}); // this results in the values pointer being null which results in a keys + // only build const auto& test_key = std::make_optional(key_string); int selector = -1; @@ -386,8 +416,10 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] values_out_ptr, n_elems, n_segments, - start_offsets_cccl, - end_offsets_cccl, + // start_offsets_cccl, + // end_offsets_cccl, + start_offsets_ptr, + end_offsets_ptr, is_overwrite_okay, &selector, build_cache, @@ -440,6 +472,8 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] REQUIRE(expected_values == std::vector(output_vals)); } +// These tests with custom types are currently failing TODO: add issue +#ifdef NEVER_DEFINED struct custom_pair { int key; @@ -449,50 +483,49 @@ struct custom_pair { return key == other.key && value == other.value; } - - bool operator<(const custom_pair& other) const - { - return key < other.key; - } }; struct SegmentedSort_CustomTypes_Fixture_Tag; -C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][custom_types]", test_params_tuple) +C2H_TEST("SegmentedSort works with custom types as values", "[segmented_sort][custom_types]", test_params_tuple) { - using T = c2h::get<0, TestType>; - using key_t = custom_pair; + using T = c2h::get<0, TestType>; + using key_t = typename T::KeyT; + using value_t = custom_pair; constexpr auto this_test_params = T(); constexpr bool is_descending = this_test_params.is_descending(); constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - const std::size_t n_segments = 25; - const std::size_t segment_size = 20; - const std::size_t n_elems = n_segments * segment_size; + const std::size_t n_segments = GENERATE(0, 13, take(2, random(1 << 10, 1 << 12))); + const std::size_t segment_size = GENERATE(1, 12, take(2, random(1 << 10, 1 << 12))); + + std::cout << "n_segments: " << n_segments << ", segment_size: " << segment_size << std::endl; - // Generate custom key data + const std::size_t n_elems = n_segments * segment_size; + + // Generate primitive keys + std::vector host_keys_int = generate(n_elems); std::vector host_keys(n_elems); + std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int x) { + return static_cast(x); + }); + + // Generate custom values + std::vector host_values(n_elems); for (std::size_t i = 0; i < n_elems; ++i) { - host_keys[i] = custom_pair{static_cast(i % 1000), static_cast(i % 100)}; + host_values[i] = value_t{static_cast(i % 1000), static_cast(i % 100)}; } - - // Generate float values by first generating ints and then transforming - std::vector host_values_int = generate(n_elems); - std::vector host_values(n_elems); - std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { - return static_cast(x); - }); std::vector host_keys_out(n_elems); - std::vector host_values_out(n_elems); + std::vector host_values_out(n_elems); pointer_t keys_in_ptr(host_keys); pointer_t keys_out_ptr(host_keys_out); - pointer_t values_in_ptr(host_values); - pointer_t values_out_ptr(host_values_out); + pointer_t values_in_ptr(host_values); + pointer_t values_out_ptr(host_values_out); - using SizeT = cuda::std::size_t; + using SizeT = long; std::vector segments(n_segments + 1); for (std::size_t i = 0; i <= n_segments; ++i) { @@ -509,8 +542,9 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust const std::string& key_string = KeyBuilder::join( {KeyBuilder::bool_as_key(is_descending), KeyBuilder::type_as_key(), - KeyBuilder::type_as_key(), - KeyBuilder::bool_as_key(is_overwrite_okay)}); + KeyBuilder::type_as_key(), + KeyBuilder::bool_as_key(is_overwrite_okay), + KeyBuilder::bool_as_key(n_elems == 0)}); const auto& test_key = std::make_optional(key_string); int selector = -1; @@ -531,14 +565,14 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust test_key); // Create expected result - std::vector> key_value_pairs; + std::vector> key_value_pairs; for (std::size_t i = 0; i < n_elems; ++i) { key_value_pairs.emplace_back(host_keys[i], host_values[i]); } std::vector expected_keys(n_elems); - std::vector expected_values(n_elems); + std::vector expected_values(n_elems); for (std::size_t i = 0; i < n_segments; ++i) { @@ -547,19 +581,19 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust if (is_descending) { - std::sort(key_value_pairs.begin() + segment_start, - key_value_pairs.begin() + segment_end, - [](const auto& a, const auto& b) { - return b.first < a.first; - }); + std::stable_sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return b.first < a.first; + }); } else { - std::sort(key_value_pairs.begin() + segment_start, - key_value_pairs.begin() + segment_end, - [](const auto& a, const auto& b) { - return a.first < b.first; - }); + std::stable_sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return a.first < b.first; + }); } // Extract sorted keys and values @@ -572,9 +606,11 @@ C2H_TEST("SegmentedSort works with custom types as keys", "[segmented_sort][cust auto& output_keys = (is_overwrite_okay && selector == 0) ? keys_in_ptr : keys_out_ptr; auto& output_vals = (is_overwrite_okay && selector == 0) ? values_in_ptr : values_out_ptr; + REQUIRE(expected_keys == std::vector(output_keys)); - REQUIRE(expected_values == std::vector(output_vals)); + REQUIRE(expected_values == std::vector(output_vals)); } +#endif using SizeT = unsigned long long; @@ -613,34 +649,36 @@ extern "C" __device__ unsigned long long dereference_variable_offset_it(variable struct SegmentedSort_VariableSegments_Fixture_Tag; C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][variable_segments]", test_params_tuple) { - using T = c2h::get<0, TestType>; - using key_t = std::int32_t; + using T = c2h::get<0, TestType>; constexpr auto this_test_params = T(); constexpr bool is_descending = this_test_params.is_descending(); constexpr auto order = is_descending ? CCCL_DESCENDING : CCCL_ASCENDING; constexpr bool is_overwrite_okay = this_test_params.is_overwrite_okay(); - const std::size_t n_segments = 20; + const std::size_t n_segments = GENERATE(20, 600); // Create variable segment sizes - std::vector segment_sizes = {1, 5, 10, 20, 30, 15, 8, 3, 25, 12, 7, 18, 22, 4, 35, 9, 14, 6, 28, 11}; + const std::vector base_pattern = { + 1, 5, 10, 20, 30, 50, 100, 3, 25, 600, 7, 18, 300, 4, 35, 9, 14, 700, 28, 11}; + std::vector segment_sizes; + segment_sizes.reserve(n_segments); + while (segment_sizes.size() < n_segments) + { + const std::size_t remaining = n_segments - segment_sizes.size(); + const std::size_t copy_count = std::min(remaining, base_pattern.size()); + segment_sizes.insert(segment_sizes.end(), base_pattern.begin(), base_pattern.begin() + copy_count); + } REQUIRE(segment_sizes.size() == n_segments); std::size_t n_elems = std::accumulate(segment_sizes.begin(), segment_sizes.end(), 0ULL); std::vector host_keys_int = generate(n_elems); - std::vector host_keys(n_elems); - std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int x) { - return static_cast(x); - }); + std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); // Generate float values by first generating ints and then transforming std::vector host_values_int = generate(n_elems); - std::vector host_values(n_elems); - std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int x) { - return static_cast(x); - }); + std::vector host_values(host_values_int.begin(), host_values_int.end()); std::vector host_keys_out(n_elems); std::vector host_values_out(n_elems); @@ -649,10 +687,50 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va pointer_t values_in_ptr(host_values); pointer_t values_out_ptr(host_values_out); - // Create segment offset arrays + // // Create segment offset arrays + // std::vector start_offsets(n_segments); + // std::vector end_offsets(n_segments); + + // SizeT current_offset = 0; + // for (std::size_t i = 0; i < n_segments; ++i) + // { + // start_offsets[i] = current_offset; + // current_offset += segment_sizes[i]; + // end_offsets[i] = current_offset; + // } + + // pointer_t start_offsets_ptr(start_offsets); + // pointer_t end_offsets_ptr(end_offsets); + + // const auto& [offset_state_src, offset_advance_src, offset_deref_src] = make_variable_segment_iterator_sources(); + + // iterator_t start_offset_it = + // make_iterator( + // {"variable_segment_offset_iterator_state_t", offset_state_src}, + // {"advance_variable_offset_it", offset_advance_src}, + // {"dereference_variable_offset_it", offset_deref_src}); + + // start_offset_it.state.linear_id = 0; + // start_offset_it.state.offsets = start_offsets_ptr.ptr; + + // iterator_t end_offset_it = + // make_iterator( + // {"variable_segment_offset_iterator_state_t", ""}, + // {"advance_variable_offset_it", ""}, + // {"dereference_variable_offset_it", ""}); + + // end_offset_it.state.linear_id = 0; + // end_offset_it.state.offsets = end_offsets_ptr.ptr; + + // auto cccl_start_offsets_it = static_cast(start_offset_it); + // auto cccl_end_offsets_it = static_cast(end_offset_it); + + // // set host_advance functions + // cccl_start_offsets_it.host_advance = &host_advance_linear_id; + // cccl_end_offsets_it.host_advance = &host_advance_linear_id; + std::vector start_offsets(n_segments); std::vector end_offsets(n_segments); - SizeT current_offset = 0; for (std::size_t i = 0; i < n_segments; ++i) { @@ -664,26 +742,6 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va pointer_t start_offsets_ptr(start_offsets); pointer_t end_offsets_ptr(end_offsets); - const auto& [offset_state_src, offset_advance_src, offset_deref_src] = make_variable_segment_iterator_sources(); - - iterator_t start_offset_it = - make_iterator( - {"variable_segment_offset_iterator_state_t", offset_state_src}, - {"advance_variable_offset_it", offset_advance_src}, - {"dereference_variable_offset_it", offset_deref_src}); - - start_offset_it.state.linear_id = 0; - start_offset_it.state.offsets = start_offsets_ptr.ptr; - - iterator_t end_offset_it = - make_iterator( - {"variable_segment_offset_iterator_state_t", ""}, - {"advance_variable_offset_it", ""}, - {"dereference_variable_offset_it", ""}); - - end_offset_it.state.linear_id = 0; - end_offset_it.state.offsets = end_offsets_ptr.ptr; - auto& build_cache = get_cache(); const std::string& key_string = KeyBuilder::join( {KeyBuilder::bool_as_key(is_descending), @@ -702,8 +760,8 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va values_out_ptr, n_elems, n_segments, - start_offset_it, - end_offset_it, + start_offsets_ptr, + end_offsets_ptr, is_overwrite_okay, &selector, build_cache, @@ -726,19 +784,19 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va if (is_descending) { - std::sort(key_value_pairs.begin() + segment_start, - key_value_pairs.begin() + segment_end, - [](const auto& a, const auto& b) { - return b.first < a.first; - }); + std::stable_sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return b.first < a.first; + }); } else { - std::sort(key_value_pairs.begin() + segment_start, - key_value_pairs.begin() + segment_end, - [](const auto& a, const auto& b) { - return a.first < b.first; - }); + std::stable_sort(key_value_pairs.begin() + segment_start, + key_value_pairs.begin() + segment_end, + [](const auto& a, const auto& b) { + return a.first < b.first; + }); } // Extract sorted keys and values @@ -754,3 +812,109 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va REQUIRE(expected_keys == std::vector(output_keys)); REQUIRE(expected_values == std::vector(output_vals)); } + +// struct SegmentedSort_LargeNumSegments_Fixture_Tag; +// C2H_TEST("SegmentedSort works with a large number of segments", "[segmented_sort][large_segments]") { +// using key_t = int; + +// constexpr bool is_descending = false; +// constexpr auto order = CCCL_ASCENDING; +// constexpr bool is_overwrite_okay = false; + +// const std::size_t n_segments = 1; +// const std::size_t segment_size = 27; +// const std::size_t n_elems = n_segments * segment_size; + +// std::vector host_keys = generate(n_elems); +// std::vector host_keys_out(n_elems); + +// REQUIRE(host_keys.size() == n_elems); +// REQUIRE(host_keys_out.size() == n_elems); + +// pointer_t keys_in_ptr(host_keys); +// pointer_t keys_out_ptr(host_keys_out); + +// pointer_t values_in_ptr; +// pointer_t values_out_ptr; + +// struct segment_offset_iterator_state_t +// { +// SizeT linear_id; +// SizeT segment_size; +// }; + +// static constexpr std::string_view index_ty_name = "unsigned long long"; +// static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; +// static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; +// static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; + +// const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = +// make_step_counting_iterator_sources( +// index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); + +// // start = i * segment_size +// iterator_t start_offset_it = +// make_iterator( +// {offset_iterator_state_name, offset_iterator_state_src}, +// {advance_offset_method_name, offset_iterator_advance_src}, +// {deref_offset_method_name, offset_iterator_deref_src}); +// start_offset_it.state.linear_id = 0; +// start_offset_it.state.segment_size = segment_size; + +// // end = (i + 1) * segment_size +// auto end_offset_it = +// make_iterator( +// {offset_iterator_state_name, ""}, +// {advance_offset_method_name, ""}, +// {deref_offset_method_name, ""}); +// end_offset_it.state.linear_id = 1; +// end_offset_it.state.segment_size = segment_size; + +// cccl_iterator_t start_offsets_cccl(start_offset_it); +// cccl_iterator_t end_offsets_cccl(end_offset_it); +// start_offsets_cccl.host_advance = &host_advance_linear_id; +// end_offsets_cccl.host_advance = &host_advance_linear_id; + +// auto& build_cache = get_cache(); +// const std::string& key_string = KeyBuilder::join( +// {KeyBuilder::bool_as_key(is_descending), +// KeyBuilder::type_as_key(), +// KeyBuilder::type_as_key(), +// KeyBuilder::bool_as_key(is_overwrite_okay)}); +// const auto& test_key = std::make_optional(key_string); + +// int selector = -1; + +// segmented_sort( +// order, +// keys_in_ptr, +// keys_out_ptr, +// values_in_ptr, +// values_out_ptr, +// n_elems, +// n_segments, +// start_offsets_cccl, +// end_offsets_cccl, +// is_overwrite_okay, +// &selector, +// build_cache, +// test_key); + +// std::vector expected_keys = host_keys; +// for (std::size_t i = 0; i < n_segments; ++i) +// { +// const std::size_t segment_start = i * segment_size; +// const std::size_t segment_end = segment_start + segment_size; +// if (is_descending) +// { +// std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end, std::greater()); +// } +// else +// { +// std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end); +// } +// } + +// auto& output_keys = (is_overwrite_okay && selector == 0) ? keys_in_ptr : keys_out_ptr; +// REQUIRE(expected_keys == std::vector(output_keys)); +// } From 5dd308f2171f4949c1710ce6ca43685fe8f4140b Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 10 Sep 2025 18:42:44 +0000 Subject: [PATCH 040/100] Pass segment selectors through kernel source --- .../dispatch/dispatch_segmented_sort.cuh | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 787920a7951..a944ac4641c 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -282,6 +282,23 @@ struct DeviceSegmentedSortKernelSource { return sizeof(KeyT); } + + using LargeSegmentsSelectorT = + cub::detail::segmented_sort::LargeSegmentsSelectorT; + using SmallSegmentsSelectorT = + cub::detail::segmented_sort::SmallSegmentsSelectorT; + + CUB_RUNTIME_FUNCTION static constexpr auto LargeSegmentsSelector( + OffsetT offset, BeginOffsetIteratorT begin_offset_iterator, EndOffsetIteratorT end_offset_iterator) + { + return LargeSegmentsSelectorT(offset, begin_offset_iterator, end_offset_iterator); + } + + CUB_RUNTIME_FUNCTION static constexpr auto SmallSegmentsSelector( + OffsetT offset, BeginOffsetIteratorT begin_offset_iterator, EndOffsetIteratorT end_offset_iterator) + { + return SmallSegmentsSelectorT(offset, begin_offset_iterator, end_offset_iterator); + } }; } // namespace detail::segmented_sort @@ -325,11 +342,6 @@ struct DispatchSegmentedSort static constexpr int KEYS_ONLY = ::cuda::std::is_same_v; - using LargeSegmentsSelectorT = - cub::detail::segmented_sort::LargeSegmentsSelectorT; - using SmallSegmentsSelectorT = - cub::detail::segmented_sort::SmallSegmentsSelectorT; - // Partition selects large and small groups. The middle group is not selected. static constexpr size_t num_selected_groups = 2; @@ -416,7 +428,6 @@ struct DispatchSegmentedSort //------------------------------------------------------------------------ const bool partition_segments = num_segments > wrapped_policy.PartitioningThreshold(); - // const bool partition_segments = num_segments > ActivePolicyT::PARTITIONING_THRESHOLD; cub::detail::temporary_storage::layout<5> temporary_storage_layout; @@ -446,10 +457,9 @@ struct DispatchSegmentedSort size_t three_way_partition_temp_storage_bytes{}; - LargeSegmentsSelectorT large_segments_selector( - wrapped_policy.MediumPolicyItemsPerTile(), d_begin_offsets, d_end_offsets); - - SmallSegmentsSelectorT small_segments_selector( + auto large_segments_selector = + kernel_source.LargeSegmentsSelector(wrapped_policy.MediumPolicyItemsPerTile(), d_begin_offsets, d_end_offsets); + auto small_segments_selector = kernel_source.SmallSegmentsSelector( wrapped_policy.SmallPolicyItemsPerTile() + 1, d_begin_offsets, d_end_offsets); auto device_partition_temp_storage = keys_slot->create_alias(); @@ -478,8 +488,8 @@ struct DispatchSegmentedSort decltype(small_segments_indices.get()), decltype(medium_indices_iterator), decltype(group_sizes.get()), - LargeSegmentsSelectorT, - SmallSegmentsSelectorT, + decltype(large_segments_selector), + decltype(small_segments_selector), PartitionOffsetT, PartitionPolicyHub, PartitionKernelSource, @@ -597,8 +607,6 @@ struct DispatchSegmentedSort : (is_num_passes_odd) ? values_allocation.get() : d_values.Alternate()); - using MaxPolicyT = typename PolicyHub::MaxPolicy; - if (partition_segments) { // Partition input segments into size groups and assign specialized @@ -715,8 +723,8 @@ private: size_t three_way_partition_temp_storage_bytes, cub::detail::device_double_buffer& d_keys_double_buffer, cub::detail::device_double_buffer& d_values_double_buffer, - LargeSegmentsSelectorT& large_segments_selector, - SmallSegmentsSelectorT& small_segments_selector, + KernelSource::LargeSegmentsSelectorT& large_segments_selector, + KernelSource::SmallSegmentsSelectorT& small_segments_selector, cub::detail::temporary_storage::alias& device_partition_temp_storage, cub::detail::temporary_storage::alias& large_and_medium_segments_indices, cub::detail::temporary_storage::alias& small_segments_indices, @@ -761,8 +769,8 @@ private: decltype(small_segments_indices.get()), decltype(medium_indices_iterator), decltype(group_sizes.get()), - LargeSegmentsSelectorT, - SmallSegmentsSelectorT, + decltype(large_segments_selector), + decltype(small_segments_selector), PartitionOffsetT, PartitionPolicyHub, PartitionKernelSource, From 6162d17aee4aa4e2a6bfcf4a4e712a34d7bffc33 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 10 Sep 2025 19:42:24 +0000 Subject: [PATCH 041/100] remove merge leftovers --- cub/cub/device/dispatch/kernels/three_way_partition.cuh | 7 ------- 1 file changed, 7 deletions(-) diff --git a/cub/cub/device/dispatch/kernels/three_way_partition.cuh b/cub/cub/device/dispatch/kernels/three_way_partition.cuh index d81f5b81e8b..4aac47f3e29 100644 --- a/cub/cub/device/dispatch/kernels/three_way_partition.cuh +++ b/cub/cub/device/dispatch/kernels/three_way_partition.cuh @@ -1,9 +1,5 @@ // SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -<<<<<<< HEAD -// SPDX-License-Identifier: BSD-3-Clause -======= // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception ->>>>>>> main #pragma once @@ -23,7 +19,6 @@ CUB_NAMESPACE_BEGIN namespace detail::three_way_partition { -<<<<<<< HEAD // Offset type used to instantiate the stream three-way-partition-kernel and agent to index the items within one // partition @@ -104,8 +99,6 @@ public: } }; -======= ->>>>>>> main /****************************************************************************** * Kernel entry points *****************************************************************************/ From 773119c5785f33b78ac3f0fa8883e05d4e1db7f9 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 12 Sep 2025 19:56:37 +0000 Subject: [PATCH 042/100] Pass large and small selector ops through kernel source. This is consistent with how dispatch three way partition is called and wil eventually allow us to properly copy the states of iterators --- c/parallel/src/segmented_sort.cu | 179 ++++++++++++++++++++---- c/parallel/test/test_segmented_sort.cpp | 4 +- 2 files changed, 152 insertions(+), 31 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index a325e40fced..e563f5a7a05 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -24,6 +24,7 @@ #include // printf #include "jit_templates/templates/input_iterator.h" +#include "jit_templates/templates/operation.h" #include "jit_templates/templates/output_iterator.h" #include "jit_templates/traits.h" #include "util/context.h" @@ -172,21 +173,130 @@ struct segmented_sort_kernel_source return build.key_type.size; } - using LargeSegmentsSelectorT = cub::detail::segmented_sort::LargeSegmentsSelectorT; - using SmallSegmentsSelectorT = cub::detail::segmented_sort::SmallSegmentsSelectorT; + using LargeSegmentsSelectorT = indirect_arg_t; + using SmallSegmentsSelectorT = indirect_arg_t; - auto LargeSegmentsSelector( + indirect_arg_t LargeSegmentsSelector( OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const { - return LargeSegmentsSelectorT( - offset, *reinterpret_cast(begin_offset_iterator.ptr), *reinterpret_cast(end_offset_iterator.ptr)); + cccl_op_t op = LargeSegmentsSelectorOp(offset, begin_offset_iterator, end_offset_iterator); + return indirect_arg_t{op}; } - auto SmallSegmentsSelector( + indirect_arg_t SmallSegmentsSelector( OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const { - return SmallSegmentsSelectorT( - offset, *reinterpret_cast(begin_offset_iterator.ptr), *reinterpret_cast(end_offset_iterator.ptr)); + cccl_op_t op = SmallSegmentsSelectorOp(offset, begin_offset_iterator, end_offset_iterator); + return indirect_arg_t{op}; + } + + void SetSegmentOffset(cccl_op_t& selector, long long base_segment_offset) const + { + auto* st = reinterpret_cast(selector.state); + st->base_segment_offset = base_segment_offset; + } + + void SetSegmentOffset(indirect_arg_t& selector, long long base_segment_offset) const + { + auto* st = reinterpret_cast(selector.ptr); + st->base_segment_offset = base_segment_offset; + } + + struct selector_state_t + { + long long threshold; + const long long* begin_offsets; + const long long* end_offsets; + long long base_segment_offset; + }; + + // Return stateful cccl_op_t predicates equivalent to the CUB selectors above. + // These embed C++ source for a device function and capture state (threshold and offset arrays). + static cccl_op_t LargeSegmentsSelectorOp( + OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) + { + // Persist state storage and code across the returned cccl_op_t lifetime + static selector_state_t state{}; + state.threshold = static_cast(offset); + state.begin_offsets = reinterpret_cast(*reinterpret_cast(begin_offset_iterator.ptr)); + state.end_offsets = reinterpret_cast(*reinterpret_cast(end_offset_iterator.ptr)); + state.base_segment_offset = 0; + + static std::string code; + code = std::string{ + R"XXX(#include +extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) +{ + struct state_t { + long long threshold; + const long long* begin_offsets; + const long long* end_offsets; + long long base_segment_offset; + }; + + auto* st = static_cast(state_ptr); + using local_segment_index_t = ::cuda::std::uint32_t; + const local_segment_index_t sid = *static_cast(arg_ptr); + const long long begin = st->begin_offsets[st->base_segment_offset + sid]; + const long long end = st->end_offsets[st->base_segment_offset + sid]; + const bool pred = (end - begin) > st->threshold; + *reinterpret_cast(result_ptr) = pred; +} +)XXX"}; + + cccl_op_t op{}; + op.type = cccl_op_kind_t::CCCL_STATEFUL; + op.name = "cccl_large_segments_selector_op"; + op.code = code.c_str(); + op.code_size = code.size(); + op.code_type = CCCL_OP_CPP_SOURCE; + op.size = sizeof(state); + op.alignment = alignof(selector_state_t); + op.state = &state; + return op; + } + + static cccl_op_t SmallSegmentsSelectorOp( + OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) + { + static selector_state_t state{}; + state.threshold = static_cast(offset); + state.begin_offsets = reinterpret_cast(*reinterpret_cast(begin_offset_iterator.ptr)); + state.end_offsets = reinterpret_cast(*reinterpret_cast(end_offset_iterator.ptr)); + state.base_segment_offset = 0; + + static std::string code; + code = std::string{ + R"XXX(#include +extern "C" __device__ void cccl_small_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) +{ + struct state_t { + long long threshold; + const long long* begin_offsets; + const long long* end_offsets; + long long base_segment_offset; + }; + auto* st = static_cast(state_ptr); + using local_segment_index_t = ::cuda::std::uint32_t; + const local_segment_index_t sid = *static_cast(arg_ptr); + const long long begin = st->begin_offsets[st->base_segment_offset + sid]; + const long long end = st->end_offsets[st->base_segment_offset + sid]; + const bool pred = (end - begin) < st->threshold; + + *reinterpret_cast(result_ptr) = pred; +} +)XXX"}; + + cccl_op_t op{}; + op.type = cccl_op_kind_t::CCCL_STATEFUL; + op.name = "cccl_small_segments_selector_op"; + op.code = code.c_str(); + op.code_size = code.size(); + op.code_type = CCCL_OP_CPP_SOURCE; + op.size = sizeof(state); + op.alignment = alignof(selector_state_t); + op.state = &state; + return op; } }; @@ -201,8 +311,7 @@ std::string get_three_way_partition_init_kernel_name() num_selected_it_t); // 1 } -std::string -get_three_way_partition_kernel_name(std::string_view start_offset_iterator_t, std::string_view end_offset_iterator_t) +std::string get_three_way_partition_kernel_name(std::string_view large_selector_t, std::string_view small_selector_t) { std::string chained_policy_t; check(nvrtcGetTypeName(&chained_policy_t)); @@ -218,18 +327,6 @@ get_three_way_partition_kernel_name(std::string_view start_offset_iterator_t, st std::string offset_t; check(nvrtcGetTypeName(&offset_t)); - std::string select_first_part_op_t = std::format( - "cub::detail::segmented_sort::LargeSegmentsSelectorT<{0}, {1}, {2}>", - offset_t, // 0 - start_offset_iterator_t, // 1 - end_offset_iterator_t); // 2 - - std::string select_second_part_op_t = std::format( - "cub::detail::segmented_sort::SmallSegmentsSelectorT<{0}, {1}, {2}>", - offset_t, // 0 - start_offset_iterator_t, // 1 - end_offset_iterator_t); // 2 - constexpr std::string_view per_partition_offset_t = "cub::detail::three_way_partition::per_partition_offset_t"; constexpr std::string_view streaming_context_t = "cub::detail::three_way_partition::streaming_context_t"; @@ -244,8 +341,8 @@ get_three_way_partition_kernel_name(std::string_view start_offset_iterator_t, st unselected_out_it_t, // 4 (UnselectedOutputIteratorT) num_selected_it_t, // 5 (NumSelectedIteratorT) scan_tile_state_t, // 6 (ScanTileStateT) - select_first_part_op_t, // 7 (SelectFirstPartOp) - select_second_part_op_t, // 8 (SelectSecondPartOp) + large_selector_t, // 7 (SelectFirstPartOp) + small_selector_t, // 8 (SelectSecondPartOp) per_partition_offset_t, // 9 (OffsetT) streaming_context_t); // 10 (StreamingContextT) } @@ -418,6 +515,8 @@ struct segmented_sort_values_input_iterator_tag; struct segmented_sort_values_output_iterator_tag; struct segmented_sort_start_offset_iterator_tag; struct segmented_sort_end_offset_iterator_tag; +struct segmented_sort_large_selector_tag; +struct segmented_sort_small_selector_tag; CUresult cccl_device_segmented_sort_build( cccl_device_segmented_sort_build_result_t* build_ptr, @@ -525,6 +624,21 @@ CUresult cccl_device_segmented_sort_build( const std::string value_t = keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); + // Build selector operations as cccl_op_t and generate their functor wrappers + cccl_op_t large_selector_op = + segmented_sort::segmented_sort_kernel_source::LargeSegmentsSelectorOp(0, start_offset_it, end_offset_it); + cccl_op_t small_selector_op = + segmented_sort::segmented_sort_kernel_source::SmallSegmentsSelectorOp(0, start_offset_it, end_offset_it); + + cccl_type_info bool_t{sizeof(bool), alignof(bool), cccl_type_enum::CCCL_BOOLEAN}; + cccl_type_info u32_t{sizeof(::cuda::std::uint32_t), alignof(::cuda::std::uint32_t), cccl_type_enum::CCCL_UINT32}; + + const auto [large_selector_name, large_selector_src] = get_specialization( + template_id(), large_selector_op, bool_t, u32_t); + + const auto [small_selector_name, small_selector_src] = get_specialization( + template_id(), small_selector_op, bool_t, u32_t); + const std::string dependent_definitions_src = std::format( R"XXX( struct __align__({1}) storage_t {{ @@ -539,6 +653,8 @@ struct __align__({3}) items_storage_t {{ {7} {8} {9} +{10} +{11} )XXX", keys_in_it.value_type.size, // 0 keys_in_it.value_type.alignment, // 1 @@ -549,7 +665,9 @@ struct __align__({3}) items_storage_t {{ values_in_iterator_src, // 6 values_out_iterator_src, // 7 start_offset_iterator_src, // 8 - end_offset_iterator_src); // 9 + end_offset_iterator_src, // 9 + large_selector_src, // 10 + small_selector_src); // 11 // Runtime parameter tuning const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); @@ -654,7 +772,7 @@ struct device_three_way_partition_policy {{ std::string three_way_partition_init_kernel_name = segmented_sort::get_three_way_partition_init_kernel_name(); std::string three_way_partition_kernel_name = - segmented_sort::get_three_way_partition_kernel_name(start_offset_iterator_name, end_offset_iterator_name); + segmented_sort::get_three_way_partition_kernel_name(large_selector_name, small_selector_name); std::string segmented_sort_fallback_kernel_lowered_name; std::string segmented_sort_kernel_small_lowered_name; @@ -680,8 +798,8 @@ struct device_three_way_partition_policy {{ const char* lopts[num_lto_args] = {"-lto", arch.c_str()}; // Collect all LTO-IRs to be linked. - nvrtc_ltoir_list ltoir_list; - nvrtc_ltoir_list_appender appender{ltoir_list}; + nvrtc_linkable_list linkable_list; + nvrtc_linkable_list_appender appender{linkable_list}; // add iterator definitions appender.add_iterator_definition(keys_in_it); @@ -694,6 +812,9 @@ struct device_three_way_partition_policy {{ appender.add_iterator_definition(start_offset_it); appender.add_iterator_definition(end_offset_it); + appender.append_operation(large_selector_op); + appender.append_operation(small_selector_op); + nvrtc_link_result result = begin_linking_nvrtc_program(num_lto_args, lopts) ->add_program(nvrtc_translation_unit{final_src.c_str(), name}) @@ -709,7 +830,7 @@ struct device_three_way_partition_policy {{ ->get_name({three_way_partition_init_kernel_name, three_way_partition_init_kernel_lowered_name}) ->get_name({three_way_partition_kernel_name, three_way_partition_kernel_lowered_name}) ->link_program() - ->add_link_list(ltoir_list) + ->add_link_list(linkable_list) ->finalize_program(); // populate build struct members diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 1e5a9d98b11..cda67636e39 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -612,7 +612,7 @@ C2H_TEST("SegmentedSort works with custom types as values", "[segmented_sort][cu } #endif -using SizeT = unsigned long long; +using SizeT = signed long long; struct variable_segment_offset_iterator_state_t { @@ -821,7 +821,7 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va // constexpr auto order = CCCL_ASCENDING; // constexpr bool is_overwrite_okay = false; -// const std::size_t n_segments = 1; +// const std::size_t n_segments = 501; // const std::size_t segment_size = 27; // const std::size_t n_elems = n_segments * segment_size; From 9e2cb0f53830ac1ba609153804d24605c1127c59 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 12 Sep 2025 20:06:12 +0000 Subject: [PATCH 043/100] Fix merge leftover and set offset through kernel source --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 8f496e22e17..21b96fa0c1b 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -300,6 +300,13 @@ struct DeviceSegmentedSortKernelSource { return SmallSegmentsSelectorT(offset, begin_offset_iterator, end_offset_iterator); } + + template + CUB_RUNTIME_FUNCTION static constexpr void + SetSegmentOffset(SelectorT& selector, global_segment_offset_t base_segment_offset) + { + selector.base_segment_offset = base_segment_offset; + } }; } // namespace detail::segmented_sort @@ -327,7 +334,7 @@ template < THRUST_NS_QUALIFIER::counting_iterator, cub::detail::segmented_sort::local_segment_index_t*, cub::detail::segmented_sort::local_segment_index_t*, - THRUST_NS_QUALIFIER::reverse_iterator, + ::cuda::std::reverse_iterator, cub::detail::segmented_sort::local_segment_index_t*, detail::three_way_partition::ScanTileStateT, cub::detail::segmented_sort::LargeSegmentsSelectorT, @@ -747,8 +754,8 @@ private: ? static_cast(num_segments - current_seg_offset) : num_segments_per_invocation_limit; - large_segments_selector.base_segment_offset = current_seg_offset; - small_segments_selector.base_segment_offset = current_seg_offset; + kernel_source.SetSegmentOffset(large_segments_selector, current_seg_offset); + kernel_source.SetSegmentOffset(small_segments_selector, current_seg_offset); BeginOffsetIteratorT current_begin_offset = d_begin_offsets; EndOffsetIteratorT current_end_offset = d_end_offsets; From d584bdc596d96bd6f898327c823448db00823659 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sat, 13 Sep 2025 16:10:58 +0000 Subject: [PATCH 044/100] Clean up segmented sort c parallel tests --- c/parallel/test/test_segmented_sort.cpp | 202 +----------------------- 1 file changed, 4 insertions(+), 198 deletions(-) diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index cda67636e39..79e4a40e92d 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -28,6 +28,8 @@ using item_t = float; using BuildResultT = cccl_device_segmented_sort_build_result_t; +using SizeT = long; + struct segmented_sort_cleanup { CUresult operator()(BuildResultT* build_data) const noexcept @@ -209,9 +211,8 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes pointer_t values_in; pointer_t values_out; - // // Always use signed long long for offset iterator since negative advances are possible - // using SizeT = unsigned long long; - // static constexpr std::string_view index_ty_name = "unsigned long long"; + // TODO: Using a step counting iterator does not work right now. + // static constexpr std::string_view index_ty_name = "signed long long"; // struct segment_offset_iterator_state_t // { @@ -250,9 +251,6 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes // start_offsets_cccl.host_advance = &host_advance_linear_id; // end_offsets_cccl.host_advance = &host_advance_linear_id; - // Provide device arrays of start/end offsets instead of custom iterators - using SizeT = unsigned long long; - std::vector start_offsets(n_segments); std::vector end_offsets(n_segments); for (std::size_t i = 0; i < n_segments; ++i) @@ -343,48 +341,6 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] pointer_t values_in_ptr(host_values); pointer_t values_out_ptr(host_values_out); - // using SizeT = signed long long; - // static constexpr std::string_view index_ty_name = "signed long long"; - - // struct segment_offset_iterator_state_t - // { - // SizeT linear_id; - // SizeT segment_size; - // }; - - // static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; - // static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; - // static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; - - // const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = - // make_step_counting_iterator_sources( - // index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); - - // iterator_t start_offset_it = - // make_iterator( - // {offset_iterator_state_name, offset_iterator_state_src}, - // {advance_offset_method_name, offset_iterator_advance_src}, - // {deref_offset_method_name, offset_iterator_deref_src}); - - // start_offset_it.state.linear_id = 0; - // start_offset_it.state.segment_size = segment_size; - - // iterator_t end_offset_it = - // make_iterator( - // {offset_iterator_state_name, ""}, {advance_offset_method_name, ""}, {deref_offset_method_name, ""}); - - // end_offset_it.state.linear_id = 1; - // end_offset_it.state.segment_size = segment_size; - - // // Provide host-advance callbacks for offset iterators - // auto start_offsets_cccl = static_cast(start_offset_it); - // auto end_offsets_cccl = static_cast(end_offset_it); - // start_offsets_cccl.host_advance = &host_advance_linear_id; - // end_offsets_cccl.host_advance = &host_advance_linear_id; - - // Provide device arrays of start/end offsets instead of custom iterators - using SizeT = unsigned long long; - std::vector start_offsets(n_segments); std::vector end_offsets(n_segments); for (std::size_t i = 0; i < n_segments; ++i) @@ -612,8 +568,6 @@ C2H_TEST("SegmentedSort works with custom types as values", "[segmented_sort][cu } #endif -using SizeT = signed long long; - struct variable_segment_offset_iterator_state_t { SizeT linear_id; @@ -687,48 +641,6 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va pointer_t values_in_ptr(host_values); pointer_t values_out_ptr(host_values_out); - // // Create segment offset arrays - // std::vector start_offsets(n_segments); - // std::vector end_offsets(n_segments); - - // SizeT current_offset = 0; - // for (std::size_t i = 0; i < n_segments; ++i) - // { - // start_offsets[i] = current_offset; - // current_offset += segment_sizes[i]; - // end_offsets[i] = current_offset; - // } - - // pointer_t start_offsets_ptr(start_offsets); - // pointer_t end_offsets_ptr(end_offsets); - - // const auto& [offset_state_src, offset_advance_src, offset_deref_src] = make_variable_segment_iterator_sources(); - - // iterator_t start_offset_it = - // make_iterator( - // {"variable_segment_offset_iterator_state_t", offset_state_src}, - // {"advance_variable_offset_it", offset_advance_src}, - // {"dereference_variable_offset_it", offset_deref_src}); - - // start_offset_it.state.linear_id = 0; - // start_offset_it.state.offsets = start_offsets_ptr.ptr; - - // iterator_t end_offset_it = - // make_iterator( - // {"variable_segment_offset_iterator_state_t", ""}, - // {"advance_variable_offset_it", ""}, - // {"dereference_variable_offset_it", ""}); - - // end_offset_it.state.linear_id = 0; - // end_offset_it.state.offsets = end_offsets_ptr.ptr; - - // auto cccl_start_offsets_it = static_cast(start_offset_it); - // auto cccl_end_offsets_it = static_cast(end_offset_it); - - // // set host_advance functions - // cccl_start_offsets_it.host_advance = &host_advance_linear_id; - // cccl_end_offsets_it.host_advance = &host_advance_linear_id; - std::vector start_offsets(n_segments); std::vector end_offsets(n_segments); SizeT current_offset = 0; @@ -812,109 +724,3 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va REQUIRE(expected_keys == std::vector(output_keys)); REQUIRE(expected_values == std::vector(output_vals)); } - -// struct SegmentedSort_LargeNumSegments_Fixture_Tag; -// C2H_TEST("SegmentedSort works with a large number of segments", "[segmented_sort][large_segments]") { -// using key_t = int; - -// constexpr bool is_descending = false; -// constexpr auto order = CCCL_ASCENDING; -// constexpr bool is_overwrite_okay = false; - -// const std::size_t n_segments = 501; -// const std::size_t segment_size = 27; -// const std::size_t n_elems = n_segments * segment_size; - -// std::vector host_keys = generate(n_elems); -// std::vector host_keys_out(n_elems); - -// REQUIRE(host_keys.size() == n_elems); -// REQUIRE(host_keys_out.size() == n_elems); - -// pointer_t keys_in_ptr(host_keys); -// pointer_t keys_out_ptr(host_keys_out); - -// pointer_t values_in_ptr; -// pointer_t values_out_ptr; - -// struct segment_offset_iterator_state_t -// { -// SizeT linear_id; -// SizeT segment_size; -// }; - -// static constexpr std::string_view index_ty_name = "unsigned long long"; -// static constexpr std::string_view offset_iterator_state_name = "segment_offset_iterator_state_t"; -// static constexpr std::string_view advance_offset_method_name = "advance_offset_it"; -// static constexpr std::string_view deref_offset_method_name = "dereference_offset_it"; - -// const auto& [offset_iterator_state_src, offset_iterator_advance_src, offset_iterator_deref_src] = -// make_step_counting_iterator_sources( -// index_ty_name, offset_iterator_state_name, advance_offset_method_name, deref_offset_method_name); - -// // start = i * segment_size -// iterator_t start_offset_it = -// make_iterator( -// {offset_iterator_state_name, offset_iterator_state_src}, -// {advance_offset_method_name, offset_iterator_advance_src}, -// {deref_offset_method_name, offset_iterator_deref_src}); -// start_offset_it.state.linear_id = 0; -// start_offset_it.state.segment_size = segment_size; - -// // end = (i + 1) * segment_size -// auto end_offset_it = -// make_iterator( -// {offset_iterator_state_name, ""}, -// {advance_offset_method_name, ""}, -// {deref_offset_method_name, ""}); -// end_offset_it.state.linear_id = 1; -// end_offset_it.state.segment_size = segment_size; - -// cccl_iterator_t start_offsets_cccl(start_offset_it); -// cccl_iterator_t end_offsets_cccl(end_offset_it); -// start_offsets_cccl.host_advance = &host_advance_linear_id; -// end_offsets_cccl.host_advance = &host_advance_linear_id; - -// auto& build_cache = get_cache(); -// const std::string& key_string = KeyBuilder::join( -// {KeyBuilder::bool_as_key(is_descending), -// KeyBuilder::type_as_key(), -// KeyBuilder::type_as_key(), -// KeyBuilder::bool_as_key(is_overwrite_okay)}); -// const auto& test_key = std::make_optional(key_string); - -// int selector = -1; - -// segmented_sort( -// order, -// keys_in_ptr, -// keys_out_ptr, -// values_in_ptr, -// values_out_ptr, -// n_elems, -// n_segments, -// start_offsets_cccl, -// end_offsets_cccl, -// is_overwrite_okay, -// &selector, -// build_cache, -// test_key); - -// std::vector expected_keys = host_keys; -// for (std::size_t i = 0; i < n_segments; ++i) -// { -// const std::size_t segment_start = i * segment_size; -// const std::size_t segment_end = segment_start + segment_size; -// if (is_descending) -// { -// std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end, std::greater()); -// } -// else -// { -// std::sort(expected_keys.begin() + segment_start, expected_keys.begin() + segment_end); -// } -// } - -// auto& output_keys = (is_overwrite_okay && selector == 0) ? keys_in_ptr : keys_out_ptr; -// REQUIRE(expected_keys == std::vector(output_keys)); -// } From d657f1627e209f3c6239e4417c05dad9e4bd7a04 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sun, 14 Sep 2025 22:00:08 +0000 Subject: [PATCH 045/100] Implement dynamic dispatch for three_way_partition --- cub/cub/agent/agent_three_way_partition.cuh | 13 + cub/cub/detail/ptx-json/value.h | 5 +- .../dispatch/dispatch_three_way_partition.cuh | 274 +++++++----------- .../dispatch/kernels/three_way_partition.cuh | 80 +++++ .../tuning/tuning_three_way_partition.cuh | 37 +++ 5 files changed, 237 insertions(+), 172 deletions(-) diff --git a/cub/cub/agent/agent_three_way_partition.cuh b/cub/cub/agent/agent_three_way_partition.cuh index 5ace6dd9177..dc7402cac27 100644 --- a/cub/cub/agent/agent_three_way_partition.cuh +++ b/cub/cub/agent/agent_three_way_partition.cuh @@ -76,9 +76,22 @@ struct AgentThreeWayPartitionPolicy }; }; +#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) namespace detail { +CUB_DETAIL_POLICY_WRAPPER_DEFINE( + ThreeWayPartitionAgentPolicy, + (GenericAgentPolicy), + (BLOCK_THREADS, BlockThreads, int), + (ITEMS_PER_THREAD, ItemsPerThread, int), + (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), + (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), + (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm)) +} // namespace detail +#endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) +namespace detail +{ namespace three_way_partition { diff --git a/cub/cub/detail/ptx-json/value.h b/cub/cub/detail/ptx-json/value.h index 6aefd1c351a..311bc26d08c 100644 --- a/cub/cub/detail/ptx-json/value.h +++ b/cub/cub/detail/ptx-json/value.h @@ -71,8 +71,9 @@ struct value } }; -template -struct value +// Integral constants (matches both signed and unsigned integrals) +template +struct value, void>> { __forceinline__ __device__ static void emit() { diff --git a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index 850ea69d8bf..3b85b5f6aaf 100644 --- a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -41,7 +41,6 @@ #include #include #include -#include #include #include @@ -60,78 +59,42 @@ CUB_NAMESPACE_BEGIN namespace detail::three_way_partition { -// Offset type used to instantiate the stream three-way-partition-kernel and agent to index the items within one -// partition -using per_partition_offset_t = ::cuda::std::int32_t; -template -class streaming_context_t +template +struct DeviceThreeWayPartitionKernelSource { -private: - bool first_partition = true; - bool last_partition = false; - TotalNumItemsT total_previous_num_items{}; - - // We use a double-buffer for keeping track of the number of previously selected items - TotalNumItemsT* d_num_selected_in = nullptr; - TotalNumItemsT* d_num_selected_out = nullptr; - -public: - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE - streaming_context_t(TotalNumItemsT* d_num_selected_in, TotalNumItemsT* d_num_selected_out, bool is_last_partition) - : last_partition(is_last_partition) - , d_num_selected_in(d_num_selected_in) - , d_num_selected_out(d_num_selected_out) - {} - - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(TotalNumItemsT num_items, bool next_partition_is_the_last) + CUB_DEFINE_KERNEL_GETTER(ThreeWayPartitionInitKernel, + DeviceThreeWayPartitionInitKernel); + + CUB_DEFINE_KERNEL_GETTER( + ThreeWayPartitionKernel, + DeviceThreeWayPartitionKernel< + MaxPolicyT, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + per_partition_offset_t, + streaming_context_t>); + + CUB_RUNTIME_FUNCTION static constexpr size_t OffsetSize() { - ::cuda::std::swap(d_num_selected_in, d_num_selected_out); - first_partition = false; - last_partition = next_partition_is_the_last; - total_previous_num_items += num_items; - }; - - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT input_offset() const - { - return first_partition ? TotalNumItemsT{0} : total_previous_num_items; - }; - - _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_first() const - { - return first_partition ? TotalNumItemsT{0} : d_num_selected_in[0]; - }; - - _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_second() const - { - return first_partition ? TotalNumItemsT{0} : d_num_selected_in[1]; - }; - - _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_rejected() const - { - return first_partition ? TotalNumItemsT{0} : d_num_selected_in[2]; - ; - }; - - template - _CCCL_DEVICE _CCCL_FORCEINLINE void update_num_selected( - NumSelectedIteratorT user_num_selected_out_it, - TotalNumItemsT num_selected_first, - TotalNumItemsT num_selected_second, - TotalNumItemsT num_items_in_partition) const - { - if (last_partition) - { - user_num_selected_out_it[0] = num_previously_selected_first() + num_selected_first; - user_num_selected_out_it[1] = num_previously_selected_second() + num_selected_second; - } - else - { - d_num_selected_out[0] = num_previously_selected_first() + num_selected_first; - d_num_selected_out[1] = num_previously_selected_second() + num_selected_second; - d_num_selected_out[2] = - num_previously_rejected() + (num_items_in_partition - num_selected_second - num_selected_first); - } + return sizeof(OffsetT); } }; } // namespace detail::three_way_partition @@ -140,16 +103,31 @@ public: * Dispatch ******************************************************************************/ -template , detail::three_way_partition::per_partition_offset_t>> +template < + typename InputIteratorT, + typename FirstOutputIteratorT, + typename SecondOutputIteratorT, + typename UnselectedOutputIteratorT, + typename NumSelectedIteratorT, + typename SelectFirstPartOp, + typename SelectSecondPartOp, + typename OffsetT, + typename PolicyHub = detail::three_way_partition::policy_hub, + detail::three_way_partition::per_partition_offset_t>, + typename KernelSource = detail::three_way_partition::DeviceThreeWayPartitionKernelSource< + typename PolicyHub::MaxPolicy, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + detail::three_way_partition::ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + detail::three_way_partition::per_partition_offset_t, + detail::three_way_partition::streaming_context_t, + OffsetT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> struct DispatchThreeWayPartitionIf { /***************************************************************************** @@ -164,9 +142,7 @@ struct DispatchThreeWayPartitionIf using streaming_context_t = detail::three_way_partition::streaming_context_t; - using AccumPackHelperT = detail::three_way_partition::accumulator_pack_t; - using AccumPackT = typename AccumPackHelperT::pack_t; - using ScanTileStateT = cub::ScanTileState; + using ScanTileStateT = detail::three_way_partition::ScanTileStateT; static constexpr int INIT_KERNEL_THREADS = 256; @@ -181,31 +157,8 @@ struct DispatchThreeWayPartitionIf SelectSecondPartOp select_second_part_op; OffsetT num_items; cudaStream_t stream; - - CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchThreeWayPartitionIf( - void* d_temp_storage, - size_t& temp_storage_bytes, - InputIteratorT d_in, - FirstOutputIteratorT d_first_part_out, - SecondOutputIteratorT d_second_part_out, - UnselectedOutputIteratorT d_unselected_out, - NumSelectedIteratorT d_num_selected_out, - SelectFirstPartOp select_first_part_op, - SelectSecondPartOp select_second_part_op, - OffsetT num_items, - cudaStream_t stream) - : d_temp_storage(d_temp_storage) - , temp_storage_bytes(temp_storage_bytes) - , d_in(d_in) - , d_first_part_out(d_first_part_out) - , d_second_part_out(d_second_part_out) - , d_unselected_out(d_unselected_out) - , d_num_selected_out(d_num_selected_out) - , select_first_part_op(select_first_part_op) - , select_second_part_op(select_second_part_op) - , num_items(num_items) - , stream(stream) - {} + KernelSource kernel_source; + KernelLauncherFactory launcher_factory; /***************************************************************************** * Dispatch entrypoints @@ -213,13 +166,15 @@ struct DispatchThreeWayPartitionIf template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t - Invoke(ScanInitKernelPtrT three_way_partition_init_kernel, SelectIfKernelPtrT three_way_partition_kernel) + Invoke(ActivePolicyT policy, + ScanInitKernelPtrT three_way_partition_init_kernel, + SelectIfKernelPtrT three_way_partition_kernel) { cudaError error = cudaSuccess; - constexpr int block_threads = ActivePolicyT::ThreeWayPartitionPolicy::BLOCK_THREADS; - constexpr int items_per_thread = ActivePolicyT::ThreeWayPartitionPolicy::ITEMS_PER_THREAD; - constexpr int tile_size = block_threads * items_per_thread; + const int block_threads = policy.ThreeWayPartition().BlockThreads(); + const int items_per_thread = policy.ThreeWayPartition().ItemsPerThread(); + const int tile_size = block_threads * items_per_thread; // The maximum number of items for which we will ever invoke the kernel (i.e. largest partition size) auto const max_partition_size = static_cast( @@ -236,7 +191,7 @@ struct DispatchThreeWayPartitionIf constexpr ::cuda::std::size_t num_counters_per_pass = 3; constexpr ::cuda::std::size_t num_streaming_counters = 2 * num_counters_per_pass; ::cuda::std::size_t streaming_selection_storage_bytes = - (num_partitions > 1) ? num_streaming_counters * sizeof(OffsetT) : ::cuda::std::size_t{0}; + (num_partitions > 1) ? num_streaming_counters * kernel_source.OffsetSize() : ::cuda::std::size_t{0}; // Specify temporary storage allocation requirements size_t allocation_sizes[2] = {0ULL, streaming_selection_storage_bytes}; @@ -289,7 +244,7 @@ struct DispatchThreeWayPartitionIf } // Log three_way_partition_init_kernel configuration - int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); + const int init_grid_size = _CUDA_VSTD::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); #ifdef CUB_DEBUG_LOG _CubLog("Invoking three_way_partition_init_kernel<<<%d, %d, 0, %lld>>>()\n", @@ -299,7 +254,7 @@ struct DispatchThreeWayPartitionIf #endif // CUB_DEBUG_LOG // Invoke three_way_partition_init_kernel to initialize tile descriptors - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream) .doit(three_way_partition_init_kernel, tile_status, current_num_tiles, d_num_selected_out); // Check for failure to launch @@ -328,9 +283,10 @@ struct DispatchThreeWayPartitionIf { // Get SM occupancy for select_if_kernel int range_select_sm_occupancy; - error = CubDebug(MaxSmOccupancy(range_select_sm_occupancy, // out - three_way_partition_kernel, - block_threads)); + error = CubDebug(launcher_factory.MaxSmOccupancy( + range_select_sm_occupancy, // out + three_way_partition_kernel, + block_threads)); if (cudaSuccess != error) { return error; @@ -347,7 +303,7 @@ struct DispatchThreeWayPartitionIf #endif // CUB_DEBUG_LOG // Invoke select_if_kernel - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(current_num_tiles, block_threads, 0, stream) + launcher_factory(current_num_tiles, block_threads, 0, stream) .doit(three_way_partition_kernel, d_in, d_first_part_out, @@ -383,28 +339,16 @@ struct DispatchThreeWayPartitionIf } template - CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke() + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) { - using MaxPolicyT = typename PolicyHub::MaxPolicy; - return Invoke( - detail::three_way_partition::DeviceThreeWayPartitionInitKernel, - detail::three_way_partition::DeviceThreeWayPartitionKernel< - MaxPolicyT, - InputIteratorT, - FirstOutputIteratorT, - SecondOutputIteratorT, - UnselectedOutputIteratorT, - NumSelectedIteratorT, - ScanTileStateT, - SelectFirstPartOp, - SelectSecondPartOp, - per_partition_offset_t, - streaming_context_t>); + const auto wrapped_policy = detail::three_way_partition::MakeThreeWayPartitionPolicyWrapper(active_policy); + return Invoke(wrapped_policy, kernel_source.ThreeWayPartitionInitKernel(), kernel_source.ThreeWayPartitionKernel()); } /** * Internal dispatch routine */ + template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -416,44 +360,34 @@ struct DispatchThreeWayPartitionIf SelectFirstPartOp select_first_part_op, SelectSecondPartOp select_second_part_op, OffsetT num_items, - cudaStream_t stream) + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) { - using MaxPolicyT = typename PolicyHub::MaxPolicy; - - cudaError error = cudaSuccess; - - do + // Get PTX version + int ptx_version = 0; + if (cudaError error = CubDebug(launcher_factory.PtxVersion(ptx_version)); cudaSuccess != error) { - // Get PTX version - int ptx_version = 0; - error = CubDebug(cub::PtxVersion(ptx_version)); - if (cudaSuccess != error) - { - break; - } - - DispatchThreeWayPartitionIf dispatch( - d_temp_storage, - temp_storage_bytes, - d_in, - d_first_part_out, - d_second_part_out, - d_unselected_out, - d_num_selected_out, - select_first_part_op, - select_second_part_op, - num_items, - stream); - - // Dispatch - error = CubDebug(MaxPolicyT::Invoke(ptx_version, dispatch)); - if (cudaSuccess != error) - { - break; - } - } while (0); + return error; + } - return error; + DispatchThreeWayPartitionIf dispatch{ + d_temp_storage, + temp_storage_bytes, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + d_num_selected_out, + select_first_part_op, + select_second_part_op, + num_items, + stream, + kernel_source, + launcher_factory}; + + return CubDebug(max_policy.Invoke(ptx_version, dispatch)); } }; diff --git a/cub/cub/device/dispatch/kernels/three_way_partition.cuh b/cub/cub/device/dispatch/kernels/three_way_partition.cuh index 4efc62b8bf3..4aac47f3e29 100644 --- a/cub/cub/device/dispatch/kernels/three_way_partition.cuh +++ b/cub/cub/device/dispatch/kernels/three_way_partition.cuh @@ -19,6 +19,86 @@ CUB_NAMESPACE_BEGIN namespace detail::three_way_partition { + +// Offset type used to instantiate the stream three-way-partition-kernel and agent to index the items within one +// partition +using per_partition_offset_t = ::cuda::std::int32_t; + +using AccumPackHelperT = detail::three_way_partition::accumulator_pack_t; +using AccumPackT = typename AccumPackHelperT::pack_t; +using ScanTileStateT = cub::ScanTileState; + +template +class streaming_context_t +{ +private: + bool first_partition = true; + bool last_partition = false; + TotalNumItemsT total_previous_num_items{}; + + // We use a double-buffer for keeping track of the number of previously selected items + TotalNumItemsT* d_num_selected_in = nullptr; + TotalNumItemsT* d_num_selected_out = nullptr; + +public: + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE + streaming_context_t(TotalNumItemsT* d_num_selected_in, TotalNumItemsT* d_num_selected_out, bool is_last_partition) + : last_partition(is_last_partition) + , d_num_selected_in(d_num_selected_in) + , d_num_selected_out(d_num_selected_out) + {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(TotalNumItemsT num_items, bool next_partition_is_the_last) + { + ::cuda::std::swap(d_num_selected_in, d_num_selected_out); + first_partition = false; + last_partition = next_partition_is_the_last; + total_previous_num_items += num_items; + }; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT input_offset() const + { + return first_partition ? TotalNumItemsT{0} : total_previous_num_items; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_first() const + { + return first_partition ? TotalNumItemsT{0} : d_num_selected_in[0]; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected_second() const + { + return first_partition ? TotalNumItemsT{0} : d_num_selected_in[1]; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_rejected() const + { + return first_partition ? TotalNumItemsT{0} : d_num_selected_in[2]; + ; + }; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void update_num_selected( + NumSelectedIteratorT user_num_selected_out_it, + TotalNumItemsT num_selected_first, + TotalNumItemsT num_selected_second, + TotalNumItemsT num_items_in_partition) const + { + if (last_partition) + { + user_num_selected_out_it[0] = num_previously_selected_first() + num_selected_first; + user_num_selected_out_it[1] = num_previously_selected_second() + num_selected_second; + } + else + { + d_num_selected_out[0] = num_previously_selected_first() + num_selected_first; + d_num_selected_out[1] = num_previously_selected_second() + num_selected_second; + d_num_selected_out[2] = + num_previously_rejected() + (num_items_in_partition - num_selected_second - num_selected_first); + } + } +}; + /****************************************************************************** * Kernel entry points *****************************************************************************/ diff --git a/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh b/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh index cf2d3740f54..a1ab2b68451 100644 --- a/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh @@ -51,6 +51,43 @@ namespace detail { namespace three_way_partition { + +template +struct ThreeWayPartitionPolicyWrapper : PolicyT +{ + CUB_RUNTIME_FUNCTION ThreeWayPartitionPolicyWrapper(PolicyT base) + : PolicyT(base) + {} +}; + +template +struct ThreeWayPartitionPolicyWrapper> + : StaticPolicyT +{ + CUB_RUNTIME_FUNCTION ThreeWayPartitionPolicyWrapper(StaticPolicyT base) + : StaticPolicyT(base) + {} + + CUB_DEFINE_SUB_POLICY_GETTER(ThreeWayPartition) + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedPolicy() + { + using namespace ptx_json; + using DelayCtor = typename StaticPolicyT::ThreeWayPartitionPolicy::detail::delay_constructor_t; + return object() = ThreeWayPartition().EncodedPolicy(), + key<"ThreeWayPartitionPolicyDelayConstructor">() = + typename detail::delay_constructor_json::type()>(); + } +#endif +}; + +template +CUB_RUNTIME_FUNCTION ThreeWayPartitionPolicyWrapper MakeThreeWayPartitionPolicyWrapper(PolicyT policy) +{ + return ThreeWayPartitionPolicyWrapper{policy}; +} + enum class input_size { _1, From b798d3b6f6d855e6fb34b4e7091d4a6756d11ae3 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sun, 14 Sep 2025 22:24:53 +0000 Subject: [PATCH 046/100] Replace threshold with actual offset type --- c/parallel/src/segmented_sort.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index e563f5a7a05..2f56d26d9fd 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -204,7 +204,7 @@ struct segmented_sort_kernel_source struct selector_state_t { - long long threshold; + OffsetT threshold; const long long* begin_offsets; const long long* end_offsets; long long base_segment_offset; @@ -217,7 +217,7 @@ struct segmented_sort_kernel_source { // Persist state storage and code across the returned cccl_op_t lifetime static selector_state_t state{}; - state.threshold = static_cast(offset); + state.threshold = offset; state.begin_offsets = reinterpret_cast(*reinterpret_cast(begin_offset_iterator.ptr)); state.end_offsets = reinterpret_cast(*reinterpret_cast(end_offset_iterator.ptr)); state.base_segment_offset = 0; @@ -260,7 +260,7 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) { static selector_state_t state{}; - state.threshold = static_cast(offset); + state.threshold = offset; state.begin_offsets = reinterpret_cast(*reinterpret_cast(begin_offset_iterator.ptr)); state.end_offsets = reinterpret_cast(*reinterpret_cast(end_offset_iterator.ptr)); state.base_segment_offset = 0; From a1f4355877aac1e950c986a349e9096876c4b3d2 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sun, 14 Sep 2025 22:37:15 +0000 Subject: [PATCH 047/100] Use global_segment_offset_t type instead of long long --- c/parallel/src/segmented_sort.cu | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 2f56d26d9fd..caa89692637 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -207,7 +207,7 @@ struct segmented_sort_kernel_source OffsetT threshold; const long long* begin_offsets; const long long* end_offsets; - long long base_segment_offset; + cub::detail::segmented_sort::global_segment_offset_t base_segment_offset; }; // Return stateful cccl_op_t predicates equivalent to the CUB selectors above. @@ -224,14 +224,18 @@ struct segmented_sort_kernel_source static std::string code; code = std::string{ - R"XXX(#include + R"XXX( +#include +#include + +using cub::detail::segmented_sort::global_segment_offset_t; extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) { struct state_t { long long threshold; const long long* begin_offsets; const long long* end_offsets; - long long base_segment_offset; + global_segment_offset_t base_segment_offset; }; auto* st = static_cast(state_ptr); @@ -267,14 +271,18 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons static std::string code; code = std::string{ - R"XXX(#include + R"XXX( +#include +#include + +using cub::detail::segmented_sort::global_segment_offset_t; extern "C" __device__ void cccl_small_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) { struct state_t { long long threshold; const long long* begin_offsets; const long long* end_offsets; - long long base_segment_offset; + global_segment_offset_t base_segment_offset; }; auto* st = static_cast(state_ptr); using local_segment_index_t = ::cuda::std::uint32_t; From cf175f35a9ec29c74ef35362a65b86d53700fce1 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sun, 14 Sep 2025 23:18:37 +0000 Subject: [PATCH 048/100] Use void* for iterator types --- c/parallel/src/segmented_sort.cu | 34 +++++++++++++++++--------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index caa89692637..17ba5824deb 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -205,8 +205,8 @@ struct segmented_sort_kernel_source struct selector_state_t { OffsetT threshold; - const long long* begin_offsets; - const long long* end_offsets; + void* begin_offsets; + void* end_offsets; cub::detail::segmented_sort::global_segment_offset_t base_segment_offset; }; @@ -217,9 +217,11 @@ struct segmented_sort_kernel_source { // Persist state storage and code across the returned cccl_op_t lifetime static selector_state_t state{}; - state.threshold = offset; - state.begin_offsets = reinterpret_cast(*reinterpret_cast(begin_offset_iterator.ptr)); - state.end_offsets = reinterpret_cast(*reinterpret_cast(end_offset_iterator.ptr)); + state.threshold = offset; + // If offsets are raw device pointers, unwrap the stored pointer-to-pointer + // from the iterator state so device code can index it directly. + state.begin_offsets = *static_cast(begin_offset_iterator.ptr); + state.end_offsets = *static_cast(end_offset_iterator.ptr); state.base_segment_offset = 0; static std::string code; @@ -233,16 +235,16 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons { struct state_t { long long threshold; - const long long* begin_offsets; - const long long* end_offsets; + void* begin_offsets; + void* end_offsets; global_segment_offset_t base_segment_offset; }; auto* st = static_cast(state_ptr); using local_segment_index_t = ::cuda::std::uint32_t; const local_segment_index_t sid = *static_cast(arg_ptr); - const long long begin = st->begin_offsets[st->base_segment_offset + sid]; - const long long end = st->end_offsets[st->base_segment_offset + sid]; + const long long begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; + const long long end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; const bool pred = (end - begin) > st->threshold; *reinterpret_cast(result_ptr) = pred; } @@ -264,9 +266,9 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) { static selector_state_t state{}; - state.threshold = offset; - state.begin_offsets = reinterpret_cast(*reinterpret_cast(begin_offset_iterator.ptr)); - state.end_offsets = reinterpret_cast(*reinterpret_cast(end_offset_iterator.ptr)); + state.threshold = offset; + state.begin_offsets = *static_cast(begin_offset_iterator.ptr); + state.end_offsets = *static_cast(end_offset_iterator.ptr); state.base_segment_offset = 0; static std::string code; @@ -280,15 +282,15 @@ extern "C" __device__ void cccl_small_segments_selector_op(void* state_ptr, cons { struct state_t { long long threshold; - const long long* begin_offsets; - const long long* end_offsets; + void* begin_offsets; + void* end_offsets; global_segment_offset_t base_segment_offset; }; auto* st = static_cast(state_ptr); using local_segment_index_t = ::cuda::std::uint32_t; const local_segment_index_t sid = *static_cast(arg_ptr); - const long long begin = st->begin_offsets[st->base_segment_offset + sid]; - const long long end = st->end_offsets[st->base_segment_offset + sid]; + const long long begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; + const long long end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; const bool pred = (end - begin) < st->threshold; *reinterpret_cast(result_ptr) = pred; From a08541b7116597f6905d4e8eb22b329db2ab870f Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sun, 14 Sep 2025 23:58:22 +0000 Subject: [PATCH 049/100] Make selector op states part of build instead of static storage --- c/parallel/include/cccl/c/segmented_sort.h | 2 + c/parallel/src/segmented_sort.cu | 97 ++++++++++++++-------- 2 files changed, 64 insertions(+), 35 deletions(-) diff --git a/c/parallel/include/cccl/c/segmented_sort.h b/c/parallel/include/cccl/c/segmented_sort.h index 1950f481455..4fe621a7a3e 100644 --- a/c/parallel/include/cccl/c/segmented_sort.h +++ b/c/parallel/include/cccl/c/segmented_sort.h @@ -31,6 +31,8 @@ typedef struct cccl_device_segmented_sort_build_result_t CUlibrary library; cccl_type_info key_type; cccl_type_info offset_type; + void* large_segments_selector_op_state; + void* small_segments_selector_op_state; CUkernel segmented_sort_fallback_kernel; CUkernel segmented_sort_kernel_small; CUkernel segmented_sort_kernel_large; diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 17ba5824deb..e523f708c2e 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -151,6 +151,14 @@ std::string get_device_segmented_sort_kernel_large_name( offset_t); // 5 } +struct selector_state_t +{ + OffsetT threshold; + void* begin_offsets; + void* end_offsets; + cub::detail::segmented_sort::global_segment_offset_t base_segment_offset; +}; + struct segmented_sort_kernel_source { cccl_device_segmented_sort_build_result_t& build; @@ -179,14 +187,22 @@ struct segmented_sort_kernel_source indirect_arg_t LargeSegmentsSelector( OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const { - cccl_op_t op = LargeSegmentsSelectorOp(offset, begin_offset_iterator, end_offset_iterator); + cccl_op_t op = LargeSegmentsSelectorOp( + offset, + begin_offset_iterator, + end_offset_iterator, + static_cast(build.large_segments_selector_op_state)); return indirect_arg_t{op}; } indirect_arg_t SmallSegmentsSelector( OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const { - cccl_op_t op = SmallSegmentsSelectorOp(offset, begin_offset_iterator, end_offset_iterator); + cccl_op_t op = SmallSegmentsSelectorOp( + offset, + begin_offset_iterator, + end_offset_iterator, + static_cast(build.small_segments_selector_op_state)); return indirect_arg_t{op}; } @@ -202,27 +218,30 @@ struct segmented_sort_kernel_source st->base_segment_offset = base_segment_offset; } - struct selector_state_t + static void initialize_state( + selector_state_t* state, + OffsetT offset, + indirect_iterator_t begin_offset_iterator, + indirect_iterator_t end_offset_iterator) { - OffsetT threshold; - void* begin_offsets; - void* end_offsets; - cub::detail::segmented_sort::global_segment_offset_t base_segment_offset; - }; + state->threshold = offset; + // If offsets are raw device pointers, unwrap the stored pointer-to-pointer + // from the iterator state so device code can index it directly. + state->begin_offsets = *static_cast(begin_offset_iterator.ptr); + state->end_offsets = *static_cast(end_offset_iterator.ptr); + state->base_segment_offset = 0; + } // Return stateful cccl_op_t predicates equivalent to the CUB selectors above. // These embed C++ source for a device function and capture state (threshold and offset arrays). static cccl_op_t LargeSegmentsSelectorOp( - OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) + OffsetT offset, + indirect_iterator_t begin_offset_iterator, + indirect_iterator_t end_offset_iterator, + selector_state_t* state) { // Persist state storage and code across the returned cccl_op_t lifetime - static selector_state_t state{}; - state.threshold = offset; - // If offsets are raw device pointers, unwrap the stored pointer-to-pointer - // from the iterator state so device code can index it directly. - state.begin_offsets = *static_cast(begin_offset_iterator.ptr); - state.end_offsets = *static_cast(end_offset_iterator.ptr); - state.base_segment_offset = 0; + initialize_state(state, offset, begin_offset_iterator, end_offset_iterator); static std::string code; code = std::string{ @@ -256,20 +275,19 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons op.code = code.c_str(); op.code_size = code.size(); op.code_type = CCCL_OP_CPP_SOURCE; - op.size = sizeof(state); + op.size = sizeof(selector_state_t); op.alignment = alignof(selector_state_t); - op.state = &state; + op.state = state; return op; } static cccl_op_t SmallSegmentsSelectorOp( - OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) + OffsetT offset, + indirect_iterator_t begin_offset_iterator, + indirect_iterator_t end_offset_iterator, + selector_state_t* state) { - static selector_state_t state{}; - state.threshold = offset; - state.begin_offsets = *static_cast(begin_offset_iterator.ptr); - state.end_offsets = *static_cast(end_offset_iterator.ptr); - state.base_segment_offset = 0; + initialize_state(state, offset, begin_offset_iterator, end_offset_iterator); static std::string code; code = std::string{ @@ -303,9 +321,9 @@ extern "C" __device__ void cccl_small_segments_selector_op(void* state_ptr, cons op.code = code.c_str(); op.code_size = code.size(); op.code_type = CCCL_OP_CPP_SOURCE; - op.size = sizeof(state); + op.size = sizeof(selector_state_t); op.alignment = alignof(selector_state_t); - op.state = &state; + op.state = state; return op; } }; @@ -634,11 +652,14 @@ CUresult cccl_device_segmented_sort_build( const std::string value_t = keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); + auto* large_segments_selector_op_state = new segmented_sort::selector_state_t{}; + auto* small_segments_selector_op_state = new segmented_sort::selector_state_t{}; + // Build selector operations as cccl_op_t and generate their functor wrappers - cccl_op_t large_selector_op = - segmented_sort::segmented_sort_kernel_source::LargeSegmentsSelectorOp(0, start_offset_it, end_offset_it); - cccl_op_t small_selector_op = - segmented_sort::segmented_sort_kernel_source::SmallSegmentsSelectorOp(0, start_offset_it, end_offset_it); + cccl_op_t large_selector_op = segmented_sort::segmented_sort_kernel_source::LargeSegmentsSelectorOp( + 0, start_offset_it, end_offset_it, large_segments_selector_op_state); + cccl_op_t small_selector_op = segmented_sort::segmented_sort_kernel_source::SmallSegmentsSelectorOp( + 0, start_offset_it, end_offset_it, small_segments_selector_op_state); cccl_type_info bool_t{sizeof(bool), alignof(bool), cccl_type_enum::CCCL_BOOLEAN}; cccl_type_info u32_t{sizeof(::cuda::std::uint32_t), alignof(::cuda::std::uint32_t), cccl_type_enum::CCCL_UINT32}; @@ -858,10 +879,12 @@ struct device_three_way_partition_policy {{ check(cuLibraryGetKernel( &build_ptr->three_way_partition_kernel, build_ptr->library, three_way_partition_kernel_lowered_name.c_str())); - build_ptr->cc = cc; - build_ptr->cubin = (void*) result.data.release(); - build_ptr->cubin_size = result.size; - build_ptr->key_type = keys_in_it.value_type; + build_ptr->cc = cc; + build_ptr->large_segments_selector_op_state = large_segments_selector_op_state; + build_ptr->small_segments_selector_op_state = small_segments_selector_op_state; + build_ptr->cubin = (void*) result.data.release(); + build_ptr->cubin_size = result.size; + build_ptr->key_type = keys_in_it.value_type; build_ptr->offset_type = cccl_type_info{sizeof(OffsetT), alignof(OffsetT), cccl_type_enum::CCCL_INT64}; // Use the runtime policy extracted via from_json build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_tuning_policy{ @@ -1023,7 +1046,11 @@ CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_res // allocation behind cubin is owned by unique_ptr with delete[] deleter now std::unique_ptr cubin(reinterpret_cast(build_ptr->cubin)); - // Clean up the runtime policy + // Clean up the selector op states + delete static_cast(build_ptr->large_segments_selector_op_state); + delete static_cast(build_ptr->small_segments_selector_op_state); + + // Clean up the runtime policies delete static_cast(build_ptr->runtime_policy); delete static_cast(build_ptr->partition_runtime_policy); check(cuLibraryUnload(build_ptr->library)); From c284814010c192de7389af5528ca69b56d85008a Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 15 Sep 2025 00:04:40 +0000 Subject: [PATCH 050/100] Use existing type alias instead of redefining one --- c/parallel/src/segmented_sort.cu | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index e523f708c2e..493a2458dfd 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -249,7 +249,9 @@ struct segmented_sort_kernel_source #include #include +using cub::detail::segmented_sort::local_segment_index_t; using cub::detail::segmented_sort::global_segment_offset_t; + extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) { struct state_t { @@ -260,7 +262,6 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons }; auto* st = static_cast(state_ptr); - using local_segment_index_t = ::cuda::std::uint32_t; const local_segment_index_t sid = *static_cast(arg_ptr); const long long begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; const long long end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; @@ -295,7 +296,9 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons #include #include +using cub::detail::segmented_sort::local_segment_index_t; using cub::detail::segmented_sort::global_segment_offset_t; + extern "C" __device__ void cccl_small_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) { struct state_t { @@ -305,7 +308,6 @@ extern "C" __device__ void cccl_small_segments_selector_op(void* state_ptr, cons global_segment_offset_t base_segment_offset; }; auto* st = static_cast(state_ptr); - using local_segment_index_t = ::cuda::std::uint32_t; const local_segment_index_t sid = *static_cast(arg_ptr); const long long begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; const long long end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; From 727f13e7af5318662150338d94fb46478e08accc Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 15 Sep 2025 00:37:44 +0000 Subject: [PATCH 051/100] Fix dangling pointer error in indirect arg --- c/parallel/src/util/indirect_arg.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/c/parallel/src/util/indirect_arg.h b/c/parallel/src/util/indirect_arg.h index 97977ea2ca9..d74341a09f7 100644 --- a/c/parallel/src/util/indirect_arg.h +++ b/c/parallel/src/util/indirect_arg.h @@ -25,7 +25,7 @@ struct indirect_arg_t {} indirect_arg_t(cccl_op_t& op) - : ptr(op.type == cccl_op_kind_t::CCCL_STATEFUL ? op.state : this) + : ptr(op.type == cccl_op_kind_t::CCCL_STATEFUL ? op.state : &op) {} indirect_arg_t(cccl_value_t& val) From 12d987d4a39f48c854e1c863ec73e66c34bbced5 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 15 Sep 2025 00:56:34 +0000 Subject: [PATCH 052/100] Avoid calling static function and don't store op state since it is already stored in op --- c/parallel/include/cccl/c/segmented_sort.h | 4 +- c/parallel/src/segmented_sort.cu | 62 +++++++++------------- 2 files changed, 28 insertions(+), 38 deletions(-) diff --git a/c/parallel/include/cccl/c/segmented_sort.h b/c/parallel/include/cccl/c/segmented_sort.h index 4fe621a7a3e..4ba99db820a 100644 --- a/c/parallel/include/cccl/c/segmented_sort.h +++ b/c/parallel/include/cccl/c/segmented_sort.h @@ -31,8 +31,8 @@ typedef struct cccl_device_segmented_sort_build_result_t CUlibrary library; cccl_type_info key_type; cccl_type_info offset_type; - void* large_segments_selector_op_state; - void* small_segments_selector_op_state; + cccl_op_t large_segments_selector_op; + cccl_op_t small_segments_selector_op; CUkernel segmented_sort_fallback_kernel; CUkernel segmented_sort_kernel_small; CUkernel segmented_sort_kernel_large; diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 493a2458dfd..14a45d56d23 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -157,6 +157,16 @@ struct selector_state_t void* begin_offsets; void* end_offsets; cub::detail::segmented_sort::global_segment_offset_t base_segment_offset; + + void initialize(OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) + { + threshold = offset; + // If offsets are raw device pointers, unwrap the stored pointer-to-pointer + // from the iterator state so device code can index it directly. + begin_offsets = *static_cast(begin_offset_iterator.ptr); + end_offsets = *static_cast(end_offset_iterator.ptr); + base_segment_offset = 0; + } }; struct segmented_sort_kernel_source @@ -187,23 +197,17 @@ struct segmented_sort_kernel_source indirect_arg_t LargeSegmentsSelector( OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const { - cccl_op_t op = LargeSegmentsSelectorOp( - offset, - begin_offset_iterator, - end_offset_iterator, - static_cast(build.large_segments_selector_op_state)); - return indirect_arg_t{op}; + static_cast(build.large_segments_selector_op.state) + ->initialize(offset, begin_offset_iterator, end_offset_iterator); + return indirect_arg_t(build.large_segments_selector_op); } indirect_arg_t SmallSegmentsSelector( OffsetT offset, indirect_iterator_t begin_offset_iterator, indirect_iterator_t end_offset_iterator) const { - cccl_op_t op = SmallSegmentsSelectorOp( - offset, - begin_offset_iterator, - end_offset_iterator, - static_cast(build.small_segments_selector_op_state)); - return indirect_arg_t{op}; + static_cast(build.small_segments_selector_op.state) + ->initialize(offset, begin_offset_iterator, end_offset_iterator); + return indirect_arg_t(build.small_segments_selector_op); } void SetSegmentOffset(cccl_op_t& selector, long long base_segment_offset) const @@ -218,20 +222,6 @@ struct segmented_sort_kernel_source st->base_segment_offset = base_segment_offset; } - static void initialize_state( - selector_state_t* state, - OffsetT offset, - indirect_iterator_t begin_offset_iterator, - indirect_iterator_t end_offset_iterator) - { - state->threshold = offset; - // If offsets are raw device pointers, unwrap the stored pointer-to-pointer - // from the iterator state so device code can index it directly. - state->begin_offsets = *static_cast(begin_offset_iterator.ptr); - state->end_offsets = *static_cast(end_offset_iterator.ptr); - state->base_segment_offset = 0; - } - // Return stateful cccl_op_t predicates equivalent to the CUB selectors above. // These embed C++ source for a device function and capture state (threshold and offset arrays). static cccl_op_t LargeSegmentsSelectorOp( @@ -241,7 +231,7 @@ struct segmented_sort_kernel_source selector_state_t* state) { // Persist state storage and code across the returned cccl_op_t lifetime - initialize_state(state, offset, begin_offset_iterator, end_offset_iterator); + state->initialize(offset, begin_offset_iterator, end_offset_iterator); static std::string code; code = std::string{ @@ -288,7 +278,7 @@ extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, cons indirect_iterator_t end_offset_iterator, selector_state_t* state) { - initialize_state(state, offset, begin_offset_iterator, end_offset_iterator); + state->initialize(offset, begin_offset_iterator, end_offset_iterator); static std::string code; code = std::string{ @@ -881,12 +871,12 @@ struct device_three_way_partition_policy {{ check(cuLibraryGetKernel( &build_ptr->three_way_partition_kernel, build_ptr->library, three_way_partition_kernel_lowered_name.c_str())); - build_ptr->cc = cc; - build_ptr->large_segments_selector_op_state = large_segments_selector_op_state; - build_ptr->small_segments_selector_op_state = small_segments_selector_op_state; - build_ptr->cubin = (void*) result.data.release(); - build_ptr->cubin_size = result.size; - build_ptr->key_type = keys_in_it.value_type; + build_ptr->cc = cc; + build_ptr->large_segments_selector_op = large_selector_op; + build_ptr->small_segments_selector_op = small_selector_op; + build_ptr->cubin = (void*) result.data.release(); + build_ptr->cubin_size = result.size; + build_ptr->key_type = keys_in_it.value_type; build_ptr->offset_type = cccl_type_info{sizeof(OffsetT), alignof(OffsetT), cccl_type_enum::CCCL_INT64}; // Use the runtime policy extracted via from_json build_ptr->runtime_policy = new segmented_sort::segmented_sort_runtime_tuning_policy{ @@ -1049,8 +1039,8 @@ CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_res std::unique_ptr cubin(reinterpret_cast(build_ptr->cubin)); // Clean up the selector op states - delete static_cast(build_ptr->large_segments_selector_op_state); - delete static_cast(build_ptr->small_segments_selector_op_state); + delete static_cast(build_ptr->large_segments_selector_op.state); + delete static_cast(build_ptr->small_segments_selector_op.state); // Clean up the runtime policies delete static_cast(build_ptr->runtime_policy); From d956f96c640ec1e2874ebff58859170b98247312 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 15 Sep 2025 02:19:59 +0000 Subject: [PATCH 053/100] Refactor to avoid code duplication --- c/parallel/src/segmented_sort.cu | 171 ++++++++++++------------------- 1 file changed, 67 insertions(+), 104 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 14a45d56d23..a803718c728 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -21,8 +21,6 @@ #include // std::string_view #include // std::is_same_v -#include // printf - #include "jit_templates/templates/input_iterator.h" #include "jit_templates/templates/operation.h" #include "jit_templates/templates/output_iterator.h" @@ -169,6 +167,66 @@ struct selector_state_t } }; +cccl_op_t make_segments_selector_op( + OffsetT offset, + cccl_iterator_t begin_offset_iterator, + cccl_iterator_t end_offset_iterator, + const char* selector_op_name, + const char* comparison) +{ + cccl_op_t selector_op{}; + selector_state_t* selector_op_state = new selector_state_t{}; + std::string offset_t; + check(nvrtcGetTypeName(&offset_t)); + + const std::string code = std::format( + R"XXX( +#include +#include + +using cub::detail::segmented_sort::local_segment_index_t; +using cub::detail::segmented_sort::global_segment_offset_t; + +extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* result_ptr) +{{ + struct state_t {{ + {1} threshold; + void* begin_offsets; + void* end_offsets; + global_segment_offset_t base_segment_offset; + }}; + + auto* st = static_cast(state_ptr); + const local_segment_index_t sid = *static_cast(arg_ptr); + const {2} begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; + const {3} end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; + const bool pred = (end - begin) {4} st->threshold; + *reinterpret_cast(result_ptr) = pred; +}} +)XXX", + selector_op_name, + offset_t, + cccl_type_enum_to_name(begin_offset_iterator.value_type.type), + cccl_type_enum_to_name(end_offset_iterator.value_type.type), + comparison); + + selector_op.type = cccl_op_kind_t::CCCL_STATEFUL; + selector_op.name = selector_op_name; + // Allocate persistent storage for the generated source code + char* code_buf = new char[code.size() + 1]; + std::memcpy(code_buf, code.c_str(), code.size() + 1); + selector_op.code = code_buf; + selector_op.code_size = code.size(); + selector_op.code_type = CCCL_OP_CPP_SOURCE; + selector_op.size = sizeof(selector_state_t); + selector_op.alignment = alignof(selector_state_t); + selector_op.state = selector_op_state; + + selector_op_state->initialize(offset, begin_offset_iterator, end_offset_iterator); + + return selector_op; +} + struct segmented_sort_kernel_source { cccl_device_segmented_sort_build_result_t& build; @@ -221,103 +279,6 @@ struct segmented_sort_kernel_source auto* st = reinterpret_cast(selector.ptr); st->base_segment_offset = base_segment_offset; } - - // Return stateful cccl_op_t predicates equivalent to the CUB selectors above. - // These embed C++ source for a device function and capture state (threshold and offset arrays). - static cccl_op_t LargeSegmentsSelectorOp( - OffsetT offset, - indirect_iterator_t begin_offset_iterator, - indirect_iterator_t end_offset_iterator, - selector_state_t* state) - { - // Persist state storage and code across the returned cccl_op_t lifetime - state->initialize(offset, begin_offset_iterator, end_offset_iterator); - - static std::string code; - code = std::string{ - R"XXX( -#include -#include - -using cub::detail::segmented_sort::local_segment_index_t; -using cub::detail::segmented_sort::global_segment_offset_t; - -extern "C" __device__ void cccl_large_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) -{ - struct state_t { - long long threshold; - void* begin_offsets; - void* end_offsets; - global_segment_offset_t base_segment_offset; - }; - - auto* st = static_cast(state_ptr); - const local_segment_index_t sid = *static_cast(arg_ptr); - const long long begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; - const long long end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; - const bool pred = (end - begin) > st->threshold; - *reinterpret_cast(result_ptr) = pred; -} -)XXX"}; - - cccl_op_t op{}; - op.type = cccl_op_kind_t::CCCL_STATEFUL; - op.name = "cccl_large_segments_selector_op"; - op.code = code.c_str(); - op.code_size = code.size(); - op.code_type = CCCL_OP_CPP_SOURCE; - op.size = sizeof(selector_state_t); - op.alignment = alignof(selector_state_t); - op.state = state; - return op; - } - - static cccl_op_t SmallSegmentsSelectorOp( - OffsetT offset, - indirect_iterator_t begin_offset_iterator, - indirect_iterator_t end_offset_iterator, - selector_state_t* state) - { - state->initialize(offset, begin_offset_iterator, end_offset_iterator); - - static std::string code; - code = std::string{ - R"XXX( -#include -#include - -using cub::detail::segmented_sort::local_segment_index_t; -using cub::detail::segmented_sort::global_segment_offset_t; - -extern "C" __device__ void cccl_small_segments_selector_op(void* state_ptr, const void* arg_ptr, void* result_ptr) -{ - struct state_t { - long long threshold; - void* begin_offsets; - void* end_offsets; - global_segment_offset_t base_segment_offset; - }; - auto* st = static_cast(state_ptr); - const local_segment_index_t sid = *static_cast(arg_ptr); - const long long begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; - const long long end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; - const bool pred = (end - begin) < st->threshold; - - *reinterpret_cast(result_ptr) = pred; -} -)XXX"}; - - cccl_op_t op{}; - op.type = cccl_op_kind_t::CCCL_STATEFUL; - op.name = "cccl_small_segments_selector_op"; - op.code = code.c_str(); - op.code_size = code.size(); - op.code_type = CCCL_OP_CPP_SOURCE; - op.size = sizeof(selector_state_t); - op.alignment = alignof(selector_state_t); - op.state = state; - return op; - } }; std::string get_three_way_partition_init_kernel_name() @@ -647,11 +608,10 @@ CUresult cccl_device_segmented_sort_build( auto* large_segments_selector_op_state = new segmented_sort::selector_state_t{}; auto* small_segments_selector_op_state = new segmented_sort::selector_state_t{}; - // Build selector operations as cccl_op_t and generate their functor wrappers - cccl_op_t large_selector_op = segmented_sort::segmented_sort_kernel_source::LargeSegmentsSelectorOp( - 0, start_offset_it, end_offset_it, large_segments_selector_op_state); - cccl_op_t small_selector_op = segmented_sort::segmented_sort_kernel_source::SmallSegmentsSelectorOp( - 0, start_offset_it, end_offset_it, small_segments_selector_op_state); + cccl_op_t large_selector_op = segmented_sort::make_segments_selector_op( + 0, start_offset_it, end_offset_it, "cccl_large_segments_selector_op", ">"); + cccl_op_t small_selector_op = segmented_sort::make_segments_selector_op( + 0, start_offset_it, end_offset_it, "cccl_small_segments_selector_op", "<"); cccl_type_info bool_t{sizeof(bool), alignof(bool), cccl_type_enum::CCCL_BOOLEAN}; cccl_type_info u32_t{sizeof(::cuda::std::uint32_t), alignof(::cuda::std::uint32_t), cccl_type_enum::CCCL_UINT32}; @@ -1042,6 +1002,9 @@ CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_res delete static_cast(build_ptr->large_segments_selector_op.state); delete static_cast(build_ptr->small_segments_selector_op.state); + delete[] const_cast(build_ptr->large_segments_selector_op.code); + delete[] const_cast(build_ptr->small_segments_selector_op.code); + // Clean up the runtime policies delete static_cast(build_ptr->runtime_policy); delete static_cast(build_ptr->partition_runtime_policy); From 4c36550a5b6dae496d2db0bbf65f9b6e0c404b54 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 15 Sep 2025 02:47:17 +0000 Subject: [PATCH 054/100] Continue cleaning up code --- c/parallel/src/segmented_sort.cu | 54 ++++++++----------------- c/parallel/test/test_segmented_sort.cpp | 2 - 2 files changed, 16 insertions(+), 40 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index a803718c728..0c298edde0f 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -184,11 +184,11 @@ cccl_op_t make_segments_selector_op( #include #include -using cub::detail::segmented_sort::local_segment_index_t; -using cub::detail::segmented_sort::global_segment_offset_t; - extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* result_ptr) {{ + using cub::detail::segmented_sort::local_segment_index_t; + using cub::detail::segmented_sort::global_segment_offset_t; + struct state_t {{ {1} threshold; void* begin_offsets; @@ -201,7 +201,7 @@ extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* resul const {2} begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; const {3} end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; const bool pred = (end - begin) {4} st->threshold; - *reinterpret_cast(result_ptr) = pred; + *static_cast(result_ptr) = pred; }} )XXX", selector_op_name, @@ -212,7 +212,7 @@ extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* resul selector_op.type = cccl_op_kind_t::CCCL_STATEFUL; selector_op.name = selector_op_name; - // Allocate persistent storage for the generated source code + // Allocate persistent storage for the generated source code. TODO: can we use LTO-IR instead? char* code_buf = new char[code.size() + 1]; std::memcpy(code_buf, code.c_str(), code.size() + 1); selector_op.code = code_buf; @@ -268,12 +268,6 @@ struct segmented_sort_kernel_source return indirect_arg_t(build.small_segments_selector_op); } - void SetSegmentOffset(cccl_op_t& selector, long long base_segment_offset) const - { - auto* st = reinterpret_cast(selector.state); - st->base_segment_offset = base_segment_offset; - } - void SetSegmentOffset(indirect_arg_t& selector, long long base_segment_offset) const { auto* st = reinterpret_cast(selector.ptr); @@ -562,27 +556,6 @@ CUresult cccl_device_segmented_sort_build( template_id(), values_out_it, values_in_it.value_type); values_out_iterator_name = vo_name; values_out_iterator_src = vo_src; - - // For STORAGE values, ensure pointer types in iterator names/sources use items_storage_t* - if (values_in_it.value_type.type == cccl_type_enum::CCCL_STORAGE) - { - auto replace_all = [](std::string& s, const std::string& from, const std::string& to) { - if (from.empty()) - { - return; - } - size_t pos = 0; - while ((pos = s.find(from, pos)) != std::string::npos) - { - s.replace(pos, from.length(), to); - pos += to.length(); - } - }; - replace_all(values_in_iterator_src, "storage_t", "items_storage_t"); - replace_all(values_out_iterator_src, "storage_t", "items_storage_t"); - replace_all(values_in_iterator_name, "storage_t", "items_storage_t"); - replace_all(values_out_iterator_name, "storage_t", "items_storage_t"); - } } else { @@ -613,14 +586,17 @@ CUresult cccl_device_segmented_sort_build( cccl_op_t small_selector_op = segmented_sort::make_segments_selector_op( 0, start_offset_it, end_offset_it, "cccl_small_segments_selector_op", "<"); - cccl_type_info bool_t{sizeof(bool), alignof(bool), cccl_type_enum::CCCL_BOOLEAN}; - cccl_type_info u32_t{sizeof(::cuda::std::uint32_t), alignof(::cuda::std::uint32_t), cccl_type_enum::CCCL_UINT32}; + cccl_type_info selector_result_t{sizeof(bool), alignof(bool), cccl_type_enum::CCCL_BOOLEAN}; + cccl_type_info selector_input_t{ + sizeof(cub::detail::segmented_sort::local_segment_index_t), + alignof(cub::detail::segmented_sort::local_segment_index_t), + cccl_type_enum::CCCL_UINT32}; const auto [large_selector_name, large_selector_src] = get_specialization( - template_id(), large_selector_op, bool_t, u32_t); + template_id(), large_selector_op, selector_result_t, selector_input_t); const auto [small_selector_name, small_selector_src] = get_specialization( - template_id(), small_selector_op, bool_t, u32_t); + template_id(), small_selector_op, selector_result_t, selector_input_t); const std::string dependent_definitions_src = std::format( R"XXX( @@ -667,7 +643,9 @@ struct __align__({3}) items_storage_t {{ value_t); // 1 static constexpr std::string_view ptx_query_tu_src_tmpl = R"XXXX( +#include #include +#include #include {0} {1} @@ -695,8 +673,8 @@ struct __align__({3}) items_storage_t {{ R"XXXX(cub::detail::three_way_partition::MakeThreeWayPartitionPolicyWrapper(cub::detail::three_way_partition::policy_hub<{0}, {1}>::MaxPolicy::ActivePolicy{{}}))XXXX"; const auto partition_policy_wrapper_expr = std::format( partition_policy_wrapper_expr_tmpl, - "::cuda::std::uint32_t", // This is local_segment_index_t defined in segmented_sort.cuh - "::cuda::std::int32_t"); // This is per_partition_offset_t defined in segmented_sort.cuh + "cub::detail::segmented_sort::local_segment_index_t", + "cub::detail::three_way_partition::per_partition_offset_t"); nlohmann::json partition_policy = get_policy(partition_policy_wrapper_expr, ptx_query_tu_src, ptx_args); diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 79e4a40e92d..a52193d363e 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -456,8 +456,6 @@ C2H_TEST("SegmentedSort works with custom types as values", "[segmented_sort][cu const std::size_t n_segments = GENERATE(0, 13, take(2, random(1 << 10, 1 << 12))); const std::size_t segment_size = GENERATE(1, 12, take(2, random(1 << 10, 1 << 12))); - std::cout << "n_segments: " << n_segments << ", segment_size: " << segment_size << std::endl; - const std::size_t n_elems = n_segments * segment_size; // Generate primitive keys From 795d59e06e65b3e58dbcfc0d9f479c14b0f41279 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 15 Sep 2025 03:16:50 +0000 Subject: [PATCH 055/100] Replace CPP_SOURCE op with LTOIR --- c/parallel/src/segmented_sort.cu | 62 ++++++++++++++++---------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 0c298edde0f..20ed77af919 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -172,7 +172,11 @@ cccl_op_t make_segments_selector_op( cccl_iterator_t begin_offset_iterator, cccl_iterator_t end_offset_iterator, const char* selector_op_name, - const char* comparison) + const char* comparison, + const char** compile_args, + size_t num_compile_args, + const char** lto_opts, + size_t num_lto_opts) { cccl_op_t selector_op{}; selector_state_t* selector_op_state = new selector_state_t{}; @@ -181,7 +185,6 @@ cccl_op_t make_segments_selector_op( const std::string code = std::format( R"XXX( -#include #include extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* result_ptr) @@ -212,12 +215,14 @@ extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* resul selector_op.type = cccl_op_kind_t::CCCL_STATEFUL; selector_op.name = selector_op_name; - // Allocate persistent storage for the generated source code. TODO: can we use LTO-IR instead? - char* code_buf = new char[code.size() + 1]; - std::memcpy(code_buf, code.c_str(), code.size() + 1); - selector_op.code = code_buf; - selector_op.code_size = code.size(); - selector_op.code_type = CCCL_OP_CPP_SOURCE; + auto [lto_size, lto_buf] = + begin_linking_nvrtc_program(static_cast(num_lto_opts), lto_opts) + ->add_program(nvrtc_translation_unit{code.c_str(), selector_op_name}) + ->compile_program({compile_args, num_compile_args}) + ->get_program_ltoir(); + selector_op.code = lto_buf.release(); + selector_op.code_size = lto_size; + selector_op.code_type = CCCL_OP_LTOIR; selector_op.size = sizeof(selector_state_t); selector_op.alignment = alignof(selector_state_t); selector_op.state = selector_op_state; @@ -578,13 +583,27 @@ CUresult cccl_device_segmented_sort_build( const std::string value_t = keys_only ? "cub::NullType" : cccl_type_enum_to_name(values_in_it.value_type.type); - auto* large_segments_selector_op_state = new segmented_sort::selector_state_t{}; - auto* small_segments_selector_op_state = new segmented_sort::selector_state_t{}; + const std::string arch = std::format("-arch=sm_{0}{1}", cc_major, cc_minor); + + constexpr size_t num_args = 9; + const char* args[num_args] = { + arch.c_str(), + cub_path, + thrust_path, + libcudacxx_path, + ctk_path, + "-rdc=true", + "-dlto", + "-DCUB_DISABLE_CDP", + "-std=c++20"}; + + constexpr size_t num_lto_args = 2; + const char* lopts[num_lto_args] = {"-lto", arch.c_str()}; cccl_op_t large_selector_op = segmented_sort::make_segments_selector_op( - 0, start_offset_it, end_offset_it, "cccl_large_segments_selector_op", ">"); + 0, start_offset_it, end_offset_it, "cccl_large_segments_selector_op", ">", args, num_args, lopts, num_lto_args); cccl_op_t small_selector_op = segmented_sort::make_segments_selector_op( - 0, start_offset_it, end_offset_it, "cccl_small_segments_selector_op", "<"); + 0, start_offset_it, end_offset_it, "cccl_small_segments_selector_op", "<", args, num_args, lopts, num_lto_args); cccl_type_info selector_result_t{sizeof(bool), alignof(bool), cccl_type_enum::CCCL_BOOLEAN}; cccl_type_info selector_input_t{ @@ -628,7 +647,6 @@ struct __align__({3}) items_storage_t {{ large_selector_src, // 10 small_selector_src); // 11 - // Runtime parameter tuning const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); constexpr size_t ptx_num_args = 5; @@ -685,7 +703,6 @@ struct __align__({3}) items_storage_t {{ const std::string three_way_partition_policy_delay_constructor = segmented_sort::get_three_way_partition_policy_delay_constructor(partition_policy); - // Inject delay constructor alias into the ThreeWayPartitionPolicy struct string const std::string injected_three_way_partition_policy_str = segmented_sort::inject_delay_constructor_into_three_way_policy( three_way_partition_policy_str, three_way_partition_policy_delay_constructor); @@ -741,23 +758,6 @@ struct device_three_way_partition_policy {{ std::string three_way_partition_init_kernel_lowered_name; std::string three_way_partition_kernel_lowered_name; - const std::string arch = std::format("-arch=sm_{0}{1}", cc_major, cc_minor); - - constexpr size_t num_args = 9; - const char* args[num_args] = { - arch.c_str(), - cub_path, - thrust_path, - libcudacxx_path, - ctk_path, - "-rdc=true", - "-dlto", - "-DCUB_DISABLE_CDP", - "-std=c++20"}; - - constexpr size_t num_lto_args = 2; - const char* lopts[num_lto_args] = {"-lto", arch.c_str()}; - // Collect all LTO-IRs to be linked. nvrtc_linkable_list linkable_list; nvrtc_linkable_list_appender appender{linkable_list}; From bc0695fa458d00889a16eea00a23db074fa6980c Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 17 Sep 2025 23:35:27 +0000 Subject: [PATCH 056/100] Add missing util_device include --- cub/cub/agent/agent_three_way_partition.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/cub/cub/agent/agent_three_way_partition.cuh b/cub/cub/agent/agent_three_way_partition.cuh index dc7402cac27..1e428d84557 100644 --- a/cub/cub/agent/agent_three_way_partition.cuh +++ b/cub/cub/agent/agent_three_way_partition.cuh @@ -44,6 +44,7 @@ #include #include #include +#include #include #include From dfd3ec92538c9e38be0201d464f2436ae454266b Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 17 Sep 2025 23:36:18 +0000 Subject: [PATCH 057/100] Remove OffsetSize from kernel source --- cub/cub/device/dispatch/dispatch_three_way_partition.cuh | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index 3b85b5f6aaf..e21a89f3043 100644 --- a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -91,11 +91,6 @@ struct DeviceThreeWayPartitionKernelSource SelectSecondPartOp, per_partition_offset_t, streaming_context_t>); - - CUB_RUNTIME_FUNCTION static constexpr size_t OffsetSize() - { - return sizeof(OffsetT); - } }; } // namespace detail::three_way_partition @@ -191,7 +186,7 @@ struct DispatchThreeWayPartitionIf constexpr ::cuda::std::size_t num_counters_per_pass = 3; constexpr ::cuda::std::size_t num_streaming_counters = 2 * num_counters_per_pass; ::cuda::std::size_t streaming_selection_storage_bytes = - (num_partitions > 1) ? num_streaming_counters * kernel_source.OffsetSize() : ::cuda::std::size_t{0}; + (num_partitions > 1) ? num_streaming_counters * sizeof(OffsetT) : ::cuda::std::size_t{0}; // Specify temporary storage allocation requirements size_t allocation_sizes[2] = {0ULL, streaming_selection_storage_bytes}; From 9acc485b424be253764d94b1f4f38b305b8baf5e Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 17 Sep 2025 23:36:37 +0000 Subject: [PATCH 058/100] Add missing enable_if include --- cub/cub/detail/ptx-json/value.h | 1 + 1 file changed, 1 insertion(+) diff --git a/cub/cub/detail/ptx-json/value.h b/cub/cub/detail/ptx-json/value.h index 311bc26d08c..967df605802 100644 --- a/cub/cub/detail/ptx-json/value.h +++ b/cub/cub/detail/ptx-json/value.h @@ -29,6 +29,7 @@ #include +#include #include #include From cb66b0631caecd93b8cf37c3d8241c5f5b66c803 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 17 Sep 2025 23:37:09 +0000 Subject: [PATCH 059/100] Add ptx_json for delay constructors --- cub/cub/agent/single_pass_scan_operators.cuh | 30 ++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/cub/cub/agent/single_pass_scan_operators.cuh b/cub/cub/agent/single_pass_scan_operators.cuh index 230873b6470..5f29121310c 100644 --- a/cub/cub/agent/single_pass_scan_operators.cuh +++ b/cub/cub/agent/single_pass_scan_operators.cuh @@ -498,6 +498,36 @@ using default_reduce_by_key_delay_constructor_t = reduce_by_key_delay_constructor_t<350, 450>, default_delay_constructor_t>>; +#if defined(CUB_ENABLE_POLICY_PTX_JSON) +# include + +// ptx-json encoders for delay constructor types. Unlike the other agent policy +// member variables, this is defined as a type alias so we can't use the +// CUB_DETAIL_POLICY_WRAPPER_DEFINE macro to embed it with ptx-json. To work +// around this, we define the ptx-json encoders here. These can then be used in +// the policy wrapper's EncodedPolicy member function to explicitly encode the +// delay constructor. + +template +struct delay_constructor_json; + +template +struct delay_constructor_json> +{ + using type = + ptx_json::object() = ptx_json::value(), + ptx_json::key<"delay">() = ptx_json::value(), + ptx_json::key<"l2_write_latency">() = ptx_json::value()>; +}; + +template +struct delay_constructor_json> +{ + using type = ptx_json::object() = ptx_json::value(), + ptx_json::key<"l2_write_latency">() = ptx_json::value()>; +}; +#endif // CUB_ENABLE_POLICY_PTX_JSON + /** * @brief Alias template for a ScanTileState specialized for a given value type, `T`, and memory order `Order`. * From 89a612cce7f65ca67e9cc2faf950bcaa30549510 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 18 Sep 2025 22:57:42 +0000 Subject: [PATCH 060/100] Add check for offset iterator types --- c/parallel/src/segmented_sort.cu | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 20ed77af919..89bb826e2dc 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -532,6 +532,15 @@ CUresult cccl_device_segmented_sort_build( return CUDA_ERROR_UNKNOWN; } + if (cccl_iterator_kind_t::CCCL_POINTER != start_offset_it.type + || cccl_iterator_kind_t::CCCL_POINTER != end_offset_it.type) + { + fflush(stderr); + printf("\nERROR in cccl_device_segmented_sort_build(): start_offset_it and end_offset_it must be a pointer\n "); + fflush(stdout); + return CUDA_ERROR_UNKNOWN; + } + try { const char* name = "device_segmented_sort"; From 591858aed79d3b9957220ccdee3bd4f75d4852bb Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 19 Sep 2025 13:29:04 +0000 Subject: [PATCH 061/100] Remove unused code --- c/parallel/test/test_segmented_sort.cpp | 32 ------------------------- 1 file changed, 32 deletions(-) diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index a52193d363e..858767d060c 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -566,38 +566,6 @@ C2H_TEST("SegmentedSort works with custom types as values", "[segmented_sort][cu } #endif -struct variable_segment_offset_iterator_state_t -{ - SizeT linear_id; - const SizeT* offsets; -}; - -static std::tuple make_variable_segment_iterator_sources() -{ - static constexpr std::string_view it_state_src = R"XXX( -struct variable_segment_offset_iterator_state_t { - unsigned long long linear_id; - const unsigned long long* offsets; -}; -)XXX"; - - static constexpr std::string_view it_advance_src = R"XXX( -extern "C" __device__ void advance_variable_offset_it(variable_segment_offset_iterator_state_t* state, unsigned long long offset) -{ - state->linear_id += offset; -} -)XXX"; - - static constexpr std::string_view it_deref_src = R"XXX( -extern "C" __device__ unsigned long long dereference_variable_offset_it(variable_segment_offset_iterator_state_t* state) -{ - return state->offsets[state->linear_id]; -} -)XXX"; - - return std::make_tuple(std::string(it_state_src), std::string(it_advance_src), std::string(it_deref_src)); -} - struct SegmentedSort_VariableSegments_Fixture_Tag; C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][variable_segments]", test_params_tuple) { From a23aab02c532eec2d6b90246a1625a7315cafdc1 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 19 Sep 2025 17:21:53 +0000 Subject: [PATCH 062/100] Move sorting algorithms to a new directory and add segmented_sort bindings --- .../cccl/parallel/experimental/_bindings.pyi | 23 +++ .../parallel/experimental/_bindings_impl.pyx | 147 +++++++++++++++++- .../experimental/algorithms/__init__.py | 10 +- .../algorithms/{ => _sort}/_merge_sort.py | 20 +-- .../algorithms/{ => _sort}/_radix_sort.py | 67 ++------ .../algorithms/_sort/_sort_common.py | 52 +++++++ 6 files changed, 245 insertions(+), 74 deletions(-) rename python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/{ => _sort}/_merge_sort.py (94%) rename python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/{ => _sort}/_radix_sort.py (81%) create mode 100644 python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_sort_common.py diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings.pyi b/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings.pyi index f7b69463872..e89e7c54054 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings.pyi +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings.pyi @@ -403,3 +403,26 @@ class DeviceHistogramBuildResult: row_stride_samples: int, stream, ) -> None: ... + +# ----------------- +# DeviceSegmentedSort +# ----------------- + +class DeviceSegmentedSortBuildResult: + def __init__(self): ... + def compute( + self, + temp_storage_ptr: int | None, + temp_storage_nbytes: int, + d_in_keys: Iterator, + d_out_keys: Iterator, + d_in_values: Iterator, + d_out_values: Iterator, + num_items: int, + num_segments: int, + d_begin_offsets: Iterator, + d_end_offsets: Iterator, + is_overwrite_okay: bool, + selector: int, + stream, + ) -> tuple[int, int]: ... diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx b/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx index 4bb573d09b8..28c99af789f 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx @@ -1615,7 +1615,7 @@ cdef class DeviceRadixSortBuildResult: if status != 0: raise RuntimeError( - f"Failed executing ascending radix_sort, error code: {status}" + f"Failed executing radix_sort, error code: {status}" ) return storage_sz, selector_int @@ -1953,5 +1953,150 @@ cdef class DeviceHistogramBuildResult: return storage_sz + def _get_cubin(self): + return self.build_data.cubin[:self.build_data.cubin_size] + +# ------------------- +# DeviceSegmentedSort +# ------------------- + +cdef extern from "cccl/c/segmented_sort.h": + cdef struct cccl_device_segmented_sort_build_result_t 'cccl_device_segmented_sort_build_result_t': + const char* cubin + size_t cubin_size + + cdef CUresult cccl_device_segmented_sort_build( + cccl_device_segmented_sort_build_result_t *build_ptr, + cccl_sort_order_t sort_order, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_keys_out, + cccl_iterator_t d_values_in, + cccl_iterator_t d_values_out, + cccl_iterator_t begin_offset_in, + cccl_iterator_t end_offset_in, + int, + int, + const char *, + const char *, + const char *, + const char * + ) nogil + + cdef CUresult cccl_device_segmented_sort( + cccl_device_segmented_sort_build_result_t build, + void* d_temp_storage, + size_t* temp_storage_bytes, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_keys_out, + cccl_iterator_t d_values_in, + cccl_iterator_t d_values_out, + int64_t num_items, + int64_t num_segments, + cccl_iterator_t start_offset_in, + cccl_iterator_t end_offset_in, + bool is_overwrite_okay, + int* selector, + CUstream stream + ) nogil + + cdef CUresult cccl_device_segmented_sort_cleanup( + cccl_device_segmented_sort_build_result_t* build_ptr + ) nogil + +cdef class DeviceSegmentedSortBuildResult: + cdef cccl_device_segmented_sort_build_result_t build_data + + def __dealloc__(DeviceSegmentedSortBuildResult self): + cdef CUresult status = -1 + with nogil: + status = cccl_device_segmented_sort_cleanup(&self.build_data) + if (status != 0): + print(f"Return code {status} encountered during segmented_sort result cleanup") + + def __cinit__( + DeviceSegmentedSortBuildResult self, + cccl_sort_order_t order, + Iterator d_keys_in, + Iterator d_values_in, + Iterator begin_offset_in, + Iterator end_offset_in, + CommonData common_data, + ): + cdef CUresult status = -1 + cdef int cc_major = common_data.get_cc_major() + cdef int cc_minor = common_data.get_cc_minor() + cdef const char *cub_path = common_data.cub_path_get_c_str() + cdef const char *thrust_path = common_data.thrust_path_get_c_str() + cdef const char *libcudacxx_path = common_data.libcudacxx_path_get_c_str() + cdef const char *ctk_path = common_data.ctk_path_get_c_str() + + memset(&self.build_data, 0, sizeof(cccl_device_segmented_sort_build_result_t)) + with nogil: + status = cccl_device_segmented_sort_build( + &self.build_data, + order, + d_keys_in.iter_data, + d_values_in.iter_data, + begin_offset_in.iter_data, + end_offset_in.iter_data, + cc_major, + cc_minor, + cub_path, + thrust_path, + libcudacxx_path, + ctk_path, + ) + if status != 0: + raise RuntimeError( + f"Failed building segmented_sort, error code: {status}" + ) + + cpdef tuple compute( + DeviceSegmentedSortBuildResult self, + temp_storage_ptr, + temp_storage_bytes, + Iterator d_keys_in, + Iterator d_keys_out, + Iterator d_values_in, + Iterator d_values_out, + size_t num_items, + size_t num_segments, + Iterator start_offset_in, + Iterator end_offset_in, + bint is_overwrite_okay, + selector, + stream + ): + cdef CUresult status = -1 + cdef void *storage_ptr = (temp_storage_ptr) if temp_storage_ptr else NULL + cdef size_t storage_sz = temp_storage_bytes + cdef int selector_int = selector + cdef CUstream c_stream = (stream) if stream else NULL + + with nogil: + status = cccl_device_segmented_sort( + self.build_data, + storage_ptr, + &storage_sz, + d_keys_in.iter_data, + d_keys_out.iter_data, + d_values_in.iter_data, + d_values_out.iter_data, + num_items, + num_segments, + start_offset_in.iter_data, + end_offset_in.iter_data, + is_overwrite_okay, + &selector_int, + c_stream + ) + + if status != 0: + raise RuntimeError( + f"Failed executing segmented_sort, error code: {status}" + ) + return storage_sz, selector_int + + def _get_cubin(self): return self.build_data.cubin[:self.build_data.cubin_size] diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py index d288f3e5195..437db5b2bc0 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py @@ -5,11 +5,6 @@ from ._histogram import histogram_even as histogram_even from ._histogram import make_histogram_even as make_histogram_even -from ._merge_sort import make_merge_sort as make_merge_sort -from ._merge_sort import merge_sort as merge_sort -from ._radix_sort import DoubleBuffer, SortOrder -from ._radix_sort import make_radix_sort as make_radix_sort -from ._radix_sort import radix_sort as radix_sort from ._reduce import make_reduce_into as make_reduce_into from ._reduce import reduce_into as reduce_into from ._scan import exclusive_scan as exclusive_scan @@ -18,6 +13,11 @@ from ._scan import make_inclusive_scan as make_inclusive_scan from ._segmented_reduce import make_segmented_reduce as make_segmented_reduce from ._segmented_reduce import segmented_reduce +from ._sort._merge_sort import make_merge_sort as make_merge_sort +from ._sort._merge_sort import merge_sort as merge_sort +from ._sort._radix_sort import make_radix_sort as make_radix_sort +from ._sort._radix_sort import radix_sort as radix_sort +from ._sort._sort_common import DoubleBuffer, SortOrder from ._transform import binary_transform, unary_transform from ._transform import make_binary_transform as make_binary_transform from ._transform import make_unary_transform as make_unary_transform diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_merge_sort.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_merge_sort.py similarity index 94% rename from python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_merge_sort.py rename to python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_merge_sort.py index eca99b3b148..59f4b8a3156 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_merge_sort.py +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_merge_sort.py @@ -7,19 +7,19 @@ import numba -from .. import _bindings -from .. import _cccl_interop as cccl -from .._caching import CachableFunction, cache_with_key -from .._cccl_interop import call_build, set_cccl_iterator_state -from .._utils import protocols -from .._utils.protocols import ( +from ... import _bindings +from ... import _cccl_interop as cccl +from ..._caching import CachableFunction, cache_with_key +from ..._cccl_interop import call_build, set_cccl_iterator_state +from ..._utils import protocols +from ..._utils.protocols import ( get_data_pointer, validate_and_get_stream, ) -from .._utils.temp_storage_buffer import TempStorageBuffer -from ..iterators._iterators import IteratorBase -from ..op import OpKind -from ..typing import DeviceArrayLike +from ..._utils.temp_storage_buffer import TempStorageBuffer +from ...iterators._iterators import IteratorBase +from ...op import OpKind +from ...typing import DeviceArrayLike def make_cache_key( diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_radix_sort.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_radix_sort.py similarity index 81% rename from python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_radix_sort.py rename to python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_radix_sort.py index ef1dee953a5..8d6a0535133 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_radix_sort.py +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_radix_sort.py @@ -3,37 +3,18 @@ # # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -from enum import Enum -from typing import Tuple - -from .. import _bindings -from .. import _cccl_interop as cccl -from .._caching import cache_with_key -from .._cccl_interop import call_build, set_cccl_iterator_state -from .._utils.protocols import ( +from ... import _bindings +from ... import _cccl_interop as cccl +from ..._caching import cache_with_key +from ..._cccl_interop import call_build, set_cccl_iterator_state +from ..._utils.protocols import ( get_data_pointer, get_dtype, validate_and_get_stream, ) -from .._utils.temp_storage_buffer import TempStorageBuffer -from ..typing import DeviceArrayLike - - -class SortOrder(Enum): - ASCENDING = 0 - DESCENDING = 1 - - -class DoubleBuffer: - def __init__(self, d_current: DeviceArrayLike, d_alternate: DeviceArrayLike): - self.d_buffers = [d_current, d_alternate] - self.selector = 0 - - def current(self): - return self.d_buffers[self.selector] - - def alternate(self): - return self.d_buffers[1 - self.selector] +from ..._utils.temp_storage_buffer import TempStorageBuffer +from ...typing import DeviceArrayLike +from ._sort_common import DoubleBuffer, SortOrder, _get_arrays def make_cache_key( @@ -65,32 +46,6 @@ def make_cache_key( ) -def _get_arrays( - d_in_keys: DeviceArrayLike | DoubleBuffer, - d_out_keys: DeviceArrayLike | DoubleBuffer | None, - d_in_values: DeviceArrayLike | DoubleBuffer | None, - d_out_values: DeviceArrayLike | None, -) -> Tuple: - if isinstance(d_in_keys, DoubleBuffer): - d_in_keys_array = d_in_keys.current() - d_out_keys_array = d_in_keys.alternate() - - if d_in_values is not None: - assert isinstance(d_in_values, DoubleBuffer) - d_in_values_array = d_in_values.current() - d_out_values_array = d_in_values.alternate() - else: - d_in_values_array = None - d_out_values_array = None - else: - d_in_keys_array = d_in_keys - d_in_values_array = d_in_values - d_out_keys_array = d_out_keys - d_out_values_array = d_out_values - - return d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array - - class _RadixSort: __slots__ = [ "d_in_keys_cccl", @@ -104,7 +59,7 @@ class _RadixSort: def __init__( self, d_in_keys: DeviceArrayLike | DoubleBuffer, - d_out_keys: DeviceArrayLike | DoubleBuffer | None, + d_out_keys: DeviceArrayLike | None, d_in_values: DeviceArrayLike | DoubleBuffer | None, d_out_values: DeviceArrayLike | None, order: SortOrder, @@ -294,10 +249,6 @@ def radix_sort( end_bit, stream, ) - # Use the appropriate array for namespace - prefer d_out_keys, fallback to d_in_keys - ref_array = d_out_keys if d_out_keys is not None else d_in_keys - if isinstance(ref_array, DoubleBuffer): - ref_array = ref_array.current() tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream) sorter( tmp_storage, diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_sort_common.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_sort_common.py new file mode 100644 index 00000000000..8c2caa505eb --- /dev/null +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_sort_common.py @@ -0,0 +1,52 @@ +# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +from enum import Enum +from typing import Tuple + +from ...typing import DeviceArrayLike + + +class SortOrder(Enum): + ASCENDING = 0 + DESCENDING = 1 + + +class DoubleBuffer: + def __init__(self, d_current: DeviceArrayLike, d_alternate: DeviceArrayLike): + self.d_buffers = [d_current, d_alternate] + self.selector = 0 + + def current(self): + return self.d_buffers[self.selector] + + def alternate(self): + return self.d_buffers[1 - self.selector] + + +def _get_arrays( + d_in_keys: DeviceArrayLike | DoubleBuffer, + d_out_keys: DeviceArrayLike | None, + d_in_values: DeviceArrayLike | DoubleBuffer | None, + d_out_values: DeviceArrayLike | None, +) -> Tuple[DeviceArrayLike, DeviceArrayLike, DeviceArrayLike, DeviceArrayLike]: + if isinstance(d_in_keys, DoubleBuffer): + d_in_keys_array = d_in_keys.current() + d_out_keys_array = d_in_keys.alternate() + + if d_in_values is not None: + assert isinstance(d_in_values, DoubleBuffer) + d_in_values_array = d_in_values.current() + d_out_values_array = d_in_values.alternate() + else: + d_in_values_array = None + d_out_values_array = None + else: + d_in_keys_array = d_in_keys + d_in_values_array = d_in_values + d_out_keys_array = d_out_keys + d_out_values_array = d_out_values + + return d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array From 4406abccbfc983b331c6f1983ee67564ae2ccade Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 19 Sep 2025 17:27:04 +0000 Subject: [PATCH 063/100] Fix cython compilation errors --- .../cuda/cccl/parallel/experimental/_bindings_impl.pyx | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx b/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx index 28c99af789f..c4dbae4271a 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/_bindings_impl.pyx @@ -1970,8 +1970,6 @@ cdef extern from "cccl/c/segmented_sort.h": cccl_sort_order_t sort_order, cccl_iterator_t d_keys_in, cccl_iterator_t d_keys_out, - cccl_iterator_t d_values_in, - cccl_iterator_t d_values_out, cccl_iterator_t begin_offset_in, cccl_iterator_t end_offset_in, int, @@ -1994,7 +1992,7 @@ cdef extern from "cccl/c/segmented_sort.h": int64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, - bool is_overwrite_okay, + bint is_overwrite_okay, int* selector, CUstream stream ) nogil From 253c8042d3e8b1b9a7399e401728ec7dd4892a87 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 19 Sep 2025 17:27:46 +0000 Subject: [PATCH 064/100] Add initial python wrappers for segmented_sort --- .../algorithms/_sort/_segmented_sort.py | 272 ++++++++++++++++ .../examples/sort/segmented_sort_basic.py | 54 +++ .../examples/sort/segmented_sort_buffer.py | 60 ++++ .../examples/sort/segmented_sort_object.py | 76 +++++ .../tests/parallel/test_segmented_sort.py | 307 ++++++++++++++++++ 5 files changed, 769 insertions(+) create mode 100644 python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py create mode 100644 python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py create mode 100644 python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py create mode 100644 python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py create mode 100644 python/cuda_cccl/tests/parallel/test_segmented_sort.py diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py new file mode 100644 index 00000000000..7136c27b098 --- /dev/null +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py @@ -0,0 +1,272 @@ +# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + + +from ... import _bindings +from ... import _cccl_interop as cccl +from ..._caching import cache_with_key +from ..._cccl_interop import call_build, set_cccl_iterator_state +from ..._utils.protocols import ( + get_data_pointer, + get_dtype, + validate_and_get_stream, +) +from ..._utils.temp_storage_buffer import TempStorageBuffer +from ...typing import DeviceArrayLike +from ._sort_common import DoubleBuffer, SortOrder, _get_arrays + + +class _SegmentedSort: + __slots__ = [ + "build_result", + "d_in_keys_cccl", + "d_out_keys_cccl", + "d_in_values_cccl", + "d_out_values_cccl", + "start_offsets_in_cccl", + "end_offsets_in_cccl", + ] + + def __init__( + self, + d_in_keys: DeviceArrayLike | DoubleBuffer, + d_out_keys: DeviceArrayLike | None, + d_in_values: DeviceArrayLike | DoubleBuffer | None, + d_out_values: DeviceArrayLike | None, + start_offsets_in: DeviceArrayLike, + end_offsets_in: DeviceArrayLike, + order: SortOrder, + ): + self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys) + self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys) + self.d_in_values_cccl = cccl.to_cccl_input_iter(d_in_values) + self.d_out_values_cccl = cccl.to_cccl_output_iter(d_out_values) + self.start_offsets_in_cccl = cccl.to_cccl_input_iter(start_offsets_in) + self.end_offsets_in_cccl = cccl.to_cccl_input_iter(end_offsets_in) + + cccl.cccl_iterator_set_host_advance( + self.start_offsets_in_cccl, start_offsets_in + ) + cccl.cccl_iterator_set_host_advance(self.end_offsets_in_cccl, end_offsets_in) + + self.build_result = call_build( + _bindings.DeviceSegmentedSortBuildResult, + self.d_in_keys_cccl, + self.d_out_keys_cccl, + self.d_in_values_cccl, + self.d_out_values_cccl, + self.start_offsets_in_cccl, + self.end_offsets_in_cccl, + order, + ) + + def __call__( + self, + temp_storage, + d_in_keys, + d_out_keys, + d_in_values, + d_out_values, + num_items, + start_offsets_in, + end_offsets_in, + stream=None, + ): + set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys) + set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys) + set_cccl_iterator_state(self.d_in_values_cccl, d_in_values) + set_cccl_iterator_state(self.d_out_values_cccl, d_out_values) + set_cccl_iterator_state(self.start_offsets_in_cccl, start_offsets_in) + set_cccl_iterator_state(self.end_offsets_in_cccl, end_offsets_in) + + stream_handle = validate_and_get_stream(stream) + if temp_storage is None: + temp_storage_bytes = 0 + d_temp_storage = 0 + else: + temp_storage_bytes = temp_storage.nbytes + d_temp_storage = get_data_pointer(temp_storage) + + # Detect overwrite mode and selector, similar to radix sort + is_overwrite_okay = isinstance(d_in_keys, DoubleBuffer) + selector = -1 + + temp_storage_bytes, selector = self.build_result.compute( + d_temp_storage, + temp_storage_bytes, + self.d_in_keys_cccl, + self.d_out_keys_cccl, + self.d_in_values_cccl, + self.d_out_values_cccl, + num_items, + self.start_offsets_in_cccl, + self.end_offsets_in_cccl, + is_overwrite_okay, + selector, + stream_handle, + ) + + if is_overwrite_okay and temp_storage is not None: + assert selector in (0, 1) + assert isinstance(d_in_keys, DoubleBuffer) + d_in_keys.selector = selector + if d_in_values is not None: + assert isinstance(d_in_values, DoubleBuffer) + d_in_values.selector = selector + + return temp_storage_bytes + + +def make_cache_key( + d_in_keys: DeviceArrayLike | DoubleBuffer, + d_out_keys: DeviceArrayLike | None, + d_in_values: DeviceArrayLike | DoubleBuffer | None, + d_out_values: DeviceArrayLike | None, + start_offsets_in: DeviceArrayLike, + end_offsets_in: DeviceArrayLike, + order: SortOrder, +): + d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = ( + _get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values) + ) + + d_in_keys_key = get_dtype(d_in_keys_array) + d_out_keys_key = None if d_out_keys_array is None else get_dtype(d_out_keys_array) + d_in_values_key = ( + None if d_in_values_array is None else get_dtype(d_in_values_array) + ) + d_out_values_key = ( + None if d_out_values_array is None else get_dtype(d_out_values_array) + ) + start_offsets_in_key = get_dtype(start_offsets_in) + end_offsets_in_key = get_dtype(end_offsets_in) + + return ( + d_in_keys_key, + d_out_keys_key, + d_in_values_key, + d_out_values_key, + start_offsets_in_key, + end_offsets_in_key, + order, + ) + + +@cache_with_key(make_cache_key) +def make_segmented_sort( + d_in_keys: DeviceArrayLike | DoubleBuffer, + d_out_keys: DeviceArrayLike | None, + d_in_values: DeviceArrayLike | DoubleBuffer | None, + d_out_values: DeviceArrayLike | None, + start_offsets_in: DeviceArrayLike, + end_offsets_in: DeviceArrayLike, + order: SortOrder, +): + """ + Performs a device-wide segmented sort using the specified keys and values. + + Example: + Below, ``make_segmented_sort`` is used to create a segmented sort object that can be reused. + + .. literalinclude:: ../../python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py + :language: python + :start-after: # example-begin + + Args: + d_in_keys: Device array or DoubleBuffer containing the input keys to be sorted + d_out_keys: Device array to store the sorted keys + d_in_values: Optional Device array or DoubleBuffer containing the input values to be sorted + d_out_values: Device array to store the sorted values + start_offsets_in: Device array or iterator containing the sequence of beginning offsets + end_offsets_in: Device array or iterator containing the sequence of ending offsets + order: SortOrder specifying the order of the sort + + Returns: + A callable object that can be used to perform the segmented sort + """ + return _SegmentedSort( + d_in_keys, + d_out_keys, + d_in_values, + d_out_values, + start_offsets_in, + end_offsets_in, + order, + ) + + +def segmented_sort( + d_in_keys: DeviceArrayLike | DoubleBuffer, + d_out_keys: DeviceArrayLike | None, + d_in_values: DeviceArrayLike | DoubleBuffer | None, + d_out_values: DeviceArrayLike | None, + num_items: int, + start_offsets_in: DeviceArrayLike, + end_offsets_in: DeviceArrayLike, + order: SortOrder, + stream=None, +): + """ + Performs device-wide segmented sort. + + This function automatically handles temporary storage allocation and execution. + + Example: + Below, ``segmented_sort`` is used to perform a segmented sort. It also rearranges the values according to the keys' order. + + .. literalinclude:: ../../python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py + :language: python + :start-after: # example-begin + + + In the following example, ``segmented_sort`` is used to perform a segmented sort with a ``DoubleBuffer` for reduced temporary storage. + + .. literalinclude:: ../../python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py + :language: python + :start-after: # example-begin + + Args: + d_in_keys: Device array or DoubleBuffer containing the input keys to be sorted + d_out_keys: Device array to store the sorted keys (optional) + d_in_values: Device array or DoubleBuffer containing the input values to be sorted (optional) + d_out_values: Device array to store the sorted values (optional) + num_items: Number of items to sort + start_offsets_in: Device array or iterator containing the sequence of beginning offsets + end_offsets_in: Device array or iterator containing the sequence of ending offsets + order: Sort order (ascending or descending) + stream: CUDA stream for the operation (optional) + """ + sorter = make_segmented_sort( + d_in_keys, + d_out_keys, + d_in_values, + d_out_values, + start_offsets_in, + end_offsets_in, + order, + ) + tmp_storage_bytes = sorter( + None, + d_in_keys, + d_out_keys, + d_in_values, + d_out_values, + num_items, + start_offsets_in, + end_offsets_in, + stream, + ) + tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream) + sorter( + tmp_storage, + d_in_keys, + d_out_keys, + d_in_values, + d_out_values, + num_items, + start_offsets_in, + end_offsets_in, + stream, + ) diff --git a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py new file mode 100644 index 00000000000..4f7c9894ce6 --- /dev/null +++ b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py @@ -0,0 +1,54 @@ +# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +# example-begin +""" +Example showing how to use segmented_sort to sort keys and values within segments. +""" + +import cupy as cp +import numpy as np + +import cuda.cccl.parallel.experimental as parallel + +# Prepare input keys and values, and segment offsets. +h_in_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype="int32") +h_in_vals = np.array([90, 10, 50, 40, 20, 80, 70, 30, 60], dtype="int32") + +# 3 segments: [0,3), [3,5), [5,9) +start_offsets = np.array([0, 3, 5], dtype=np.int64) +end_offsets = np.array([3, 5, 9], dtype=np.int64) + +d_in_keys = cp.asarray(h_in_keys) +d_in_vals = cp.asarray(h_in_vals) +d_out_keys = cp.empty_like(d_in_keys) +d_out_vals = cp.empty_like(d_in_vals) + +# Perform the segmented sort (ascending within each segment). +parallel.segmented_sort( + d_in_keys, + d_out_keys, + d_in_vals, + d_out_vals, + d_in_keys.size, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + parallel.SortOrder.ASCENDING, +) + +# Verify the result. +h_out_keys = cp.asnumpy(d_out_keys) +h_out_vals = cp.asnumpy(d_out_vals) + +expected_pairs = [] +for s, e in zip(start_offsets, end_offsets): + seg_pairs = sorted(zip(h_in_keys[s:e], h_in_vals[s:e]), key=lambda kv: kv[0]) + expected_pairs.extend(seg_pairs) + +expected_keys = np.array([k for k, _ in expected_pairs], dtype=h_in_keys.dtype) +expected_vals = np.array([v for _, v in expected_pairs], dtype=h_in_vals.dtype) + +assert np.array_equal(h_out_keys, expected_keys) +assert np.array_equal(h_out_vals, expected_vals) +print(f"Segmented sort basic result - keys: {h_out_keys}, values: {h_out_vals}") diff --git a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py new file mode 100644 index 00000000000..ffdcc43516a --- /dev/null +++ b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +# example-begin +""" +Example showing how to use segmented_sort with DoubleBuffer for reduced temporary storage. +""" + +import cupy as cp +import numpy as np + +import cuda.cccl.parallel.experimental as parallel + +# Prepare input keys and values, and segment offsets. +h_in_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype="int32") +h_in_vals = np.array([90, 10, 50, 40, 20, 80, 70, 30, 60], dtype="int32") + +# 3 segments: [0,3), [3,5), [5,9) +start_offsets = np.array([0, 3, 5], dtype=np.int64) +end_offsets = np.array([3, 5, 9], dtype=np.int64) + +d_in_keys = cp.asarray(h_in_keys) +d_in_vals = cp.asarray(h_in_vals) +d_tmp_keys = cp.empty_like(d_in_keys) +d_tmp_vals = cp.empty_like(d_in_vals) + +# Create double buffers for keys and values. +keys_db = parallel.DoubleBuffer(d_in_keys, d_tmp_keys) +vals_db = parallel.DoubleBuffer(d_in_vals, d_tmp_vals) + +# Perform the segmented sort (descending within each segment). +parallel.segmented_sort( + keys_db, + None, + vals_db, + None, + d_in_keys.size, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + parallel.SortOrder.DESCENDING, +) + +# Verify the result. +h_out_keys = cp.asnumpy(keys_db.current()) +h_out_vals = cp.asnumpy(vals_db.current()) + +expected_pairs = [] +for s, e in zip(start_offsets, end_offsets): + seg_pairs = sorted( + zip(h_in_keys[s:e], h_in_vals[s:e]), key=lambda kv: kv[0], reverse=True + ) + expected_pairs.extend(seg_pairs) + +expected_keys = np.array([k for k, _ in expected_pairs], dtype=h_in_keys.dtype) +expected_vals = np.array([v for _, v in expected_pairs], dtype=h_in_vals.dtype) + +assert np.array_equal(h_out_keys, expected_keys) +assert np.array_equal(h_out_vals, expected_vals) +print(f"Segmented sort buffer result - keys: {h_out_keys}, values: {h_out_vals}") diff --git a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py new file mode 100644 index 00000000000..d87d98b851f --- /dev/null +++ b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py @@ -0,0 +1,76 @@ +# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +# example-begin +""" +Example showing how to use segmented_sort with the object API. +""" + +import cupy as cp +import numpy as np + +import cuda.cccl.parallel.experimental as parallel + +# Prepare the input and segment offsets. +dtype = np.int32 +h_input_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype=dtype) +h_input_vals = np.array([90, 10, 50, 40, 20, 80, 70, 30, 60], dtype=dtype) +start_offsets = np.array([0, 3, 5], dtype=np.int64) +end_offsets = np.array([3, 5, 9], dtype=np.int64) + +d_input_keys = cp.asarray(h_input_keys) +d_input_vals = cp.asarray(h_input_vals) +d_output_keys = cp.empty_like(d_input_keys) +d_output_vals = cp.empty_like(d_input_vals) + +# Create the segmented sort object. +sorter = parallel.make_segmented_sort( + d_input_keys, + d_output_keys, + d_input_vals, + d_output_vals, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + parallel.SortOrder.ASCENDING, +) + +# Get the temporary storage size. +temp_storage_size = sorter( + None, + d_input_keys, + d_output_keys, + d_input_vals, + d_output_vals, + len(h_input_keys), + cp.asarray(start_offsets), + cp.asarray(end_offsets), +) +d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8) + +# Perform the segmented sort. +sorter( + d_temp_storage, + d_input_keys, + d_output_keys, + d_input_vals, + d_output_vals, + len(h_input_keys), + cp.asarray(start_offsets), + cp.asarray(end_offsets), +) + +# Verify the result. +expected_pairs = [] +for s, e in zip(start_offsets, end_offsets): + seg_pairs = sorted(zip(h_input_keys[s:e], h_input_vals[s:e]), key=lambda kv: kv[0]) + expected_pairs.extend(seg_pairs) + +expected_keys = np.array([k for k, _ in expected_pairs], dtype=dtype) +expected_values = np.array([v for _, v in expected_pairs], dtype=dtype) + +actual_keys = d_output_keys.get() +actual_values = d_output_vals.get() +np.testing.assert_array_equal(actual_keys, expected_keys) +np.testing.assert_array_equal(actual_values, expected_values) +print("Segmented sort object example completed successfully") diff --git a/python/cuda_cccl/tests/parallel/test_segmented_sort.py b/python/cuda_cccl/tests/parallel/test_segmented_sort.py new file mode 100644 index 00000000000..c0950539ee9 --- /dev/null +++ b/python/cuda_cccl/tests/parallel/test_segmented_sort.py @@ -0,0 +1,307 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +from typing import Tuple + +import cupy as cp +import numba +import numpy as np +import pytest + +import cuda.cccl.parallel.experimental as parallel + +DTYPE_LIST = [ + np.uint8, + np.int16, + np.uint32, + np.int64, + np.float16, + np.float32, + np.float64, +] + + +def get_mark_by_size(num_segments: int, segment_size: int): + num_items = num_segments * segment_size + return pytest.mark.large if num_items >= (1 << 20) else tuple() + + +NUM_SEGMENTS_LIST = [0, 13, 1024, 2048] +SEGMENT_SIZE_LIST = [1, 12, 1024, 2048] + +DTYPE_SEGMENT_PARAMS = [ + pytest.param(dt, ns, ss, marks=get_mark_by_size(ns, ss)) + for dt in DTYPE_LIST + for ns in NUM_SEGMENTS_LIST + for ss in SEGMENT_SIZE_LIST +] + + +def random_array(size, dtype, max_value=None) -> np.typing.NDArray: + rng = np.random.default_rng() + if np.isdtype(dtype, "integral"): + if max_value is None: + max_value = np.iinfo(dtype).max + return rng.integers(max_value, size=size, dtype=dtype) + elif np.isdtype(dtype, "real floating"): + return np.random.uniform(low=-10.0, high=10.0, size=size).astype(dtype) + else: + raise ValueError(f"Unsupported dtype {dtype}") + + +def make_uniform_segments( + num_segments: int, segment_size: int +) -> Tuple[np.ndarray, np.ndarray]: + start_offsets = np.arange(num_segments, dtype=np.int64) * segment_size + end_offsets = start_offsets + segment_size + return start_offsets, end_offsets + + +def host_segmented_sort( + h_keys: np.ndarray, + h_vals: np.ndarray | None, + start_offsets: np.ndarray, + end_offsets: np.ndarray, + order: "parallel.SortOrder", +) -> Tuple[np.ndarray, np.ndarray | None]: + assert start_offsets.shape == end_offsets.shape + keys = h_keys.copy() + vals = None if h_vals is None else h_vals.copy() + + for s, e in zip(start_offsets, end_offsets): + if e <= s: + continue + if vals is None: + if order is parallel.SortOrder.DESCENDING: + # stable descending + signed_dtype = ( + np.dtype(keys.dtype.name.replace("uint", "int")) + if np.issubdtype(keys.dtype, np.unsignedinteger) + else keys.dtype + ) + idx = np.argsort(-keys[s:e].astype(signed_dtype), stable=True) + else: + idx = np.argsort(keys[s:e], stable=True) + keys[s:e] = keys[s:e][idx] + else: + # build pairs for stable sort + pairs = list(zip(keys[s:e], vals[s:e])) + if order is parallel.SortOrder.DESCENDING: + pairs.sort(key=lambda kv: kv[0], reverse=True) + else: + pairs.sort(key=lambda kv: kv[0]) + ks, vs = zip(*pairs) if pairs else ([], []) + keys[s:e] = np.array(ks, dtype=keys.dtype) + vals[s:e] = np.array(vs, dtype=vals.dtype) + + return keys, vals + + +@pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) +def test_segmented_sort_keys(dtype, num_segments, segment_size): + order = parallel.SortOrder.ASCENDING + num_items = num_segments * segment_size + + h_in_keys = random_array(num_items, dtype, max_value=50) + start_offsets, end_offsets = make_uniform_segments(num_segments, segment_size) + + d_in_keys = numba.cuda.to_device(h_in_keys) + d_out_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) + + parallel.segmented_sort( + d_in_keys, + d_out_keys, + None, + None, + num_items, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + order, + ) + + h_out_keys = d_out_keys.copy_to_host() + expected_keys, _ = host_segmented_sort( + h_in_keys, None, start_offsets, end_offsets, order + ) + + np.testing.assert_array_equal(h_out_keys, expected_keys) + + +@pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) +def test_segmented_sort_pairs(dtype, num_segments, segment_size): + order = parallel.SortOrder.DESCENDING + num_items = num_segments * segment_size + + h_in_keys = random_array( + num_items, dtype, max_value=50 if np.isdtype(dtype, "integral") else None + ) + h_in_vals = random_array(num_items, np.float32) + + start_offsets, end_offsets = make_uniform_segments(num_segments, segment_size) + + d_in_keys = numba.cuda.to_device(h_in_keys) + d_in_vals = numba.cuda.to_device(h_in_vals) + d_out_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) + d_out_vals = numba.cuda.to_device(np.empty_like(h_in_vals)) + + parallel.segmented_sort( + d_in_keys, + d_out_keys, + d_in_vals, + d_out_vals, + num_items, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + order, + ) + + h_out_keys = d_out_keys.copy_to_host() + h_out_vals = d_out_vals.copy_to_host() + + expected_keys, expected_vals = host_segmented_sort( + h_in_keys, h_in_vals, start_offsets, end_offsets, order + ) + + np.testing.assert_array_equal(h_out_keys, expected_keys) + np.testing.assert_array_equal(h_out_vals, expected_vals) + + +@pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) +def test_segmented_sort_keys_double_buffer(dtype, num_segments, segment_size): + order = parallel.SortOrder.ASCENDING + num_items = num_segments * segment_size + + h_in_keys = random_array(num_items, dtype, max_value=20) + start_offsets, end_offsets = make_uniform_segments(num_segments, segment_size) + + d_in_keys = numba.cuda.to_device(h_in_keys) + d_tmp_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) + keys_db = parallel.DoubleBuffer(d_in_keys, d_tmp_keys) + + parallel.segmented_sort( + keys_db, + None, + None, + None, + num_items, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + order, + ) + + h_out_keys = keys_db.current().copy_to_host() + expected_keys, _ = host_segmented_sort( + h_in_keys, None, start_offsets, end_offsets, order + ) + np.testing.assert_array_equal(h_out_keys, expected_keys) + + +@pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) +def test_segmented_sort_pairs_double_buffer(dtype, num_segments, segment_size): + order = parallel.SortOrder.DESCENDING + num_items = num_segments * segment_size + + h_in_keys = random_array( + num_items, dtype, max_value=50 if np.isdtype(dtype, "integral") else None + ) + h_in_vals = random_array(num_items, np.float32) + + start_offsets, end_offsets = make_uniform_segments(num_segments, segment_size) + + d_in_keys = numba.cuda.to_device(h_in_keys) + d_in_vals = numba.cuda.to_device(h_in_vals) + d_tmp_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) + d_tmp_vals = numba.cuda.to_device(np.empty_like(h_in_vals)) + + keys_db = parallel.DoubleBuffer(d_in_keys, d_tmp_keys) + vals_db = parallel.DoubleBuffer(d_in_vals, d_tmp_vals) + + parallel.segmented_sort( + keys_db, + None, + vals_db, + None, + num_items, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + order, + ) + + h_out_keys = keys_db.current().copy_to_host() + h_out_vals = vals_db.current().copy_to_host() + + expected_keys, expected_vals = host_segmented_sort( + h_in_keys, h_in_vals, start_offsets, end_offsets, order + ) + np.testing.assert_array_equal(h_out_keys, expected_keys) + np.testing.assert_array_equal(h_out_vals, expected_vals) + + +@pytest.mark.parametrize("num_segments", [20, 600]) +def test_segmented_sort_variable_segment_sizes(num_segments): + order = parallel.SortOrder.ASCENDING + base_pattern = [ + 1, + 5, + 10, + 20, + 30, + 50, + 100, + 3, + 25, + 600, + 7, + 18, + 300, + 4, + 35, + 9, + 14, + 700, + 28, + 11, + ] + segment_sizes = [] + while len(segment_sizes) < num_segments: + remaining = num_segments - len(segment_sizes) + copy_count = min(remaining, len(base_pattern)) + segment_sizes.extend(base_pattern[:copy_count]) + + start_offsets = np.zeros(num_segments, dtype=np.int64) + end_offsets = np.zeros(num_segments, dtype=np.int64) + current = 0 + for i, sz in enumerate(segment_sizes): + start_offsets[i] = current + current += sz + end_offsets[i] = current + num_items = current + + h_in_keys = random_array(num_items, np.int32, max_value=100) + h_in_vals = random_array(num_items, np.float32) + + d_in_keys = numba.cuda.to_device(h_in_keys) + d_in_vals = numba.cuda.to_device(h_in_vals) + d_out_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) + d_out_vals = numba.cuda.to_device(np.empty_like(h_in_vals)) + + parallel.segmented_sort( + d_in_keys, + d_out_keys, + d_in_vals, + d_out_vals, + num_items, + cp.asarray(start_offsets), + cp.asarray(end_offsets), + order, + ) + + h_out_keys = d_out_keys.copy_to_host() + h_out_vals = d_out_vals.copy_to_host() + expected_keys, expected_vals = host_segmented_sort( + h_in_keys, h_in_vals, start_offsets, end_offsets, order + ) + + np.testing.assert_array_equal(h_out_keys, expected_keys) + np.testing.assert_array_equal(h_out_vals, expected_vals) From dd783ba252f4d90b9c035d8af4a4999a1099f7bc Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 19 Sep 2025 17:31:36 +0000 Subject: [PATCH 065/100] Add missing imports --- python/cuda_cccl/cuda/cccl/parallel/experimental/__init__.py | 4 ++++ .../cuda/cccl/parallel/experimental/algorithms/__init__.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/__init__.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/__init__.py index 414932201ad..53c99b03529 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/__init__.py +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/__init__.py @@ -17,12 +17,14 @@ make_radix_sort, make_reduce_into, make_segmented_reduce, + make_segmented_sort, make_unary_transform, make_unique_by_key, merge_sort, radix_sort, reduce_into, segmented_reduce, + segmented_sort, unary_transform, unique_by_key, ) @@ -55,6 +57,7 @@ "make_radix_sort", "make_reduce_into", "make_segmented_reduce", + "make_segmented_sort", "make_unary_transform", "make_unique_by_key", "merge_sort", @@ -63,6 +66,7 @@ "reduce_into", "ReverseIterator", "segmented_reduce", + "segmented_sort", "SortOrder", "TransformIterator", "unary_transform", diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py index 437db5b2bc0..dd7a479f0e2 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/__init__.py @@ -17,6 +17,8 @@ from ._sort._merge_sort import merge_sort as merge_sort from ._sort._radix_sort import make_radix_sort as make_radix_sort from ._sort._radix_sort import radix_sort as radix_sort +from ._sort._segmented_sort import make_segmented_sort as make_segmented_sort +from ._sort._segmented_sort import segmented_sort as segmented_sort from ._sort._sort_common import DoubleBuffer, SortOrder from ._transform import binary_transform, unary_transform from ._transform import make_binary_transform as make_binary_transform @@ -45,6 +47,8 @@ "make_segmented_reduce", "unique_by_key", "make_unique_by_key", + "segmented_sort", + "make_segmented_sort", "DoubleBuffer", "SortOrder", ] From 97fbfc48e57bed583b140e1d3da8da512970198d Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 19 Sep 2025 19:22:50 +0000 Subject: [PATCH 066/100] Adjust segmented_sort build to not need the output arrays --- c/parallel/include/cccl/c/segmented_sort.h | 22 +++- c/parallel/src/segmented_sort.cu | 133 ++++++++++++--------- c/parallel/test/test_segmented_sort.cpp | 6 +- 3 files changed, 99 insertions(+), 62 deletions(-) diff --git a/c/parallel/include/cccl/c/segmented_sort.h b/c/parallel/include/cccl/c/segmented_sort.h index 4ba99db820a..530b71bc170 100644 --- a/c/parallel/include/cccl/c/segmented_sort.h +++ b/c/parallel/include/cccl/c/segmented_sort.h @@ -48,9 +48,7 @@ CCCL_C_API CUresult cccl_device_segmented_sort_build( cccl_device_segmented_sort_build_result_t* build, cccl_sort_order_t sort_order, cccl_iterator_t d_keys_in, - cccl_iterator_t d_keys_out, cccl_iterator_t d_values_in, - cccl_iterator_t d_values_out, cccl_iterator_t begin_offset_in, cccl_iterator_t end_offset_in, int cc_major, @@ -60,6 +58,22 @@ CCCL_C_API CUresult cccl_device_segmented_sort_build( const char* libcudacxx_path, const char* ctk_path); +// Extended version with build configuration (mirrors radix_sort build_ex) +CCCL_C_API CUresult cccl_device_segmented_sort_build_ex( + cccl_device_segmented_sort_build_result_t* build, + cccl_sort_order_t sort_order, + cccl_iterator_t d_keys_in, + cccl_iterator_t d_values_in, + cccl_iterator_t begin_offset_in, + cccl_iterator_t end_offset_in, + int cc_major, + int cc_minor, + const char* cub_path, + const char* thrust_path, + const char* libcudacxx_path, + const char* ctk_path, + cccl_build_config* config); + CCCL_C_API CUresult cccl_device_segmented_sort( cccl_device_segmented_sort_build_result_t build, void* d_temp_storage, @@ -68,8 +82,8 @@ CCCL_C_API CUresult cccl_device_segmented_sort( cccl_iterator_t d_keys_out, cccl_iterator_t d_values_in, cccl_iterator_t d_values_out, - int64_t num_items, - int64_t num_segments, + uint64_t num_items, + uint64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, bool is_overwrite_okay, diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 89bb826e2dc..a70062e3470 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -35,6 +35,7 @@ #include #include #include +#include struct device_segmented_sort_policy; struct device_three_way_partition_policy; @@ -498,13 +499,11 @@ struct segmented_sort_end_offset_iterator_tag; struct segmented_sort_large_selector_tag; struct segmented_sort_small_selector_tag; -CUresult cccl_device_segmented_sort_build( +CUresult cccl_device_segmented_sort_build_ex( cccl_device_segmented_sort_build_result_t* build_ptr, cccl_sort_order_t sort_order, cccl_iterator_t keys_in_it, - cccl_iterator_t keys_out_it, cccl_iterator_t values_in_it, - cccl_iterator_t values_out_it, cccl_iterator_t start_offset_it, cccl_iterator_t end_offset_it, int cc_major, @@ -512,22 +511,15 @@ CUresult cccl_device_segmented_sort_build( const char* cub_path, const char* thrust_path, const char* libcudacxx_path, - const char* ctk_path) + const char* ctk_path, + cccl_build_config* config) { CUresult error = CUDA_SUCCESS; - if (keys_in_it.value_type.type != keys_out_it.value_type.type) + if (cccl_iterator_kind_t::CCCL_POINTER != keys_in_it.type || cccl_iterator_kind_t::CCCL_POINTER != values_in_it.type) { fflush(stderr); - printf("\nERROR in cccl_device_segmented_sort_build(): keys_in_it and keys_out_it must have the same type\n "); - fflush(stdout); - return CUDA_ERROR_UNKNOWN; - } - - if (values_in_it.value_type.type != values_out_it.value_type.type) - { - fflush(stderr); - printf("\nERROR in cccl_device_segmented_sort_build(): values_in_it and values_out_it must have the same type\n "); + printf("\nERROR in cccl_device_segmented_sort_build(): keys_in_it and values_in_it must be a pointer\n "); fflush(stdout); return CUDA_ERROR_UNKNOWN; } @@ -550,14 +542,9 @@ CUresult cccl_device_segmented_sort_build( const auto [keys_in_iterator_name, keys_in_iterator_src] = get_specialization(template_id(), keys_in_it); - const auto [keys_out_iterator_name, keys_out_iterator_src] = - get_specialization( - template_id(), keys_out_it, keys_out_it.value_type); - const bool keys_only = values_in_it.type == cccl_iterator_kind_t::CCCL_POINTER && values_in_it.state == nullptr; std::string values_in_iterator_name, values_in_iterator_src; - std::string values_out_iterator_name, values_out_iterator_src; if (!keys_only) { @@ -565,18 +552,11 @@ CUresult cccl_device_segmented_sort_build( template_id(), values_in_it); values_in_iterator_name = vi_name; values_in_iterator_src = vi_src; - - const auto [vo_name, vo_src] = get_specialization( - template_id(), values_out_it, values_in_it.value_type); - values_out_iterator_name = vo_name; - values_out_iterator_src = vo_src; } else { - values_in_iterator_name = "cub::NullType*"; - values_out_iterator_name = "cub::NullType*"; - values_in_iterator_src = ""; - values_out_iterator_src = ""; + values_in_iterator_name = "cub::NullType*"; + values_in_iterator_src = ""; } const auto [start_offset_iterator_name, start_offset_iterator_src] = @@ -594,8 +574,7 @@ CUresult cccl_device_segmented_sort_build( const std::string arch = std::format("-arch=sm_{0}{1}", cc_major, cc_minor); - constexpr size_t num_args = 9; - const char* args[num_args] = { + std::vector args = { arch.c_str(), cub_path, thrust_path, @@ -606,13 +585,31 @@ CUresult cccl_device_segmented_sort_build( "-DCUB_DISABLE_CDP", "-std=c++20"}; + cccl::detail::extend_args_with_build_config(args, config); + constexpr size_t num_lto_args = 2; const char* lopts[num_lto_args] = {"-lto", arch.c_str()}; cccl_op_t large_selector_op = segmented_sort::make_segments_selector_op( - 0, start_offset_it, end_offset_it, "cccl_large_segments_selector_op", ">", args, num_args, lopts, num_lto_args); + 0, + start_offset_it, + end_offset_it, + "cccl_large_segments_selector_op", + ">", + args.data(), + args.size(), + lopts, + num_lto_args); cccl_op_t small_selector_op = segmented_sort::make_segments_selector_op( - 0, start_offset_it, end_offset_it, "cccl_small_segments_selector_op", "<", args, num_args, lopts, num_lto_args); + 0, + start_offset_it, + end_offset_it, + "cccl_small_segments_selector_op", + "<", + args.data(), + args.size(), + lopts, + num_lto_args); cccl_type_info selector_result_t{sizeof(bool), alignof(bool), cccl_type_enum::CCCL_BOOLEAN}; cccl_type_info selector_input_t{ @@ -640,21 +637,17 @@ struct __align__({3}) items_storage_t {{ {7} {8} {9} -{10} -{11} )XXX", keys_in_it.value_type.size, // 0 keys_in_it.value_type.alignment, // 1 values_in_it.value_type.size, // 2 values_in_it.value_type.alignment, // 3 keys_in_iterator_src, // 4 - keys_out_iterator_src, // 5 - values_in_iterator_src, // 6 - values_out_iterator_src, // 7 - start_offset_iterator_src, // 8 - end_offset_iterator_src, // 9 - large_selector_src, // 10 - small_selector_src); // 11 + values_in_iterator_src, // 5 + start_offset_iterator_src, // 6 + end_offset_iterator_src, // 7 + large_selector_src, // 8 + small_selector_src); // 9 const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); @@ -773,11 +766,9 @@ struct device_three_way_partition_policy {{ // add iterator definitions appender.add_iterator_definition(keys_in_it); - appender.add_iterator_definition(keys_out_it); if (!keys_only) { appender.add_iterator_definition(values_in_it); - appender.add_iterator_definition(values_out_it); } appender.add_iterator_definition(start_offset_it); appender.add_iterator_definition(end_offset_it); @@ -793,7 +784,7 @@ struct device_three_way_partition_policy {{ ->add_expression({segmented_sort_kernel_large_name}) ->add_expression({three_way_partition_init_kernel_name}) ->add_expression({three_way_partition_kernel_name}) - ->compile_program({args, num_args}) + ->compile_program({args.data(), args.size()}) ->get_name({segmented_sort_fallback_kernel_name, segmented_sort_fallback_kernel_lowered_name}) ->get_name({segmented_sort_kernel_small_name, segmented_sort_kernel_small_lowered_name}) ->get_name({segmented_sort_kernel_large_name, segmented_sort_kernel_large_lowered_name}) @@ -843,6 +834,36 @@ struct device_three_way_partition_policy {{ return error; } +CUresult cccl_device_segmented_sort_build( + cccl_device_segmented_sort_build_result_t* build_ptr, + cccl_sort_order_t sort_order, + cccl_iterator_t keys_in_it, + cccl_iterator_t values_in_it, + cccl_iterator_t start_offset_it, + cccl_iterator_t end_offset_it, + int cc_major, + int cc_minor, + const char* cub_path, + const char* thrust_path, + const char* libcudacxx_path, + const char* ctk_path) +{ + return cccl_device_segmented_sort_build_ex( + build_ptr, + sort_order, + keys_in_it, + values_in_it, + start_offset_it, + end_offset_it, + cc_major, + cc_minor, + cub_path, + thrust_path, + libcudacxx_path, + ctk_path, + nullptr); +} + template CUresult cccl_device_segmented_sort_impl( cccl_device_segmented_sort_build_result_t build, @@ -852,14 +873,22 @@ CUresult cccl_device_segmented_sort_impl( cccl_iterator_t d_keys_out, cccl_iterator_t d_values_in, cccl_iterator_t d_values_out, - int64_t num_items, - int64_t num_segments, + uint64_t num_items, + uint64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, bool is_overwrite_okay, int* selector, CUstream stream) { + if (selector == nullptr) + { + fflush(stderr); + printf("\nERROR in cccl_device_segmented_sort(): selector cannot be nullptr\n"); + fflush(stdout); + return CUDA_ERROR_UNKNOWN; + } + bool pushed = false; CUresult error = CUDA_SUCCESS; try @@ -911,12 +940,8 @@ CUresult cccl_device_segmented_sort_impl( /* partition_policy */ *reinterpret_cast(build.partition_runtime_policy)); - if (selector != nullptr) - { - *selector = d_keys_double_buffer.selector; - } - - error = static_cast(exec_status); + *selector = d_keys_double_buffer.selector; + error = static_cast(exec_status); } catch (const std::exception& exc) { @@ -943,8 +968,8 @@ CUresult cccl_device_segmented_sort( cccl_iterator_t d_keys_out, cccl_iterator_t d_values_in, cccl_iterator_t d_values_out, - int64_t num_items, - int64_t num_segments, + uint64_t num_items, + uint64_t num_segments, cccl_iterator_t start_offset_in, cccl_iterator_t end_offset_in, bool is_overwrite_okay, diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 858767d060c..a7ceb9e0dfe 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -78,9 +78,9 @@ struct segmented_sort_build BuildResultT* build_ptr, cccl_sort_order_t sort_order, cccl_iterator_t keys_in, - cccl_iterator_t keys_out, + cccl_iterator_t /*keys_out*/, cccl_iterator_t values_in, - cccl_iterator_t values_out, + cccl_iterator_t /*values_out*/, int64_t /*num_items*/, int64_t /*num_segments*/, cccl_iterator_t start_offsets, @@ -98,9 +98,7 @@ struct segmented_sort_build build_ptr, sort_order, keys_in, - keys_out, values_in, - values_out, start_offsets, end_offsets, cc_major, From 896452902d483d832c9344428b89c20af8d624b2 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sat, 20 Sep 2025 03:16:26 +0000 Subject: [PATCH 067/100] Separate num_segments and num_items properly --- .../algorithms/_sort/_segmented_sort.py | 40 +++++++++++++------ .../examples/sort/segmented_sort_basic.py | 1 + .../examples/sort/segmented_sort_buffer.py | 1 + .../examples/sort/segmented_sort_object.py | 1 + .../tests/parallel/test_segmented_sort.py | 7 +++- 5 files changed, 37 insertions(+), 13 deletions(-) diff --git a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py index 7136c27b098..2df9bcc3db7 100644 --- a/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py +++ b/python/cuda_cccl/cuda/cccl/parallel/experimental/algorithms/_sort/_segmented_sort.py @@ -39,10 +39,14 @@ def __init__( end_offsets_in: DeviceArrayLike, order: SortOrder, ): - self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys) - self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys) - self.d_in_values_cccl = cccl.to_cccl_input_iter(d_in_values) - self.d_out_values_cccl = cccl.to_cccl_output_iter(d_out_values) + d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = ( + _get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values) + ) + + self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys_array) + self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys_array) + self.d_in_values_cccl = cccl.to_cccl_input_iter(d_in_values_array) + self.d_out_values_cccl = cccl.to_cccl_output_iter(d_out_values_array) self.start_offsets_in_cccl = cccl.to_cccl_input_iter(start_offsets_in) self.end_offsets_in_cccl = cccl.to_cccl_input_iter(end_offsets_in) @@ -53,13 +57,13 @@ def __init__( self.build_result = call_build( _bindings.DeviceSegmentedSortBuildResult, + _bindings.SortOrder.ASCENDING + if order is SortOrder.ASCENDING + else _bindings.SortOrder.DESCENDING, self.d_in_keys_cccl, - self.d_out_keys_cccl, self.d_in_values_cccl, - self.d_out_values_cccl, self.start_offsets_in_cccl, self.end_offsets_in_cccl, - order, ) def __call__( @@ -70,14 +74,21 @@ def __call__( d_in_values, d_out_values, num_items, + num_segments, start_offsets_in, end_offsets_in, stream=None, ): - set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys) - set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys) - set_cccl_iterator_state(self.d_in_values_cccl, d_in_values) - set_cccl_iterator_state(self.d_out_values_cccl, d_out_values) + d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = ( + _get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values) + ) + + set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys_array) + set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys_array) + if d_in_values_array is not None: + set_cccl_iterator_state(self.d_in_values_cccl, d_in_values_array) + if d_out_values_array is not None: + set_cccl_iterator_state(self.d_out_values_cccl, d_out_values_array) set_cccl_iterator_state(self.start_offsets_in_cccl, start_offsets_in) set_cccl_iterator_state(self.end_offsets_in_cccl, end_offsets_in) @@ -101,6 +112,7 @@ def __call__( self.d_in_values_cccl, self.d_out_values_cccl, num_items, + num_segments, self.start_offsets_in_cccl, self.end_offsets_in_cccl, is_overwrite_okay, @@ -203,6 +215,7 @@ def segmented_sort( d_in_values: DeviceArrayLike | DoubleBuffer | None, d_out_values: DeviceArrayLike | None, num_items: int, + num_segments: int, start_offsets_in: DeviceArrayLike, end_offsets_in: DeviceArrayLike, order: SortOrder, @@ -232,7 +245,8 @@ def segmented_sort( d_out_keys: Device array to store the sorted keys (optional) d_in_values: Device array or DoubleBuffer containing the input values to be sorted (optional) d_out_values: Device array to store the sorted values (optional) - num_items: Number of items to sort + num_items: Total number of items to sort + num_segments: Number of segments to sort start_offsets_in: Device array or iterator containing the sequence of beginning offsets end_offsets_in: Device array or iterator containing the sequence of ending offsets order: Sort order (ascending or descending) @@ -254,6 +268,7 @@ def segmented_sort( d_in_values, d_out_values, num_items, + num_segments, start_offsets_in, end_offsets_in, stream, @@ -266,6 +281,7 @@ def segmented_sort( d_in_values, d_out_values, num_items, + num_segments, start_offsets_in, end_offsets_in, stream, diff --git a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py index 4f7c9894ce6..2d37a064a20 100644 --- a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py +++ b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py @@ -32,6 +32,7 @@ d_in_vals, d_out_vals, d_in_keys.size, + start_offsets.size, cp.asarray(start_offsets), cp.asarray(end_offsets), parallel.SortOrder.ASCENDING, diff --git a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py index ffdcc43516a..12eb7e3b096 100644 --- a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py +++ b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py @@ -36,6 +36,7 @@ vals_db, None, d_in_keys.size, + start_offsets.size, cp.asarray(start_offsets), cp.asarray(end_offsets), parallel.SortOrder.DESCENDING, diff --git a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py index d87d98b851f..76f71b42fb6 100644 --- a/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py +++ b/python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py @@ -43,6 +43,7 @@ d_input_vals, d_output_vals, len(h_input_keys), + len(start_offsets), cp.asarray(start_offsets), cp.asarray(end_offsets), ) diff --git a/python/cuda_cccl/tests/parallel/test_segmented_sort.py b/python/cuda_cccl/tests/parallel/test_segmented_sort.py index c0950539ee9..e8699756954 100644 --- a/python/cuda_cccl/tests/parallel/test_segmented_sort.py +++ b/python/cuda_cccl/tests/parallel/test_segmented_sort.py @@ -63,7 +63,7 @@ def host_segmented_sort( h_vals: np.ndarray | None, start_offsets: np.ndarray, end_offsets: np.ndarray, - order: "parallel.SortOrder", + order: parallel.SortOrder, ) -> Tuple[np.ndarray, np.ndarray | None]: assert start_offsets.shape == end_offsets.shape keys = h_keys.copy() @@ -115,6 +115,7 @@ def test_segmented_sort_keys(dtype, num_segments, segment_size): None, None, num_items, + num_segments, cp.asarray(start_offsets), cp.asarray(end_offsets), order, @@ -151,6 +152,7 @@ def test_segmented_sort_pairs(dtype, num_segments, segment_size): d_in_vals, d_out_vals, num_items, + num_segments, cp.asarray(start_offsets), cp.asarray(end_offsets), order, @@ -185,6 +187,7 @@ def test_segmented_sort_keys_double_buffer(dtype, num_segments, segment_size): None, None, num_items, + num_segments, cp.asarray(start_offsets), cp.asarray(end_offsets), order, @@ -223,6 +226,7 @@ def test_segmented_sort_pairs_double_buffer(dtype, num_segments, segment_size): vals_db, None, num_items, + num_segments, cp.asarray(start_offsets), cp.asarray(end_offsets), order, @@ -292,6 +296,7 @@ def test_segmented_sort_variable_segment_sizes(num_segments): d_in_vals, d_out_vals, num_items, + num_segments, cp.asarray(start_offsets), cp.asarray(end_offsets), order, From c376097d63081478ff2536c6e026a0679f5837da Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sat, 20 Sep 2025 03:18:11 +0000 Subject: [PATCH 068/100] Add fp16 include during policy wrapper creation --- c/parallel/src/jit_templates/template_pre.h.in | 5 +++++ c/parallel/src/segmented_sort.cu | 5 +++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/c/parallel/src/jit_templates/template_pre.h.in b/c/parallel/src/jit_templates/template_pre.h.in index 97450d7dc1f..26cddcb7f38 100644 --- a/c/parallel/src/jit_templates/template_pre.h.in +++ b/c/parallel/src/jit_templates/template_pre.h.in @@ -1,2 +1,7 @@ const char * jit_template_header_contents = R"unlikelystring( #include +/// include this to use _CCCL_HAS_NVFP16() +#include +#if _CCCL_HAS_NVFP16() +#include +#endif diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index a70062e3470..bae48a7dfb9 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -651,8 +651,9 @@ struct __align__({3}) items_storage_t {{ const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); - constexpr size_t ptx_num_args = 5; - const char* ptx_args[ptx_num_args] = {ptx_arch.c_str(), cub_path, thrust_path, libcudacxx_path, "-rdc=true"}; + constexpr size_t ptx_num_args = 6; + const char* ptx_args[ptx_num_args] = { + ptx_arch.c_str(), cub_path, thrust_path, libcudacxx_path, ctk_path, "-rdc=true"}; static constexpr std::string_view policy_wrapper_expr_tmpl = R"XXXX(cub::detail::segmented_sort::MakeSegmentedSortPolicyWrapper(cub::detail::segmented_sort::policy_hub<{0}, {1}>::MaxPolicy::ActivePolicy{{}}))XXXX"; From c4ae3e27379abc8e4de86ac8536385e754b5e396 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sun, 21 Sep 2025 23:37:45 +0000 Subject: [PATCH 069/100] Don't use num_segments == 0 in tests since this causes an issue where the build assumes keys_only since the array size is 0 --- python/cuda_cccl/tests/parallel/test_segmented_sort.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cuda_cccl/tests/parallel/test_segmented_sort.py b/python/cuda_cccl/tests/parallel/test_segmented_sort.py index e8699756954..1dda052d921 100644 --- a/python/cuda_cccl/tests/parallel/test_segmented_sort.py +++ b/python/cuda_cccl/tests/parallel/test_segmented_sort.py @@ -27,7 +27,7 @@ def get_mark_by_size(num_segments: int, segment_size: int): return pytest.mark.large if num_items >= (1 << 20) else tuple() -NUM_SEGMENTS_LIST = [0, 13, 1024, 2048] +NUM_SEGMENTS_LIST = [1, 13, 1024, 2048] SEGMENT_SIZE_LIST = [1, 12, 1024, 2048] DTYPE_SEGMENT_PARAMS = [ From bcfd6b9a8b4d20c17fe05c28a48b296152c232e2 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 22 Sep 2025 02:50:31 +0000 Subject: [PATCH 070/100] Add typename to avoid benchmark compilation error --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 21b96fa0c1b..85ccd9662cb 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -730,8 +730,8 @@ private: size_t three_way_partition_temp_storage_bytes, cub::detail::device_double_buffer& d_keys_double_buffer, cub::detail::device_double_buffer& d_values_double_buffer, - KernelSource::LargeSegmentsSelectorT& large_segments_selector, - KernelSource::SmallSegmentsSelectorT& small_segments_selector, + typename KernelSource::LargeSegmentsSelectorT& large_segments_selector, + typename KernelSource::SmallSegmentsSelectorT& small_segments_selector, cub::detail::temporary_storage::alias& device_partition_temp_storage, cub::detail::temporary_storage::alias& large_and_medium_segments_indices, cub::detail::temporary_storage::alias& small_segments_indices, From baa01a3c74957ca1240dd2a51b7c4be856eb68bf Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 22 Sep 2025 03:04:16 +0000 Subject: [PATCH 071/100] Skip SASS check when the dtype is int64 since it exists in C++ as well --- .../cuda_cccl/tests/parallel/test_segmented_sort.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/cuda_cccl/tests/parallel/test_segmented_sort.py b/python/cuda_cccl/tests/parallel/test_segmented_sort.py index 1dda052d921..dab125cc63b 100644 --- a/python/cuda_cccl/tests/parallel/test_segmented_sort.py +++ b/python/cuda_cccl/tests/parallel/test_segmented_sort.py @@ -99,7 +99,16 @@ def host_segmented_sort( @pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) -def test_segmented_sort_keys(dtype, num_segments, segment_size): +def test_segmented_sort_keys(dtype, num_segments, segment_size, monkeypatch): + # Disable SASS verification only for this test when dtype is int64 + if np.dtype(dtype) == np.dtype(np.int64): + import cuda.cccl.parallel.experimental._cccl_interop + + monkeypatch.setattr( + cuda.cccl.parallel.experimental._cccl_interop, + "_check_sass", + False, + ) order = parallel.SortOrder.ASCENDING num_items = num_segments * segment_size From 8c3e808c4fa0db85499b38ee35d36964c484365e Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 22 Sep 2025 03:19:53 +0000 Subject: [PATCH 072/100] Rename parameter list for clarity --- .../cuda_cccl/tests/parallel/test_segmented_sort.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/python/cuda_cccl/tests/parallel/test_segmented_sort.py b/python/cuda_cccl/tests/parallel/test_segmented_sort.py index dab125cc63b..1fdd1efe62a 100644 --- a/python/cuda_cccl/tests/parallel/test_segmented_sort.py +++ b/python/cuda_cccl/tests/parallel/test_segmented_sort.py @@ -27,14 +27,19 @@ def get_mark_by_size(num_segments: int, segment_size: int): return pytest.mark.large if num_items >= (1 << 20) else tuple() -NUM_SEGMENTS_LIST = [1, 13, 1024, 2048] -SEGMENT_SIZE_LIST = [1, 12, 1024, 2048] +NUM_SEGMENTS_AND_SEGMENT_SIZES = [ + (1, 1), + (1, 2048), + (13, 12), + (1024, 1024), + (2048, 1), + (2048, 13), +] DTYPE_SEGMENT_PARAMS = [ pytest.param(dt, ns, ss, marks=get_mark_by_size(ns, ss)) for dt in DTYPE_LIST - for ns in NUM_SEGMENTS_LIST - for ss in SEGMENT_SIZE_LIST + for (ns, ss) in NUM_SEGMENTS_AND_SEGMENT_SIZES ] From 79f3f4a9c1ab523d91f506b8bef2beedd4d7c66f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20=27Griwes=27=20Dominiak?= Date: Sun, 21 Sep 2025 20:22:32 -0700 Subject: [PATCH 073/100] c.parallel: enable dynamic policies in scan. Co-authored-by: Nader Al Awar --- c/parallel/include/cccl/c/scan.h | 1 + c/parallel/src/scan.cu | 139 +++++++++--------- c/parallel/src/util/runtime_policy.cpp | 3 - c/parallel/test/test_scan.cpp | 20 ++- cub/cub/agent/agent_scan.cuh | 22 +++ cub/cub/agent/single_pass_scan_operators.cuh | 81 ++++++++++ cub/cub/detail/ptx-json/value.h | 4 +- .../device/dispatch/tuning/tuning_scan.cuh | 10 ++ cub/cub/util_device.cuh | 37 +++-- 9 files changed, 220 insertions(+), 97 deletions(-) diff --git a/c/parallel/include/cccl/c/scan.h b/c/parallel/include/cccl/c/scan.h index 16febf55e03..e438753829b 100644 --- a/c/parallel/include/cccl/c/scan.h +++ b/c/parallel/include/cccl/c/scan.h @@ -35,6 +35,7 @@ typedef struct cccl_device_scan_build_result_t bool force_inclusive; size_t description_bytes_per_tile; size_t payload_bytes_per_tile; + void* runtime_policy; } cccl_device_scan_build_result_t; CCCL_C_API CUresult cccl_device_scan_build( diff --git a/c/parallel/src/scan.cu b/c/parallel/src/scan.cu index 807e239b8e3..f81e2a1685d 100644 --- a/c/parallel/src/scan.cu +++ b/c/parallel/src/scan.cu @@ -35,6 +35,7 @@ #include #include #include +#include #include #include @@ -51,61 +52,30 @@ namespace scan struct scan_runtime_tuning_policy { - int block_size; - int items_per_thread; - cub::CacheLoadModifier load_modifier; + cub::detail::RuntimeScanAgentPolicy scan; - scan_runtime_tuning_policy Scan() const + auto Scan() const { - return *this; - } - - int ItemsPerThread() const - { - return items_per_thread; - } - - int BlockThreads() const - { - return block_size; - } - - cub::CacheLoadModifier LoadModifier() const - { - return load_modifier; + return scan; } void CheckLoadModifier() const { - if (LoadModifier() == cub::CacheLoadModifier::LOAD_LDG) + if (scan.LoadModifier() == cub::CacheLoadModifier::LOAD_LDG) { throw std::runtime_error("The memory consistency model does not apply to texture " "accesses"); } } -}; -template -Tuning find_tuning(int cc, const Tuning (&tunings)[N]) -{ - for (const Tuning& tuning : tunings) + using MaxPolicy = scan_runtime_tuning_policy; + + template + cudaError_t Invoke(int, F& op) { - if (cc >= tuning.cc) - { - return tuning; - } + return op.template Invoke(*this); } - - return tunings[N - 1]; -} - -scan_runtime_tuning_policy get_policy(int /*cc*/, cccl_type_info /*accumulator_type*/) -{ - // TODO: we should update this once we figure out a way to reuse - // tuning logic from C++. Alternately, we should implement - // something better than a hardcoded default: - return {128, 4, cub::LOAD_DEFAULT}; -} +}; static cccl_type_info get_accumulator_type(cccl_op_t /*op*/, cccl_iterator_t /*input_it*/, cccl_value_t init) { @@ -235,7 +205,6 @@ CUresult cccl_device_scan_build_ex( const int cc = cc_major * 10 + cc_minor; const cccl_type_info accum_t = scan::get_accumulator_type(op, input_it, init); - const auto policy = scan::get_policy(cc, accum_t); const auto accum_cpp = cccl_type_enum_to_name(accum_t.type); const auto input_it_value_t = cccl_type_enum_to_name(input_it.value_type.type); const auto offset_t = cccl_type_enum_to_name(cccl_type_enum::CCCL_UINT64); @@ -254,41 +223,70 @@ CUresult cccl_device_scan_build_ex( struct __align__({1}) storage_t {{ char data[{0}]; }}; +{2} +{3} {4} -{5} -struct agent_policy_t {{ - static constexpr int ITEMS_PER_THREAD = {2}; - static constexpr int BLOCK_THREADS = {3}; - static constexpr cub::BlockLoadAlgorithm LOAD_ALGORITHM = cub::BLOCK_LOAD_WARP_TRANSPOSE; - static constexpr cub::CacheLoadModifier LOAD_MODIFIER = cub::LOAD_DEFAULT; - static constexpr cub::BlockStoreAlgorithm STORE_ALGORITHM = cub::BLOCK_STORE_WARP_TRANSPOSE; - static constexpr cub::BlockScanAlgorithm SCAN_ALGORITHM = cub::BLOCK_SCAN_WARP_SCANS; - struct detail {{ - using delay_constructor_t = cub::detail::default_delay_constructor_t<{7}>; - }}; -}}; -struct device_scan_policy {{ - struct ActivePolicy {{ - using ScanPolicyT = agent_policy_t; - }}; -}}; -{6} )XXX"; const std::string& src = std::format( src_template, input_it.value_type.size, // 0 input_it.value_type.alignment, // 1 - policy.items_per_thread, // 2 - policy.block_size, // 3 - input_iterator_src, // 4 - output_iterator_src, // 5 - op_src, // 6 - accum_cpp); // 7 + input_iterator_src, // 2 + output_iterator_src, // 3 + op_src); // 4 + + const auto output_it_value_t = cccl_type_enum_to_name(output_it.value_type.type); + + const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); + + std::vector ptx_args = { + ptx_arch.c_str(), cub_path, thrust_path, libcudacxx_path, ctk_path, "-rdc=true"}; + + cccl::detail::extend_args_with_build_config(ptx_args, config); + + std::string policy_hub_expr = std::format( + "cub::detail::scan::policy_hub<{}, {}, {}, {}, {}>", + input_it_value_t, + output_it_value_t, + accum_cpp, + offset_t, + "op_wrapper"); + + nlohmann::json runtime_policy = get_policy( + std::format("cub::detail::scan::MakeScanPolicyWrapper({}::MaxPolicy::ActivePolicy{{}})", policy_hub_expr), + "#include \n" + src, + ptx_args); + + auto delay_ctor_info = runtime_policy["DelayConstructor"]; + std::string delay_ctor_params; + for (auto&& param : delay_ctor_info["params"]) + { + delay_ctor_params.append(to_string(param) + ", "); + } + delay_ctor_params.erase(delay_ctor_params.size() - 2); // remove last ", " + auto delay_ctor_t = + std::format("cub::detail::{}<{}>", delay_ctor_info["name"].get(), delay_ctor_params); + + using cub::detail::RuntimeScanAgentPolicy; + auto [scan_policy, + scan_policy_str] = RuntimeScanAgentPolicy::from_json(runtime_policy, "ScanPolicyT", delay_ctor_t); + + std::string final_src = std::format( + R"XXX( +{0} +struct device_scan_policy {{ + struct ActivePolicy {{ + {1} + }}; +}}; +)XXX", + src, + scan_policy_str); #if false // CCCL_DEBUGGING_SWITCH fflush(stderr); - printf("\nCODE4NVRTC BEGIN\n%sCODE4NVRTC END\n", src.c_str()); + printf("\nCODE4NVRTC BEGIN\n%sCODE4NVRTC END\n", final_src.c_str()); fflush(stdout); #endif @@ -317,7 +315,7 @@ struct device_scan_policy {{ nvrtc_link_result result = begin_linking_nvrtc_program(num_lto_args, lopts) - ->add_program(nvrtc_translation_unit{src.c_str(), name}) + ->add_program(nvrtc_translation_unit{final_src.c_str(), name}) ->add_expression({init_kernel_name}) ->add_expression({scan_kernel_name}) ->compile_program({args.data(), args.size()}) @@ -341,6 +339,7 @@ struct device_scan_policy {{ build_ptr->force_inclusive = force_inclusive; build_ptr->description_bytes_per_tile = description_bytes_per_tile; build_ptr->payload_bytes_per_tile = payload_bytes_per_tile; + build_ptr->runtime_policy = new scan::scan_runtime_tuning_policy{scan_policy}; } catch (const std::exception& exc) { @@ -382,7 +381,7 @@ CUresult cccl_device_scan( ::cuda::std::size_t, void, EnforceInclusive, - scan::dynamic_scan_policy_t<&scan::get_policy>, + scan::scan_runtime_tuning_policy, scan::scan_kernel_source, cub::detail::CudaDriverLauncherFactory>:: Dispatch( @@ -396,7 +395,7 @@ CUresult cccl_device_scan( stream, {build}, cub::detail::CudaDriverLauncherFactory{cu_device, build.cc}, - {scan::get_accumulator_type(op, d_in, init)}); + *reinterpret_cast(build.runtime_policy)); error = static_cast(exec_status); } diff --git a/c/parallel/src/util/runtime_policy.cpp b/c/parallel/src/util/runtime_policy.cpp index f9ae69b9e53..a1a4719b4e5 100644 --- a/c/parallel/src/util/runtime_policy.cpp +++ b/c/parallel/src/util/runtime_policy.cpp @@ -27,9 +27,6 @@ get_policy(std::string_view policy_wrapper_expr, std::string_view translation_un std::string_view tag_name = "c_parallel_get_policy_tag"; std::string fixed_source = std::format( "{0}\n" - "#if _CCCL_HAS_NVFP16()\n" - "#include \n" - "#endif\n" "__global__ void ptx_json_emitting_kernel()\n" "{{\n" " [[maybe_unused]] auto wrapped = {1};\n" diff --git a/c/parallel/test/test_scan.cpp b/c/parallel/test/test_scan.cpp index f8d62dba7b9..1f7486c3dc1 100644 --- a/c/parallel/test/test_scan.cpp +++ b/c/parallel/test/test_scan.cpp @@ -39,7 +39,7 @@ auto& get_cache() return fixture::get_or_create().get_value(); } -template +template struct scan_build { CUresult operator()( @@ -75,7 +75,7 @@ struct scan_build static bool should_check_sass(int cc_major) { // TODO: add a check for NVRTC version; ref nvbug 5243118 - return (!Disable75SassCheck || cc_major > 7) && cc_major < 9; + return !(Disable75SassCheck && DisableForOtherArches) && (!Disable75SassCheck || cc_major > 7) && cc_major < 9; } }; @@ -96,7 +96,10 @@ struct scan_run } }; -template +template void scan(cccl_iterator_t input, cccl_iterator_t output, uint64_t num_items, @@ -106,8 +109,12 @@ void scan(cccl_iterator_t input, std::optional& cache, const std::optional& lookup_key) { - AlgorithmExecute, scan_cleanup, scan_run, BuildCache, KeyT>( - cache, lookup_key, inclusive, input, output, num_items, op, init); + AlgorithmExecute, + scan_cleanup, + scan_run, + BuildCache, + KeyT>(cache, lookup_key, inclusive, input, output, num_items, op, init); } // ============== @@ -486,7 +493,8 @@ C2H_TEST("Scan works with floating point types", "[scan]", floating_point_types) auto& build_cache = get_cache(); const auto& test_key = make_key(); - scan(input_ptr, output_ptr, num_items, op, init, false, build_cache, test_key); + // FIXME: figure out why scan spills to lmem for double + scan, true>(input_ptr, output_ptr, num_items, op, init, false, build_cache, test_key); const std::vector output = output_ptr; std::vector expected(num_items); diff --git a/cub/cub/agent/agent_scan.cuh b/cub/cub/agent/agent_scan.cuh index f75cd5c1ff6..d082f64166c 100644 --- a/cub/cub/agent/agent_scan.cuh +++ b/cub/cub/agent/agent_scan.cuh @@ -49,6 +49,7 @@ #include #include #include +#include #include #include @@ -111,6 +112,27 @@ struct AgentScanPolicy : ScalingType }; }; +#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) +namespace detail +{ +// Only define this when needed. +// Because of overload woes, this depends on C++20 concepts. util_device.h checks that concepts are available when +// either runtime policies or PTX JSON information are enabled, so if they are, this is always valid. The generic +// version is always defined, and that's the only one needed for regular CUB operations. +// +// TODO: enable this unconditionally once concepts are always available +CUB_DETAIL_POLICY_WRAPPER_DEFINE( + ScanAgentPolicy, + (GenericAgentPolicy), + (BLOCK_THREADS, BlockThreads, int), + (ITEMS_PER_THREAD, ItemsPerThread, int), + (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), + (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), + (STORE_ALGORITHM, StoreAlgorithm, cub::BlockStoreAlgorithm), + (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm)) +} // namespace detail +#endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) + /****************************************************************************** * Thread block abstractions ******************************************************************************/ diff --git a/cub/cub/agent/single_pass_scan_operators.cuh b/cub/cub/agent/single_pass_scan_operators.cuh index 230873b6470..4a1b3606084 100644 --- a/cub/cub/agent/single_pass_scan_operators.cuh +++ b/cub/cub/agent/single_pass_scan_operators.cuh @@ -48,6 +48,7 @@ #include #include #include +#include #include #include @@ -220,6 +221,14 @@ struct no_delay_constructor_t { return {}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), key<"params">() = array()>(); + } +#endif }; template @@ -248,6 +257,15 @@ struct reduce_by_key_delay_constructor_t { return {}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; template @@ -270,6 +288,15 @@ struct fixed_delay_constructor_t { return {}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; template @@ -295,6 +322,15 @@ struct exponential_backoff_constructor_t { return {InitialDelay}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; template @@ -333,6 +369,15 @@ struct exponential_backoff_jitter_constructor_t { return {InitialDelay, seed}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; template @@ -371,6 +416,15 @@ struct exponential_backoff_jitter_window_constructor_t { return {InitialDelay, seed}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; template @@ -412,6 +466,15 @@ struct exponential_backon_jitter_window_constructor_t max_delay >>= 1; return {max_delay, seed}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; template @@ -452,6 +515,15 @@ struct exponential_backon_jitter_constructor_t max_delay >>= 1; return {max_delay, seed}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; template @@ -480,6 +552,15 @@ struct exponential_backon_constructor_t max_delay >>= 1; return {max_delay}; } + +#if defined(CUB_ENABLE_POLICY_PTX_JSON) + _CCCL_DEVICE static constexpr auto EncodedConstructor() + { + using namespace ptx_json; + return object() = value(), + key<"params">() = array()>(); + } +#endif }; using default_no_delay_constructor_t = no_delay_constructor_t<450>; diff --git a/cub/cub/detail/ptx-json/value.h b/cub/cub/detail/ptx-json/value.h index 6aefd1c351a..7108316d968 100644 --- a/cub/cub/detail/ptx-json/value.h +++ b/cub/cub/detail/ptx-json/value.h @@ -67,11 +67,11 @@ struct value { __forceinline__ __device__ static void emit() { - value::emit(); + Nested.emit(); } }; -template +template struct value { __forceinline__ __device__ static void emit() diff --git a/cub/cub/device/dispatch/tuning/tuning_scan.cuh b/cub/cub/device/dispatch/tuning/tuning_scan.cuh index 10ec520fff8..97f5edee527 100644 --- a/cub/cub/device/dispatch/tuning/tuning_scan.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_scan.cuh @@ -469,6 +469,16 @@ struct ScanPolicyWrapper() = Scan().EncodedPolicy(), + key<"DelayConstructor">() = + StaticPolicyT::ScanPolicyT::detail::delay_constructor_t::EncodedConstructor()>(); + } +#endif }; template diff --git a/cub/cub/util_device.cuh b/cub/cub/util_device.cuh index c80ab244324..b3ac7cfe794 100644 --- a/cub/cub/util_device.cuh +++ b/cub/cub/util_device.cuh @@ -579,22 +579,27 @@ namespace detail # define CUB_DETAIL_POLICY_WRAPPER_FIELD_VALUE(field) , (int) ap._CCCL_PP_CAT(runtime_, _CCCL_PP_FIRST field) -# define CUB_DETAIL_POLICY_WRAPPER_AGENT_POLICY(concept_name, ...) \ - struct Runtime##concept_name \ - { \ - _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_FIELD, __VA_ARGS__) \ - static std::pair \ - from_json(const nlohmann::json& json, std::string_view subpolicy_name) \ - { \ - auto subpolicy = json[subpolicy_name]; \ - assert(!subpolicy.is_null()); \ - Runtime##concept_name ap; \ - _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_GET_FIELD, __VA_ARGS__) \ - return std::make_pair( \ - ap, \ - std::format("struct {} {{\n" _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_FIELD_STRING, __VA_ARGS__) "}};\n", \ - subpolicy_name _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_FIELD_VALUE, __VA_ARGS__))); \ - } \ +# define CUB_DETAIL_POLICY_WRAPPER_AGENT_POLICY(concept_name, ...) \ + struct Runtime##concept_name \ + { \ + _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_FIELD, __VA_ARGS__) \ + static std::pair \ + from_json(const nlohmann::json& json, \ + std::string_view subpolicy_name, \ + std::optional delay_cons_type = std::nullopt) \ + { \ + auto subpolicy = json[subpolicy_name]; \ + assert(!subpolicy.is_null()); \ + Runtime##concept_name ap; \ + _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_GET_FIELD, __VA_ARGS__) \ + return std::make_pair( \ + ap, \ + std::format( \ + "struct {} {{\n" _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_FIELD_STRING, __VA_ARGS__) "{} }};\n", \ + subpolicy_name _CCCL_PP_FOR_EACH(CUB_DETAIL_POLICY_WRAPPER_FIELD_VALUE, __VA_ARGS__), \ + delay_cons_type ? std::format("struct detail {{ using delay_constructor_t = {}; }}; ", *delay_cons_type) \ + : "")); \ + } \ }; #else # define CUB_DETAIL_POLICY_WRAPPER_AGENT_POLICY(...) From 418b1c8af15d8c8ca9718c6cb59d165d0f618f66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20=27Griwes=27=20Dominiak?= Date: Sun, 21 Sep 2025 22:41:26 -0700 Subject: [PATCH 074/100] Add a missing include. --- cub/cub/detail/ptx-json/value.h | 1 + 1 file changed, 1 insertion(+) diff --git a/cub/cub/detail/ptx-json/value.h b/cub/cub/detail/ptx-json/value.h index 7108316d968..13c74195cda 100644 --- a/cub/cub/detail/ptx-json/value.h +++ b/cub/cub/detail/ptx-json/value.h @@ -29,6 +29,7 @@ #include +#include #include #include From 2ec8219afecac33009e2f471c00ee22e34691e88 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Thu, 25 Sep 2025 12:52:56 +0200 Subject: [PATCH 075/100] Update cub/cub/device/dispatch/dispatch_three_way_partition.cuh --- cub/cub/device/dispatch/dispatch_three_way_partition.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index e21a89f3043..ec31151fbe0 100644 --- a/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -239,7 +239,7 @@ struct DispatchThreeWayPartitionIf } // Log three_way_partition_init_kernel configuration - const int init_grid_size = _CUDA_VSTD::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); + const int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); #ifdef CUB_DEBUG_LOG _CubLog("Invoking three_way_partition_init_kernel<<<%d, %d, 0, %lld>>>()\n", From 509a0b03c37a2879c50794e065d5c2bf9b026a69 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 25 Sep 2025 15:05:27 +0000 Subject: [PATCH 076/100] Undo some CUB changes from #5960 (they will be added in another PR) --- c/parallel/include/cccl/c/scan.h | 1 - c/parallel/src/scan.cu | 139 +++++++++--------- c/parallel/src/util/runtime_policy.cpp | 3 + c/parallel/test/test_scan.cpp | 20 +-- cub/cub/agent/agent_scan.cuh | 23 +-- .../device/dispatch/tuning/tuning_scan.cuh | 10 -- 6 files changed, 80 insertions(+), 116 deletions(-) diff --git a/c/parallel/include/cccl/c/scan.h b/c/parallel/include/cccl/c/scan.h index e438753829b..16febf55e03 100644 --- a/c/parallel/include/cccl/c/scan.h +++ b/c/parallel/include/cccl/c/scan.h @@ -35,7 +35,6 @@ typedef struct cccl_device_scan_build_result_t bool force_inclusive; size_t description_bytes_per_tile; size_t payload_bytes_per_tile; - void* runtime_policy; } cccl_device_scan_build_result_t; CCCL_C_API CUresult cccl_device_scan_build( diff --git a/c/parallel/src/scan.cu b/c/parallel/src/scan.cu index f81e2a1685d..807e239b8e3 100644 --- a/c/parallel/src/scan.cu +++ b/c/parallel/src/scan.cu @@ -35,7 +35,6 @@ #include #include #include -#include #include #include @@ -52,30 +51,61 @@ namespace scan struct scan_runtime_tuning_policy { - cub::detail::RuntimeScanAgentPolicy scan; + int block_size; + int items_per_thread; + cub::CacheLoadModifier load_modifier; - auto Scan() const + scan_runtime_tuning_policy Scan() const { - return scan; + return *this; + } + + int ItemsPerThread() const + { + return items_per_thread; + } + + int BlockThreads() const + { + return block_size; + } + + cub::CacheLoadModifier LoadModifier() const + { + return load_modifier; } void CheckLoadModifier() const { - if (scan.LoadModifier() == cub::CacheLoadModifier::LOAD_LDG) + if (LoadModifier() == cub::CacheLoadModifier::LOAD_LDG) { throw std::runtime_error("The memory consistency model does not apply to texture " "accesses"); } } +}; - using MaxPolicy = scan_runtime_tuning_policy; - - template - cudaError_t Invoke(int, F& op) +template +Tuning find_tuning(int cc, const Tuning (&tunings)[N]) +{ + for (const Tuning& tuning : tunings) { - return op.template Invoke(*this); + if (cc >= tuning.cc) + { + return tuning; + } } -}; + + return tunings[N - 1]; +} + +scan_runtime_tuning_policy get_policy(int /*cc*/, cccl_type_info /*accumulator_type*/) +{ + // TODO: we should update this once we figure out a way to reuse + // tuning logic from C++. Alternately, we should implement + // something better than a hardcoded default: + return {128, 4, cub::LOAD_DEFAULT}; +} static cccl_type_info get_accumulator_type(cccl_op_t /*op*/, cccl_iterator_t /*input_it*/, cccl_value_t init) { @@ -205,6 +235,7 @@ CUresult cccl_device_scan_build_ex( const int cc = cc_major * 10 + cc_minor; const cccl_type_info accum_t = scan::get_accumulator_type(op, input_it, init); + const auto policy = scan::get_policy(cc, accum_t); const auto accum_cpp = cccl_type_enum_to_name(accum_t.type); const auto input_it_value_t = cccl_type_enum_to_name(input_it.value_type.type); const auto offset_t = cccl_type_enum_to_name(cccl_type_enum::CCCL_UINT64); @@ -223,70 +254,41 @@ CUresult cccl_device_scan_build_ex( struct __align__({1}) storage_t {{ char data[{0}]; }}; -{2} -{3} {4} +{5} +struct agent_policy_t {{ + static constexpr int ITEMS_PER_THREAD = {2}; + static constexpr int BLOCK_THREADS = {3}; + static constexpr cub::BlockLoadAlgorithm LOAD_ALGORITHM = cub::BLOCK_LOAD_WARP_TRANSPOSE; + static constexpr cub::CacheLoadModifier LOAD_MODIFIER = cub::LOAD_DEFAULT; + static constexpr cub::BlockStoreAlgorithm STORE_ALGORITHM = cub::BLOCK_STORE_WARP_TRANSPOSE; + static constexpr cub::BlockScanAlgorithm SCAN_ALGORITHM = cub::BLOCK_SCAN_WARP_SCANS; + struct detail {{ + using delay_constructor_t = cub::detail::default_delay_constructor_t<{7}>; + }}; +}}; +struct device_scan_policy {{ + struct ActivePolicy {{ + using ScanPolicyT = agent_policy_t; + }}; +}}; +{6} )XXX"; const std::string& src = std::format( src_template, input_it.value_type.size, // 0 input_it.value_type.alignment, // 1 - input_iterator_src, // 2 - output_iterator_src, // 3 - op_src); // 4 - - const auto output_it_value_t = cccl_type_enum_to_name(output_it.value_type.type); - - const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); - - std::vector ptx_args = { - ptx_arch.c_str(), cub_path, thrust_path, libcudacxx_path, ctk_path, "-rdc=true"}; - - cccl::detail::extend_args_with_build_config(ptx_args, config); - - std::string policy_hub_expr = std::format( - "cub::detail::scan::policy_hub<{}, {}, {}, {}, {}>", - input_it_value_t, - output_it_value_t, - accum_cpp, - offset_t, - "op_wrapper"); - - nlohmann::json runtime_policy = get_policy( - std::format("cub::detail::scan::MakeScanPolicyWrapper({}::MaxPolicy::ActivePolicy{{}})", policy_hub_expr), - "#include \n" + src, - ptx_args); - - auto delay_ctor_info = runtime_policy["DelayConstructor"]; - std::string delay_ctor_params; - for (auto&& param : delay_ctor_info["params"]) - { - delay_ctor_params.append(to_string(param) + ", "); - } - delay_ctor_params.erase(delay_ctor_params.size() - 2); // remove last ", " - auto delay_ctor_t = - std::format("cub::detail::{}<{}>", delay_ctor_info["name"].get(), delay_ctor_params); - - using cub::detail::RuntimeScanAgentPolicy; - auto [scan_policy, - scan_policy_str] = RuntimeScanAgentPolicy::from_json(runtime_policy, "ScanPolicyT", delay_ctor_t); - - std::string final_src = std::format( - R"XXX( -{0} -struct device_scan_policy {{ - struct ActivePolicy {{ - {1} - }}; -}}; -)XXX", - src, - scan_policy_str); + policy.items_per_thread, // 2 + policy.block_size, // 3 + input_iterator_src, // 4 + output_iterator_src, // 5 + op_src, // 6 + accum_cpp); // 7 #if false // CCCL_DEBUGGING_SWITCH fflush(stderr); - printf("\nCODE4NVRTC BEGIN\n%sCODE4NVRTC END\n", final_src.c_str()); + printf("\nCODE4NVRTC BEGIN\n%sCODE4NVRTC END\n", src.c_str()); fflush(stdout); #endif @@ -315,7 +317,7 @@ struct device_scan_policy {{ nvrtc_link_result result = begin_linking_nvrtc_program(num_lto_args, lopts) - ->add_program(nvrtc_translation_unit{final_src.c_str(), name}) + ->add_program(nvrtc_translation_unit{src.c_str(), name}) ->add_expression({init_kernel_name}) ->add_expression({scan_kernel_name}) ->compile_program({args.data(), args.size()}) @@ -339,7 +341,6 @@ struct device_scan_policy {{ build_ptr->force_inclusive = force_inclusive; build_ptr->description_bytes_per_tile = description_bytes_per_tile; build_ptr->payload_bytes_per_tile = payload_bytes_per_tile; - build_ptr->runtime_policy = new scan::scan_runtime_tuning_policy{scan_policy}; } catch (const std::exception& exc) { @@ -381,7 +382,7 @@ CUresult cccl_device_scan( ::cuda::std::size_t, void, EnforceInclusive, - scan::scan_runtime_tuning_policy, + scan::dynamic_scan_policy_t<&scan::get_policy>, scan::scan_kernel_source, cub::detail::CudaDriverLauncherFactory>:: Dispatch( @@ -395,7 +396,7 @@ CUresult cccl_device_scan( stream, {build}, cub::detail::CudaDriverLauncherFactory{cu_device, build.cc}, - *reinterpret_cast(build.runtime_policy)); + {scan::get_accumulator_type(op, d_in, init)}); error = static_cast(exec_status); } diff --git a/c/parallel/src/util/runtime_policy.cpp b/c/parallel/src/util/runtime_policy.cpp index a1a4719b4e5..f9ae69b9e53 100644 --- a/c/parallel/src/util/runtime_policy.cpp +++ b/c/parallel/src/util/runtime_policy.cpp @@ -27,6 +27,9 @@ get_policy(std::string_view policy_wrapper_expr, std::string_view translation_un std::string_view tag_name = "c_parallel_get_policy_tag"; std::string fixed_source = std::format( "{0}\n" + "#if _CCCL_HAS_NVFP16()\n" + "#include \n" + "#endif\n" "__global__ void ptx_json_emitting_kernel()\n" "{{\n" " [[maybe_unused]] auto wrapped = {1};\n" diff --git a/c/parallel/test/test_scan.cpp b/c/parallel/test/test_scan.cpp index 1f7486c3dc1..f8d62dba7b9 100644 --- a/c/parallel/test/test_scan.cpp +++ b/c/parallel/test/test_scan.cpp @@ -39,7 +39,7 @@ auto& get_cache() return fixture::get_or_create().get_value(); } -template +template struct scan_build { CUresult operator()( @@ -75,7 +75,7 @@ struct scan_build static bool should_check_sass(int cc_major) { // TODO: add a check for NVRTC version; ref nvbug 5243118 - return !(Disable75SassCheck && DisableForOtherArches) && (!Disable75SassCheck || cc_major > 7) && cc_major < 9; + return (!Disable75SassCheck || cc_major > 7) && cc_major < 9; } }; @@ -96,10 +96,7 @@ struct scan_run } }; -template +template void scan(cccl_iterator_t input, cccl_iterator_t output, uint64_t num_items, @@ -109,12 +106,8 @@ void scan(cccl_iterator_t input, std::optional& cache, const std::optional& lookup_key) { - AlgorithmExecute, - scan_cleanup, - scan_run, - BuildCache, - KeyT>(cache, lookup_key, inclusive, input, output, num_items, op, init); + AlgorithmExecute, scan_cleanup, scan_run, BuildCache, KeyT>( + cache, lookup_key, inclusive, input, output, num_items, op, init); } // ============== @@ -493,8 +486,7 @@ C2H_TEST("Scan works with floating point types", "[scan]", floating_point_types) auto& build_cache = get_cache(); const auto& test_key = make_key(); - // FIXME: figure out why scan spills to lmem for double - scan, true>(input_ptr, output_ptr, num_items, op, init, false, build_cache, test_key); + scan(input_ptr, output_ptr, num_items, op, init, false, build_cache, test_key); const std::vector output = output_ptr; std::vector expected(num_items); diff --git a/cub/cub/agent/agent_scan.cuh b/cub/cub/agent/agent_scan.cuh index e6d4fe4f445..f75cd5c1ff6 100644 --- a/cub/cub/agent/agent_scan.cuh +++ b/cub/cub/agent/agent_scan.cuh @@ -47,8 +47,8 @@ #include #include #include +#include #include -#include #include #include @@ -111,27 +111,6 @@ struct AgentScanPolicy : ScalingType }; }; -#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) -namespace detail -{ -// Only define this when needed. -// Because of overload woes, this depends on C++20 concepts. util_device.h checks that concepts are available when -// either runtime policies or PTX JSON information are enabled, so if they are, this is always valid. The generic -// version is always defined, and that's the only one needed for regular CUB operations. -// -// TODO: enable this unconditionally once concepts are always available -CUB_DETAIL_POLICY_WRAPPER_DEFINE( - ScanAgentPolicy, - (GenericAgentPolicy), - (BLOCK_THREADS, BlockThreads, int), - (ITEMS_PER_THREAD, ItemsPerThread, int), - (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), - (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), - (STORE_ALGORITHM, StoreAlgorithm, cub::BlockStoreAlgorithm), - (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm)) -} // namespace detail -#endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) - /****************************************************************************** * Thread block abstractions ******************************************************************************/ diff --git a/cub/cub/device/dispatch/tuning/tuning_scan.cuh b/cub/cub/device/dispatch/tuning/tuning_scan.cuh index 97f5edee527..10ec520fff8 100644 --- a/cub/cub/device/dispatch/tuning/tuning_scan.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_scan.cuh @@ -469,16 +469,6 @@ struct ScanPolicyWrapper() = Scan().EncodedPolicy(), - key<"DelayConstructor">() = - StaticPolicyT::ScanPolicyT::detail::delay_constructor_t::EncodedConstructor()>(); - } -#endif }; template From fb77098e3ab77da531aec513a9f9bba0b9d49dc6 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 25 Sep 2025 16:09:39 +0000 Subject: [PATCH 077/100] Update retrieval of delay constructor --- c/parallel/src/segmented_sort.cu | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index bae48a7dfb9..dfe4b12c5fb 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -455,21 +455,16 @@ struct partition_runtime_tuning_policy std::string get_three_way_partition_policy_delay_constructor(const nlohmann::json& partition_policy) { - auto dc_json = partition_policy["ThreeWayPartitionPolicyDelayConstructor"]; - auto delay_constructor_type = dc_json["type"].get(); + auto delay_ctor_info = partition_policy["DelayConstructor"]; - if (delay_constructor_type == "fixed_delay_constructor_t") + std::string delay_ctor_params; + for (auto&& param : delay_ctor_info["params"]) { - auto delay = dc_json["delay"].get(); - auto l2_write_latency = dc_json["l2_write_latency"].get(); - return std::format("cub::detail::fixed_delay_constructor_t<{}, {}>", delay, l2_write_latency); + delay_ctor_params.append(to_string(param) + ", "); } - else if (delay_constructor_type == "no_delay_constructor_t") - { - auto l2_write_latency = dc_json["l2_write_latency"].get(); - return std::format("cub::detail::no_delay_constructor_t<{}>", l2_write_latency); - } - throw std::runtime_error("Invalid delay constructor type: " + delay_constructor_type); + delay_ctor_params.erase(delay_ctor_params.size() - 2); // remove last ", " + + return std::format("cub::detail::{}<{}>", delay_ctor_info["name"].get(), delay_ctor_params); } std::string inject_delay_constructor_into_three_way_policy( From 8d28b6cc1f736c50ba050dd604b87da2ef805abb Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Sat, 27 Sep 2025 01:01:31 +0000 Subject: [PATCH 078/100] Add comments at end of idefs --- cub/cub/agent/single_pass_scan_operators.cuh | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/cub/cub/agent/single_pass_scan_operators.cuh b/cub/cub/agent/single_pass_scan_operators.cuh index 5a0684f66e1..9786294070a 100644 --- a/cub/cub/agent/single_pass_scan_operators.cuh +++ b/cub/cub/agent/single_pass_scan_operators.cuh @@ -228,7 +228,7 @@ struct no_delay_constructor_t using namespace ptx_json; return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; template @@ -265,7 +265,7 @@ struct reduce_by_key_delay_constructor_t return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; template @@ -330,7 +330,7 @@ struct exponential_backoff_constructor_t return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; template @@ -377,7 +377,7 @@ struct exponential_backoff_jitter_constructor_t return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; template @@ -424,7 +424,7 @@ struct exponential_backoff_jitter_window_constructor_t return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; template @@ -474,7 +474,7 @@ struct exponential_backon_jitter_window_constructor_t return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; template @@ -523,7 +523,7 @@ struct exponential_backon_jitter_constructor_t return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; template @@ -560,7 +560,7 @@ struct exponential_backon_constructor_t return object() = value(), key<"params">() = array()>(); } -#endif +#endif // CUB_ENABLE_POLICY_PTX_JSON }; using default_no_delay_constructor_t = no_delay_constructor_t<450>; From 0d0baf6751f85c5be5cf0c75024920024a5312f2 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 30 Sep 2025 00:51:07 +0000 Subject: [PATCH 079/100] remove redundant ptx json from delay constructors --- cub/cub/agent/single_pass_scan_operators.cuh | 30 -------------------- 1 file changed, 30 deletions(-) diff --git a/cub/cub/agent/single_pass_scan_operators.cuh b/cub/cub/agent/single_pass_scan_operators.cuh index a12fc673db0..6a2a741b5e8 100644 --- a/cub/cub/agent/single_pass_scan_operators.cuh +++ b/cub/cub/agent/single_pass_scan_operators.cuh @@ -579,36 +579,6 @@ using default_reduce_by_key_delay_constructor_t = reduce_by_key_delay_constructor_t<350, 450>, default_delay_constructor_t>>; -#if defined(CUB_ENABLE_POLICY_PTX_JSON) -# include - -// ptx-json encoders for delay constructor types. Unlike the other agent policy -// member variables, this is defined as a type alias so we can't use the -// CUB_DETAIL_POLICY_WRAPPER_DEFINE macro to embed it with ptx-json. To work -// around this, we define the ptx-json encoders here. These can then be used in -// the policy wrapper's EncodedPolicy member function to explicitly encode the -// delay constructor. - -template -struct delay_constructor_json; - -template -struct delay_constructor_json> -{ - using type = - ptx_json::object() = ptx_json::value(), - ptx_json::key<"delay">() = ptx_json::value(), - ptx_json::key<"l2_write_latency">() = ptx_json::value()>; -}; - -template -struct delay_constructor_json> -{ - using type = ptx_json::object() = ptx_json::value(), - ptx_json::key<"l2_write_latency">() = ptx_json::value()>; -}; -#endif // CUB_ENABLE_POLICY_PTX_JSON - /** * @brief Alias template for a ScanTileState specialized for a given value type, `T`, and memory order `Order`. * From 0d47ca401867362ae77fa8fcca6cd60a38d82e38 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 30 Sep 2025 04:01:09 +0000 Subject: [PATCH 080/100] Fix variable shadowing warning in MSVC --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 85ccd9662cb..6f9688ecc75 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -785,7 +785,7 @@ private: // Signed integer type for global offsets // Check if the number of items exceeds the range covered by the selected signed offset type - cudaError_t error = ChooseOffsetT::is_exceeding_offset_type(num_items); + error = ChooseOffsetT::is_exceeding_offset_type(num_items); if (error) { return error; From 9b7cb1c86425146cdb1a4552a4a9505ec03c70c8 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 30 Sep 2025 13:12:30 +0000 Subject: [PATCH 081/100] Fix other variable shadowing warning in MSVC --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index 6f9688ecc75..e7457c13492 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -504,7 +504,7 @@ struct DispatchSegmentedSort // Signed integer type for global offsets // Check if the number of items exceeds the range covered by the selected signed offset type - cudaError_t error = ChooseOffsetT::is_exceeding_offset_type(num_items); + error = ChooseOffsetT::is_exceeding_offset_type(num_items); if (error) { return error; From 56de6c19c80e7ec0391d8548d31f07be1a7e380c Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 30 Sep 2025 13:37:12 +0000 Subject: [PATCH 082/100] Add methods to retrieve enums used for asserts --- c/parallel/src/segmented_sort.cu | 57 +++++++++++++++++--------------- 1 file changed, 30 insertions(+), 27 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index dfe4b12c5fb..e519e3ec244 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -369,33 +369,6 @@ struct segmented_sort_runtime_tuning_policy return medium_segment; } - void CheckLoadModifierIsNotLDG() const - { - if (large_segment.LoadModifier() == cub::CacheLoadModifier::LOAD_LDG) - { - throw std::runtime_error("The memory consistency model does not apply to texture accesses"); - } - } - - void CheckLoadAlgorithmIsNotStriped() const - { - if (large_segment.LoadAlgorithm() == cub::BLOCK_LOAD_STRIPED - || medium_segment.LoadAlgorithm() == cub::WARP_LOAD_STRIPED - || small_segment.LoadAlgorithm() == cub::WARP_LOAD_STRIPED) - { - throw std::runtime_error("Striped load will make this algorithm unstable"); - } - } - - void CheckStoreAlgorithmIsNotStriped() const - { - if (medium_segment.StoreAlgorithm() == cub::WARP_STORE_STRIPED - || small_segment.StoreAlgorithm() == cub::WARP_STORE_STRIPED) - { - throw std::runtime_error("Striped stores will produce unsorted results"); - } - } - int PartitioningThreshold() const { return partitioning_threshold; @@ -426,6 +399,36 @@ struct segmented_sort_runtime_tuning_policy return medium_segment.ItemsPerTile(); } + cub::CacheLoadModifier LargeSegmentLoadModifier() const + { + return large_segment.LoadModifier(); + } + + cub::BlockLoadAlgorithm LargeSegmentLoadAlgorithm() const + { + return large_segment.LoadAlgorithm(); + } + + cub::WarpLoadAlgorithm MediumSegmentLoadAlgorithm() const + { + return medium_segment.LoadAlgorithm(); + } + + cub::WarpLoadAlgorithm SmallSegmentLoadAlgorithm() const + { + return small_segment.LoadAlgorithm(); + } + + cub::WarpStoreAlgorithm MediumSegmentStoreAlgorithm() const + { + return medium_segment.StoreAlgorithm(); + } + + cub::WarpStoreAlgorithm SmallSegmentStoreAlgorithm() const + { + return small_segment.StoreAlgorithm(); + } + using MaxPolicy = segmented_sort_runtime_tuning_policy; template From 60bf5516c757b96c193fcd342ad3e36a1f7e1e2c Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 30 Sep 2025 13:39:12 +0000 Subject: [PATCH 083/100] Replace Check*() methods with CUB_DETAIL_STATIC_ISH_ASSERT --- .../dispatch/dispatch_segmented_sort.cuh | 20 +++++--- .../dispatch/tuning/tuning_segmented_sort.cuh | 51 +++++++++++-------- 2 files changed, 42 insertions(+), 29 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index e7457c13492..cd010984aaf 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -416,14 +416,18 @@ struct DispatchSegmentedSort { auto wrapped_policy = detail::segmented_sort::MakeSegmentedSortPolicyWrapper(policy); - wrapped_policy.CheckLoadModifierIsNotLDG(); - - if constexpr (!KEYS_ONLY) - { - wrapped_policy.CheckLoadAlgorithmIsNotStriped(); - } - - wrapped_policy.CheckStoreAlgorithmIsNotStriped(); + CUB_DETAIL_STATIC_ISH_ASSERT(wrapped_policy.LargeSegmentLoadModifier() != CacheLoadModifier::LOAD_LDG, + "The memory consistency model does not apply to texture accesses"); + + CUB_DETAIL_STATIC_ISH_ASSERT( + KEYS_ONLY || wrapped_policy.LargeSegmentLoadAlgorithm() != BLOCK_LOAD_STRIPED + || wrapped_policy.MediumSegmentLoadAlgorithm() != WARP_LOAD_STRIPED + || wrapped_policy.SmallSegmentLoadAlgorithm() != WARP_LOAD_STRIPED, + "Striped load will make this algorithm unstable"); + + CUB_DETAIL_STATIC_ISH_ASSERT(wrapped_policy.MediumSegmentStoreAlgorithm() != WARP_STORE_STRIPED + || wrapped_policy.SmallSegmentStoreAlgorithm() != WARP_STORE_STRIPED, + "Striped stores will produce unsorted results"); const int radix_bits = wrapped_policy.LargeSegmentRadixBits(); diff --git a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index becb8d09cbc..481ba246815 100644 --- a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -81,27 +81,6 @@ struct SegmentedSortPolicyWrapper Date: Tue, 30 Sep 2025 14:42:58 +0000 Subject: [PATCH 084/100] static cast to handle msvc warning --- cub/cub/device/dispatch/dispatch_segmented_sort.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh index cd010984aaf..e7c67d146d4 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -704,7 +704,7 @@ private: CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE int GetNumPasses(int radix_bits) { constexpr int byte_size = 8; - const int num_bits = kernel_source.KeySize() * byte_size; + const int num_bits = static_cast(kernel_source.KeySize()) * byte_size; const int num_passes = ::cuda::ceil_div(num_bits, radix_bits); return num_passes; } From 777cbd988d802215e292652c4a80930f66031463 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 1 Oct 2025 14:52:58 +0000 Subject: [PATCH 085/100] Use different OffsetT to not break windows build --- c/parallel/src/segmented_sort.cu | 2 +- c/parallel/test/test_segmented_sort.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index e519e3ec244..109d29d8200 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -39,7 +39,7 @@ struct device_segmented_sort_policy; struct device_three_way_partition_policy; -using OffsetT = long; +using OffsetT = ptrdiff_t; static_assert(std::is_same_v, OffsetT>, "OffsetT must be long"); // check we can map OffsetT to cuda::std::int64_t diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index a7ceb9e0dfe..14bb5d948ab 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -28,7 +28,7 @@ using item_t = float; using BuildResultT = cccl_device_segmented_sort_build_result_t; -using SizeT = long; +using SizeT = ptrdiff_t; struct segmented_sort_cleanup { From 754fc0ff3ccbc9d44c92d0337276f65bc18efbff Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 1 Oct 2025 15:05:36 +0000 Subject: [PATCH 086/100] Fix merge conflict --- c/parallel/test/test_segmented_reduce.cpp | 41 ----------------------- 1 file changed, 41 deletions(-) diff --git a/c/parallel/test/test_segmented_reduce.cpp b/c/parallel/test/test_segmented_reduce.cpp index 1ac2535af8d..91239275c1c 100644 --- a/c/parallel/test/test_segmented_reduce.cpp +++ b/c/parallel/test/test_segmented_reduce.cpp @@ -99,47 +99,6 @@ void segmented_reduce( // Test section // ============== -<<<<<<< HEAD -======= -static std::tuple make_step_counting_iterator_sources( - std::string_view index_ty_name, - std::string_view state_name, - std::string_view advance_fn_name, - std::string_view dereference_fn_name) -{ - static constexpr std::string_view it_state_src_tmpl = R"XXX( -struct {0} {{ - {1} linear_id; - {1} row_size; -}}; -)XXX"; - - const std::string it_state_def_src = std::format(it_state_src_tmpl, state_name, index_ty_name); - - static constexpr std::string_view it_def_src_tmpl = R"XXX( -extern "C" __device__ void {0}({1}* state, {2} offset) -{{ - state->linear_id += offset; -}} -)XXX"; - - const std::string it_advance_fn_def_src = - std::format(it_def_src_tmpl, /*0*/ advance_fn_name, state_name, index_ty_name); - - static constexpr std::string_view it_deref_src_tmpl = R"XXX( -extern "C" __device__ void {0}({1}* state, {2}* result) -{{ - *result = (state->linear_id) * (state->row_size); -}} -)XXX"; - - const std::string it_deref_fn_def_src = - std::format(it_deref_src_tmpl, dereference_fn_name, state_name, index_ty_name); - - return std::make_tuple(it_state_def_src, it_advance_fn_def_src, it_deref_fn_def_src); -} - ->>>>>>> main struct SegmentedReduce_SumOverRows_Fixture_Tag; C2H_TEST_LIST("segmented_reduce can sum over rows of matrix with integral type", "[segmented_reduce]", From fac34037472be62cc67c0b7c29362a1d240025a8 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 28 Oct 2025 22:05:50 +0000 Subject: [PATCH 087/100] Address review comments --- c/parallel/src/segmented_sort.cu | 55 ++++++++++++++++++-------------- 1 file changed, 31 insertions(+), 24 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 109d29d8200..9e00a6e1ebe 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -180,7 +180,7 @@ cccl_op_t make_segments_selector_op( size_t num_lto_opts) { cccl_op_t selector_op{}; - selector_state_t* selector_op_state = new selector_state_t{}; + auto selector_op_state = std::make_unique(); std::string offset_t; check(nvrtcGetTypeName(&offset_t)); @@ -202,8 +202,9 @@ extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* resul auto* st = static_cast(state_ptr); const local_segment_index_t sid = *static_cast(arg_ptr); - const {2} begin = static_cast(st->begin_offsets)[st->base_segment_offset + sid]; - const {3} end = static_cast(st->end_offsets)[st->base_segment_offset + sid]; + const global_segment_offset_t index = st->base_segment_offset + static_cast(sid); + const {2} begin = static_cast(st->begin_offsets)[index]; + const {3} end = static_cast(st->end_offsets)[index]; const bool pred = (end - begin) {4} st->threshold; *static_cast(result_ptr) = pred; }} @@ -226,9 +227,10 @@ extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* resul selector_op.code_type = CCCL_OP_LTOIR; selector_op.size = sizeof(selector_state_t); selector_op.alignment = alignof(selector_state_t); - selector_op.state = selector_op_state; + selector_op.state = selector_op_state.get(); selector_op_state->initialize(offset, begin_offset_iterator, end_offset_iterator); + selector_op_state.release(); return selector_op; } @@ -276,16 +278,16 @@ struct segmented_sort_kernel_source void SetSegmentOffset(indirect_arg_t& selector, long long base_segment_offset) const { - auto* st = reinterpret_cast(selector.ptr); + auto* st = static_cast(selector.ptr); st->base_segment_offset = base_segment_offset; } }; std::string get_three_way_partition_init_kernel_name() { - constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT"; + static constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT"; - constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + static constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; return std::format("cub::detail::three_way_partition::DeviceThreeWayPartitionInitKernel<{0}, {1}>", scan_tile_state_t, // 0 @@ -297,19 +299,19 @@ std::string get_three_way_partition_kernel_name(std::string_view large_selector_ std::string chained_policy_t; check(nvrtcGetTypeName(&chained_policy_t)); - constexpr std::string_view input_it_t = + static constexpr std::string_view input_it_t = "thrust::counting_iterator"; - constexpr std::string_view first_out_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; - constexpr std::string_view second_out_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; - constexpr std::string_view unselected_out_it_t = + static constexpr std::string_view first_out_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + static constexpr std::string_view second_out_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + static constexpr std::string_view unselected_out_it_t = "thrust::reverse_iterator"; - constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; - constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT"; + static constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; + static constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT"; std::string offset_t; check(nvrtcGetTypeName(&offset_t)); - constexpr std::string_view per_partition_offset_t = "cub::detail::three_way_partition::per_partition_offset_t"; - constexpr std::string_view streaming_context_t = + static constexpr std::string_view per_partition_offset_t = "cub::detail::three_way_partition::per_partition_offset_t"; + static constexpr std::string_view streaming_context_t = "cub::detail::three_way_partition::streaming_context_t"; return std::format( @@ -474,8 +476,8 @@ std::string inject_delay_constructor_into_three_way_policy( const std::string& three_way_partition_policy_str, const std::string& delay_constructor_type) { // Insert before the final closing of the struct (right before the sequence "};") - const std::string needle = "};"; - const auto pos = three_way_partition_policy_str.rfind(needle); + static constexpr std::string_view needle = "};"; + const auto pos = three_way_partition_policy_str.rfind(needle); if (pos == std::string::npos) { return three_way_partition_policy_str; // unexpected; return as-is @@ -708,7 +710,7 @@ struct __align__({3}) items_storage_t {{ segmented_sort::inject_delay_constructor_into_three_way_policy( three_way_partition_policy_str, three_way_partition_policy_delay_constructor); - constexpr std::string_view program_preamble_template = R"XXX( + static constexpr std::string_view program_preamble_template = R"XXX( #include #include #include // used in three_way_partition kernel @@ -1010,15 +1012,20 @@ CUresult cccl_device_segmented_sort_cleanup(cccl_device_segmented_sort_build_res std::unique_ptr cubin(reinterpret_cast(build_ptr->cubin)); // Clean up the selector op states - delete static_cast(build_ptr->large_segments_selector_op.state); - delete static_cast(build_ptr->small_segments_selector_op.state); + std::unique_ptr large_state( + static_cast(build_ptr->large_segments_selector_op.state)); + std::unique_ptr small_state( + static_cast(build_ptr->small_segments_selector_op.state)); - delete[] const_cast(build_ptr->large_segments_selector_op.code); - delete[] const_cast(build_ptr->small_segments_selector_op.code); + // Clean up the selector op code buffers + std::unique_ptr large_code(const_cast(build_ptr->large_segments_selector_op.code)); + std::unique_ptr small_code(const_cast(build_ptr->small_segments_selector_op.code)); // Clean up the runtime policies - delete static_cast(build_ptr->runtime_policy); - delete static_cast(build_ptr->partition_runtime_policy); + std::unique_ptr rtp( + static_cast(build_ptr->runtime_policy)); + std::unique_ptr prtp( + static_cast(build_ptr->partition_runtime_policy)); check(cuLibraryUnload(build_ptr->library)); } catch (const std::exception& exc) From 7e383a0d4c3020304d89ee01df515a7b88486c50 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 31 Oct 2025 22:51:39 +0000 Subject: [PATCH 088/100] Implement single compilation for segmented_sort. This required making some changes in the CUB policies to make sure we do not have any ambiguous overloads --- c/parallel/src/radix_sort.cu | 23 ++- c/parallel/src/segmented_sort.cu | 194 ++++++++---------- cub/cub/agent/agent_radix_sort_downsweep.cuh | 6 +- cub/cub/agent/agent_three_way_partition.cuh | 6 +- .../dispatch/tuning/tuning_segmented_sort.cuh | 36 ++-- 5 files changed, 124 insertions(+), 141 deletions(-) diff --git a/c/parallel/src/radix_sort.cu b/c/parallel/src/radix_sort.cu index 85ade712fc7..de27ae1c293 100644 --- a/c/parallel/src/radix_sort.cu +++ b/c/parallel/src/radix_sort.cu @@ -39,11 +39,11 @@ struct radix_sort_runtime_tuning_policy RuntimeRadixSortExclusiveSumAgentPolicy exclusive_sum; RuntimeRadixSortOnesweepAgentPolicy onesweep; cub::detail::RuntimeScanAgentPolicy scan; - RuntimeRadixSortDownsweepAgentPolicy downsweep; - RuntimeRadixSortDownsweepAgentPolicy alt_downsweep; + cub::detail::RuntimeRadixSortDownsweepAgentPolicy downsweep; + cub::detail::RuntimeRadixSortDownsweepAgentPolicy alt_downsweep; RuntimeRadixSortUpsweepAgentPolicy upsweep; RuntimeRadixSortUpsweepAgentPolicy alt_upsweep; - RuntimeRadixSortDownsweepAgentPolicy single_tile; + cub::detail::RuntimeRadixSortDownsweepAgentPolicy single_tile; bool is_onesweep; auto Histogram() const @@ -440,13 +440,16 @@ __device__ consteval auto& policy_generator() {{ using namespace cub::detail::radix_sort_runtime_policies; using cub::detail::RuntimeScanAgentPolicy; - auto single_tile_policy = RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "SingleTilePolicy"); - auto onesweep_policy = RuntimeRadixSortOnesweepAgentPolicy::from_json(runtime_policy, "OnesweepPolicy"); - auto upsweep_policy = RuntimeRadixSortUpsweepAgentPolicy::from_json(runtime_policy, "UpsweepPolicy"); - auto alt_upsweep_policy = RuntimeRadixSortUpsweepAgentPolicy::from_json(runtime_policy, "AltUpsweepPolicy"); - auto downsweep_policy = RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "DownsweepPolicy"); - auto alt_downsweep_policy = RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "AltDownsweepPolicy"); - auto histogram_policy = RuntimeRadixSortHistogramAgentPolicy::from_json(runtime_policy, "HistogramPolicy"); + auto single_tile_policy = + cub::detail::RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "SingleTilePolicy"); + auto onesweep_policy = RuntimeRadixSortOnesweepAgentPolicy::from_json(runtime_policy, "OnesweepPolicy"); + auto upsweep_policy = RuntimeRadixSortUpsweepAgentPolicy::from_json(runtime_policy, "UpsweepPolicy"); + auto alt_upsweep_policy = RuntimeRadixSortUpsweepAgentPolicy::from_json(runtime_policy, "AltUpsweepPolicy"); + auto downsweep_policy = + cub::detail::RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "DownsweepPolicy"); + auto alt_downsweep_policy = + cub::detail::RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "AltDownsweepPolicy"); + auto histogram_policy = RuntimeRadixSortHistogramAgentPolicy::from_json(runtime_policy, "HistogramPolicy"); auto exclusive_sum_policy = RuntimeRadixSortExclusiveSumAgentPolicy::from_json(runtime_policy, "ExclusiveSumPolicy"); auto scan_policy = RuntimeScanAgentPolicy::from_json(runtime_policy, "ScanPolicy"); diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 9e00a6e1ebe..3773d30c1c5 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -10,6 +10,7 @@ #include // cub::detail::choose_offset_t #include // cub::detail::CudaDriverLauncherFactory +#include #include // cub::DispatchSegmentedSort #include // DeviceSegmentedSort kernels #include // policy_hub @@ -28,7 +29,6 @@ #include "util/context.h" #include "util/errors.h" #include "util/indirect_arg.h" -#include "util/runtime_policy.h" #include "util/types.h" #include #include // cccl_type_info @@ -574,7 +574,7 @@ CUresult cccl_device_segmented_sort_build_ex( const std::string arch = std::format("-arch=sm_{0}{1}", cc_major, cc_minor); - std::vector args = { + std::vector selector_compilation_args = { arch.c_str(), cub_path, thrust_path, @@ -585,7 +585,7 @@ CUresult cccl_device_segmented_sort_build_ex( "-DCUB_DISABLE_CDP", "-std=c++20"}; - cccl::detail::extend_args_with_build_config(args, config); + cccl::detail::extend_args_with_build_config(selector_compilation_args, config); constexpr size_t num_lto_args = 2; const char* lopts[num_lto_args] = {"-lto", arch.c_str()}; @@ -596,8 +596,8 @@ CUresult cccl_device_segmented_sort_build_ex( end_offset_it, "cccl_large_segments_selector_op", ">", - args.data(), - args.size(), + selector_compilation_args.data(), + selector_compilation_args.size(), lopts, num_lto_args); cccl_op_t small_selector_op = segmented_sort::make_segments_selector_op( @@ -606,8 +606,8 @@ CUresult cccl_device_segmented_sort_build_ex( end_offset_it, "cccl_small_segments_selector_op", "<", - args.data(), - args.size(), + selector_compilation_args.data(), + selector_compilation_args.size(), lopts, num_lto_args); @@ -623,123 +623,80 @@ CUresult cccl_device_segmented_sort_build_ex( const auto [small_selector_name, small_selector_src] = get_specialization( template_id(), small_selector_op, selector_result_t, selector_input_t); - const std::string dependent_definitions_src = std::format( - R"XXX( -struct __align__({1}) storage_t {{ - char data[{0}]; -}}; -struct __align__({3}) items_storage_t {{ - char data[{2}]; -}}; -{4} -{5} -{6} -{7} -{8} -{9} -)XXX", - keys_in_it.value_type.size, // 0 - keys_in_it.value_type.alignment, // 1 - values_in_it.value_type.size, // 2 - values_in_it.value_type.alignment, // 3 - keys_in_iterator_src, // 4 - values_in_iterator_src, // 5 - start_offset_iterator_src, // 6 - end_offset_iterator_src, // 7 - large_selector_src, // 8 - small_selector_src); // 9 - - const std::string ptx_arch = std::format("-arch=compute_{}{}", cc_major, cc_minor); - - constexpr size_t ptx_num_args = 6; - const char* ptx_args[ptx_num_args] = { - ptx_arch.c_str(), cub_path, thrust_path, libcudacxx_path, ctk_path, "-rdc=true"}; - - static constexpr std::string_view policy_wrapper_expr_tmpl = - R"XXXX(cub::detail::segmented_sort::MakeSegmentedSortPolicyWrapper(cub::detail::segmented_sort::policy_hub<{0}, {1}>::MaxPolicy::ActivePolicy{{}}))XXXX"; - - const auto policy_wrapper_expr = std::format( - policy_wrapper_expr_tmpl, + const auto segmented_sort_policy_hub_expr = std::format( + "cub::detail::segmented_sort::policy_hub<{0}, {1}>", key_t, // 0 value_t); // 1 - static constexpr std::string_view ptx_query_tu_src_tmpl = R"XXXX( + static constexpr std::string_view three_way_partition_policy_hub_expr = + "cub::detail::three_way_partition::policy_hub"; + + const std::string final_src = std::format( + R"XXX( #include #include #include #include -{0} -{1} -)XXXX"; - - const auto ptx_query_tu_src = - std::format(ptx_query_tu_src_tmpl, jit_template_header_contents, dependent_definitions_src); - - nlohmann::json runtime_policy = get_policy(policy_wrapper_expr, ptx_query_tu_src, ptx_args); - - using cub::detail::RuntimeRadixSortDownsweepAgentPolicy; - auto [large_segment_policy, large_segment_policy_str] = - RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "LargeSegmentPolicy"); - - using cub::detail::RuntimeSubWarpMergeSortAgentPolicy; - auto [small_segment_policy, small_segment_policy_str] = - RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "SmallSegmentPolicy"); - - auto [medium_segment_policy, medium_segment_policy_str] = - RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "MediumSegmentPolicy"); - - auto partitioning_threshold = runtime_policy["PartitioningThreshold"].get(); - static constexpr std::string_view partition_policy_wrapper_expr_tmpl = - R"XXXX(cub::detail::three_way_partition::MakeThreeWayPartitionPolicyWrapper(cub::detail::three_way_partition::policy_hub<{0}, {1}>::MaxPolicy::ActivePolicy{{}}))XXXX"; - const auto partition_policy_wrapper_expr = std::format( - partition_policy_wrapper_expr_tmpl, - "cub::detail::segmented_sort::local_segment_index_t", - "cub::detail::three_way_partition::per_partition_offset_t"); - - nlohmann::json partition_policy = get_policy(partition_policy_wrapper_expr, ptx_query_tu_src, ptx_args); - - using cub::detail::RuntimeThreeWayPartitionAgentPolicy; - auto [three_way_partition_policy, three_way_partition_policy_str] = - RuntimeThreeWayPartitionAgentPolicy::from_json(partition_policy, "ThreeWayPartitionPolicy"); - - const std::string three_way_partition_policy_delay_constructor = - segmented_sort::get_three_way_partition_policy_delay_constructor(partition_policy); - - const std::string injected_three_way_partition_policy_str = - segmented_sort::inject_delay_constructor_into_three_way_policy( - three_way_partition_policy_str, three_way_partition_policy_delay_constructor); +{0} - static constexpr std::string_view program_preamble_template = R"XXX( -#include -#include #include // used in three_way_partition kernel #include // used in three_way_partition kernel #include // used in three_way_partition kernel -{0} -{1} -struct device_segmented_sort_policy {{ - struct ActivePolicy {{ - {2} - {3} - {4} - }}; + +struct __align__({2}) storage_t {{ + char data[{1}]; }}; -struct device_three_way_partition_policy {{ - struct ActivePolicy {{ - {5} - }}; +struct __align__({4}) items_storage_t {{ + char data[{3}]; }}; -)XXX"; - - std::string final_src = std::format( - program_preamble_template, +{5} +{6} +{7} +{8} +{9} +{10} +using device_segmented_sort_policy = {11}::MaxPolicy; +using device_three_way_partition_policy = {12}::MaxPolicy; + +#include +__device__ consteval auto& segmented_sort_policy_generator() {{ + return ptx_json::id() + = cub::detail::segmented_sort::SegmentedSortPolicyWrapper::EncodedPolicy(); +}} +__device__ consteval auto& three_way_partition_policy_generator() {{ + return ptx_json::id() + = cub::detail::three_way_partition::ThreeWayPartitionPolicyWrapper::EncodedPolicy(); +}} +)XXX", jit_template_header_contents, // 0 - dependent_definitions_src, // 1 - large_segment_policy_str, // 2 - small_segment_policy_str, // 3 - medium_segment_policy_str, // 4 - injected_three_way_partition_policy_str); // 5 + keys_in_it.value_type.size, // 1 + keys_in_it.value_type.alignment, // 2 + values_in_it.value_type.size, // 3 + values_in_it.value_type.alignment, // 4 + keys_in_iterator_src, // 5 + values_in_iterator_src, // 6 + start_offset_iterator_src, // 7 + end_offset_iterator_src, // 8 + large_selector_src, // 9 + small_selector_src, // 10 + segmented_sort_policy_hub_expr, // 11 + three_way_partition_policy_hub_expr); // 12 + + std::vector args = { + arch.c_str(), + cub_path, + thrust_path, + libcudacxx_path, + ctk_path, + "-rdc=true", + "-dlto", + "-DCUB_DISABLE_CDP", + "-DCUB_ENABLE_POLICY_PTX_JSON", + "-std=c++20"}; + + cccl::detail::extend_args_with_build_config(args, config); std::string segmented_sort_fallback_kernel_name = segmented_sort::get_device_segmented_sort_fallback_kernel_name( start_offset_iterator_name, end_offset_iterator_name, key_t, value_t, sort_order); @@ -810,6 +767,25 @@ struct device_three_way_partition_policy {{ check(cuLibraryGetKernel( &build_ptr->three_way_partition_kernel, build_ptr->library, three_way_partition_kernel_lowered_name.c_str())); + nlohmann::json runtime_policy = + cub::detail::ptx_json::parse("device_segmented_sort_policy", {result.data.get(), result.size}); + + using cub::detail::RuntimeRadixSortDownsweepAgentPolicy; + auto large_segment_policy = RuntimeRadixSortDownsweepAgentPolicy::from_json(runtime_policy, "LargeSegmentPolicy"); + + using cub::detail::RuntimeSubWarpMergeSortAgentPolicy; + auto small_segment_policy = RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "SmallSegmentPolicy"); + + auto medium_segment_policy = RuntimeSubWarpMergeSortAgentPolicy::from_json(runtime_policy, "MediumSegmentPolicy"); + + int partitioning_threshold = runtime_policy["PartitioningThreshold"].get(); + nlohmann::json partition_policy = + cub::detail::ptx_json::parse("device_three_way_partition_policy", {result.data.get(), result.size}); + + using cub::detail::RuntimeThreeWayPartitionAgentPolicy; + auto three_way_partition_policy = + RuntimeThreeWayPartitionAgentPolicy::from_json(partition_policy, "ThreeWayPartitionPolicy"); + build_ptr->cc = cc; build_ptr->large_segments_selector_op = large_selector_op; build_ptr->small_segments_selector_op = small_selector_op; diff --git a/cub/cub/agent/agent_radix_sort_downsweep.cuh b/cub/cub/agent/agent_radix_sort_downsweep.cuh index 96c93275f16..ef4d1d8361e 100644 --- a/cub/cub/agent/agent_radix_sort_downsweep.cuh +++ b/cub/cub/agent/agent_radix_sort_downsweep.cuh @@ -125,7 +125,7 @@ struct AgentRadixSortDownsweepPolicy : ScalingType }; #if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) -namespace detail::radix_sort_runtime_policies +namespace detail { // Only define this when needed. // Because of overload woes, this depends on C++20 concepts. util_device.h checks that concepts are available when @@ -135,7 +135,7 @@ namespace detail::radix_sort_runtime_policies // TODO: enable this unconditionally once concepts are always available CUB_DETAIL_POLICY_WRAPPER_DEFINE( RadixSortDownsweepAgentPolicy, - (RadixSortUpsweepAgentPolicy, UniqueByKeyAgentPolicy), + (cub::detail::radix_sort_runtime_policies::RadixSortUpsweepAgentPolicy, UniqueByKeyAgentPolicy), (BLOCK_THREADS, BlockThreads, int), (ITEMS_PER_THREAD, ItemsPerThread, int), (RADIX_BITS, RadixBits, int), @@ -143,7 +143,7 @@ CUB_DETAIL_POLICY_WRAPPER_DEFINE( (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier), (RANK_ALGORITHM, RankAlgorithm, cub::RadixRankAlgorithm), (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm)) -} // namespace detail::radix_sort_runtime_policies +} // namespace detail #endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) /****************************************************************************** diff --git a/cub/cub/agent/agent_three_way_partition.cuh b/cub/cub/agent/agent_three_way_partition.cuh index 0e6d6e8fdb8..1f3b58dcb0d 100644 --- a/cub/cub/agent/agent_three_way_partition.cuh +++ b/cub/cub/agent/agent_three_way_partition.cuh @@ -46,6 +46,10 @@ #include #include +#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON) +# include // for UniqueByKeyAgentPolicy +#endif + #include #include #include @@ -82,7 +86,7 @@ namespace detail { CUB_DETAIL_POLICY_WRAPPER_DEFINE( ThreeWayPartitionAgentPolicy, - (GenericAgentPolicy), + (UniqueByKeyAgentPolicy), (BLOCK_THREADS, BlockThreads, int), (ITEMS_PER_THREAD, ItemsPerThread, int), (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm), diff --git a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index c19350147ab..677a9e1481c 100644 --- a/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -49,7 +49,7 @@ namespace detail::segmented_sort template struct SegmentedSortPolicyWrapper : PolicyT { - CUB_RUNTIME_FUNCTION SegmentedSortPolicyWrapper(PolicyT base) + _CCCL_HOST_DEVICE SegmentedSortPolicyWrapper(PolicyT base) : PolicyT(base) {} }; @@ -60,81 +60,81 @@ struct SegmentedSortPolicyWrapper> : StaticPolicyT { - CUB_RUNTIME_FUNCTION SegmentedSortPolicyWrapper(StaticPolicyT base) + _CCCL_HOST_DEVICE SegmentedSortPolicyWrapper(StaticPolicyT base) : StaticPolicyT(base) {} - CUB_RUNTIME_FUNCTION static constexpr auto LargeSegment() + _CCCL_HOST_DEVICE static constexpr auto LargeSegment() { return cub::detail::MakePolicyWrapper(typename StaticPolicyT::LargeSegmentPolicy()); } - CUB_RUNTIME_FUNCTION static constexpr auto SmallSegment() + _CCCL_HOST_DEVICE static constexpr auto SmallSegment() { return cub::detail::MakePolicyWrapper(typename StaticPolicyT::SmallSegmentPolicy()); } - CUB_RUNTIME_FUNCTION static constexpr auto MediumSegment() + _CCCL_HOST_DEVICE static constexpr auto MediumSegment() { return cub::detail::MakePolicyWrapper(typename StaticPolicyT::MediumSegmentPolicy()); } - CUB_RUNTIME_FUNCTION static constexpr int PartitioningThreshold() + _CCCL_HOST_DEVICE static constexpr int PartitioningThreshold() { return StaticPolicyT::PARTITIONING_THRESHOLD; } - CUB_RUNTIME_FUNCTION static constexpr int LargeSegmentRadixBits() + _CCCL_HOST_DEVICE static constexpr int LargeSegmentRadixBits() { return StaticPolicyT::LargeSegmentPolicy::RADIX_BITS; } - CUB_RUNTIME_FUNCTION static constexpr int SegmentsPerSmallBlock() + _CCCL_HOST_DEVICE static constexpr int SegmentsPerSmallBlock() { return StaticPolicyT::SmallSegmentPolicy::SEGMENTS_PER_BLOCK; } - CUB_RUNTIME_FUNCTION static constexpr int SegmentsPerMediumBlock() + _CCCL_HOST_DEVICE static constexpr int SegmentsPerMediumBlock() { return StaticPolicyT::MediumSegmentPolicy::SEGMENTS_PER_BLOCK; } - CUB_RUNTIME_FUNCTION static constexpr int SmallPolicyItemsPerTile() + _CCCL_HOST_DEVICE static constexpr int SmallPolicyItemsPerTile() { return StaticPolicyT::SmallSegmentPolicy::ITEMS_PER_TILE; } - CUB_RUNTIME_FUNCTION static constexpr int MediumPolicyItemsPerTile() + _CCCL_HOST_DEVICE static constexpr int MediumPolicyItemsPerTile() { return StaticPolicyT::MediumSegmentPolicy::ITEMS_PER_TILE; } - CUB_RUNTIME_FUNCTION static constexpr CacheLoadModifier LargeSegmentLoadModifier() + _CCCL_HOST_DEVICE static constexpr CacheLoadModifier LargeSegmentLoadModifier() { return StaticPolicyT::LargeSegmentPolicy::LOAD_MODIFIER; } - CUB_RUNTIME_FUNCTION static constexpr BlockLoadAlgorithm LargeSegmentLoadAlgorithm() + _CCCL_HOST_DEVICE static constexpr BlockLoadAlgorithm LargeSegmentLoadAlgorithm() { return StaticPolicyT::LargeSegmentPolicy::LOAD_ALGORITHM; } - CUB_RUNTIME_FUNCTION static constexpr WarpLoadAlgorithm MediumSegmentLoadAlgorithm() + _CCCL_HOST_DEVICE static constexpr WarpLoadAlgorithm MediumSegmentLoadAlgorithm() { return StaticPolicyT::MediumSegmentPolicy::LOAD_ALGORITHM; } - CUB_RUNTIME_FUNCTION static constexpr WarpLoadAlgorithm SmallSegmentLoadAlgorithm() + _CCCL_HOST_DEVICE static constexpr WarpLoadAlgorithm SmallSegmentLoadAlgorithm() { return StaticPolicyT::SmallSegmentPolicy::LOAD_ALGORITHM; } - CUB_RUNTIME_FUNCTION static constexpr WarpStoreAlgorithm MediumSegmentStoreAlgorithm() + _CCCL_HOST_DEVICE static constexpr WarpStoreAlgorithm MediumSegmentStoreAlgorithm() { return StaticPolicyT::MediumSegmentPolicy::STORE_ALGORITHM; } - CUB_RUNTIME_FUNCTION static constexpr WarpStoreAlgorithm SmallSegmentStoreAlgorithm() + _CCCL_HOST_DEVICE static constexpr WarpStoreAlgorithm SmallSegmentStoreAlgorithm() { return StaticPolicyT::SmallSegmentPolicy::STORE_ALGORITHM; } @@ -152,7 +152,7 @@ struct SegmentedSortPolicyWrapper -CUB_RUNTIME_FUNCTION SegmentedSortPolicyWrapper MakeSegmentedSortPolicyWrapper(PolicyT policy) +_CCCL_HOST_DEVICE SegmentedSortPolicyWrapper MakeSegmentedSortPolicyWrapper(PolicyT policy) { return SegmentedSortPolicyWrapper{policy}; } From b96b1add992a0968bc4541040c78bbb6541604f8 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 31 Oct 2025 23:22:31 +0000 Subject: [PATCH 089/100] Add missing imports and fix imports in test --- python/cuda_cccl/cuda/compute/__init__.py | 4 ++ .../examples/sort/segmented_sort_basic.py | 6 +-- .../examples/sort/segmented_sort_buffer.py | 10 ++--- .../examples/sort/segmented_sort_object.py | 6 +-- .../tests/compute/test_segmented_sort.py | 38 +++++++++---------- 5 files changed, 33 insertions(+), 31 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/__init__.py b/python/cuda_cccl/cuda/compute/__init__.py index c43f6d653b9..43bef744912 100644 --- a/python/cuda_cccl/cuda/compute/__init__.py +++ b/python/cuda_cccl/cuda/compute/__init__.py @@ -17,6 +17,7 @@ make_radix_sort, make_reduce_into, make_segmented_reduce, + make_segmented_sort, make_three_way_partition, make_unary_transform, make_unique_by_key, @@ -24,6 +25,7 @@ radix_sort, reduce_into, segmented_reduce, + segmented_sort, three_way_partition, unary_transform, unique_by_key, @@ -59,6 +61,7 @@ "make_radix_sort", "make_reduce_into", "make_segmented_reduce", + "make_segmented_sort", "make_three_way_partition", "make_unary_transform", "make_unique_by_key", @@ -69,6 +72,7 @@ "reduce_into", "ReverseIterator", "segmented_reduce", + "segmented_sort", "SortOrder", "TransformIterator", "TransformOutputIterator", diff --git a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_basic.py b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_basic.py index 2d37a064a20..53d8cc2361b 100644 --- a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_basic.py +++ b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_basic.py @@ -10,7 +10,7 @@ import cupy as cp import numpy as np -import cuda.cccl.parallel.experimental as parallel +import cuda.compute # Prepare input keys and values, and segment offsets. h_in_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype="int32") @@ -26,7 +26,7 @@ d_out_vals = cp.empty_like(d_in_vals) # Perform the segmented sort (ascending within each segment). -parallel.segmented_sort( +cuda.compute.segmented_sort( d_in_keys, d_out_keys, d_in_vals, @@ -35,7 +35,7 @@ start_offsets.size, cp.asarray(start_offsets), cp.asarray(end_offsets), - parallel.SortOrder.ASCENDING, + cuda.compute.SortOrder.ASCENDING, ) # Verify the result. diff --git a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_buffer.py b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_buffer.py index 12eb7e3b096..2cece439773 100644 --- a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_buffer.py +++ b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_buffer.py @@ -10,7 +10,7 @@ import cupy as cp import numpy as np -import cuda.cccl.parallel.experimental as parallel +import cuda.compute # Prepare input keys and values, and segment offsets. h_in_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype="int32") @@ -26,11 +26,11 @@ d_tmp_vals = cp.empty_like(d_in_vals) # Create double buffers for keys and values. -keys_db = parallel.DoubleBuffer(d_in_keys, d_tmp_keys) -vals_db = parallel.DoubleBuffer(d_in_vals, d_tmp_vals) +keys_db = cuda.compute.DoubleBuffer(d_in_keys, d_tmp_keys) +vals_db = cuda.compute.DoubleBuffer(d_in_vals, d_tmp_vals) # Perform the segmented sort (descending within each segment). -parallel.segmented_sort( +cuda.compute.segmented_sort( keys_db, None, vals_db, @@ -39,7 +39,7 @@ start_offsets.size, cp.asarray(start_offsets), cp.asarray(end_offsets), - parallel.SortOrder.DESCENDING, + cuda.compute.SortOrder.DESCENDING, ) # Verify the result. diff --git a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py index 76f71b42fb6..cde29b8127e 100644 --- a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py +++ b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py @@ -10,7 +10,7 @@ import cupy as cp import numpy as np -import cuda.cccl.parallel.experimental as parallel +import cuda.compute # Prepare the input and segment offsets. dtype = np.int32 @@ -25,14 +25,14 @@ d_output_vals = cp.empty_like(d_input_vals) # Create the segmented sort object. -sorter = parallel.make_segmented_sort( +sorter = cuda.compute.make_segmented_sort( d_input_keys, d_output_keys, d_input_vals, d_output_vals, cp.asarray(start_offsets), cp.asarray(end_offsets), - parallel.SortOrder.ASCENDING, + cuda.compute.SortOrder.ASCENDING, ) # Get the temporary storage size. diff --git a/python/cuda_cccl/tests/compute/test_segmented_sort.py b/python/cuda_cccl/tests/compute/test_segmented_sort.py index 1fdd1efe62a..d6b2f0a0d87 100644 --- a/python/cuda_cccl/tests/compute/test_segmented_sort.py +++ b/python/cuda_cccl/tests/compute/test_segmented_sort.py @@ -9,7 +9,7 @@ import numpy as np import pytest -import cuda.cccl.parallel.experimental as parallel +import cuda.compute DTYPE_LIST = [ np.uint8, @@ -68,7 +68,7 @@ def host_segmented_sort( h_vals: np.ndarray | None, start_offsets: np.ndarray, end_offsets: np.ndarray, - order: parallel.SortOrder, + order: cuda.compute.SortOrder, ) -> Tuple[np.ndarray, np.ndarray | None]: assert start_offsets.shape == end_offsets.shape keys = h_keys.copy() @@ -78,7 +78,7 @@ def host_segmented_sort( if e <= s: continue if vals is None: - if order is parallel.SortOrder.DESCENDING: + if order is cuda.compute.SortOrder.DESCENDING: # stable descending signed_dtype = ( np.dtype(keys.dtype.name.replace("uint", "int")) @@ -92,7 +92,7 @@ def host_segmented_sort( else: # build pairs for stable sort pairs = list(zip(keys[s:e], vals[s:e])) - if order is parallel.SortOrder.DESCENDING: + if order is cuda.compute.SortOrder.DESCENDING: pairs.sort(key=lambda kv: kv[0], reverse=True) else: pairs.sort(key=lambda kv: kv[0]) @@ -107,14 +107,12 @@ def host_segmented_sort( def test_segmented_sort_keys(dtype, num_segments, segment_size, monkeypatch): # Disable SASS verification only for this test when dtype is int64 if np.dtype(dtype) == np.dtype(np.int64): - import cuda.cccl.parallel.experimental._cccl_interop - monkeypatch.setattr( - cuda.cccl.parallel.experimental._cccl_interop, + cuda.compute._cccl_interop, "_check_sass", False, ) - order = parallel.SortOrder.ASCENDING + order = cuda.compute.SortOrder.ASCENDING num_items = num_segments * segment_size h_in_keys = random_array(num_items, dtype, max_value=50) @@ -123,7 +121,7 @@ def test_segmented_sort_keys(dtype, num_segments, segment_size, monkeypatch): d_in_keys = numba.cuda.to_device(h_in_keys) d_out_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) - parallel.segmented_sort( + cuda.compute.segmented_sort( d_in_keys, d_out_keys, None, @@ -145,7 +143,7 @@ def test_segmented_sort_keys(dtype, num_segments, segment_size, monkeypatch): @pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) def test_segmented_sort_pairs(dtype, num_segments, segment_size): - order = parallel.SortOrder.DESCENDING + order = cuda.compute.SortOrder.DESCENDING num_items = num_segments * segment_size h_in_keys = random_array( @@ -160,7 +158,7 @@ def test_segmented_sort_pairs(dtype, num_segments, segment_size): d_out_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) d_out_vals = numba.cuda.to_device(np.empty_like(h_in_vals)) - parallel.segmented_sort( + cuda.compute.segmented_sort( d_in_keys, d_out_keys, d_in_vals, @@ -185,7 +183,7 @@ def test_segmented_sort_pairs(dtype, num_segments, segment_size): @pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) def test_segmented_sort_keys_double_buffer(dtype, num_segments, segment_size): - order = parallel.SortOrder.ASCENDING + order = cuda.compute.SortOrder.ASCENDING num_items = num_segments * segment_size h_in_keys = random_array(num_items, dtype, max_value=20) @@ -193,9 +191,9 @@ def test_segmented_sort_keys_double_buffer(dtype, num_segments, segment_size): d_in_keys = numba.cuda.to_device(h_in_keys) d_tmp_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) - keys_db = parallel.DoubleBuffer(d_in_keys, d_tmp_keys) + keys_db = cuda.compute.DoubleBuffer(d_in_keys, d_tmp_keys) - parallel.segmented_sort( + cuda.compute.segmented_sort( keys_db, None, None, @@ -216,7 +214,7 @@ def test_segmented_sort_keys_double_buffer(dtype, num_segments, segment_size): @pytest.mark.parametrize("dtype, num_segments, segment_size", DTYPE_SEGMENT_PARAMS) def test_segmented_sort_pairs_double_buffer(dtype, num_segments, segment_size): - order = parallel.SortOrder.DESCENDING + order = cuda.compute.SortOrder.DESCENDING num_items = num_segments * segment_size h_in_keys = random_array( @@ -231,10 +229,10 @@ def test_segmented_sort_pairs_double_buffer(dtype, num_segments, segment_size): d_tmp_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) d_tmp_vals = numba.cuda.to_device(np.empty_like(h_in_vals)) - keys_db = parallel.DoubleBuffer(d_in_keys, d_tmp_keys) - vals_db = parallel.DoubleBuffer(d_in_vals, d_tmp_vals) + keys_db = cuda.compute.DoubleBuffer(d_in_keys, d_tmp_keys) + vals_db = cuda.compute.DoubleBuffer(d_in_vals, d_tmp_vals) - parallel.segmented_sort( + cuda.compute.segmented_sort( keys_db, None, vals_db, @@ -258,7 +256,7 @@ def test_segmented_sort_pairs_double_buffer(dtype, num_segments, segment_size): @pytest.mark.parametrize("num_segments", [20, 600]) def test_segmented_sort_variable_segment_sizes(num_segments): - order = parallel.SortOrder.ASCENDING + order = cuda.compute.SortOrder.ASCENDING base_pattern = [ 1, 5, @@ -304,7 +302,7 @@ def test_segmented_sort_variable_segment_sizes(num_segments): d_out_keys = numba.cuda.to_device(np.empty_like(h_in_keys)) d_out_vals = numba.cuda.to_device(np.empty_like(h_in_vals)) - parallel.segmented_sort( + cuda.compute.segmented_sort( d_in_keys, d_out_keys, d_in_vals, From f7d50b97d86b6daaf4cf2713cf04c62626be1dfb Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 31 Oct 2025 23:30:53 +0000 Subject: [PATCH 090/100] Fix MSVC error --- c/parallel/test/test_segmented_sort.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 14bb5d948ab..7ef1dd36113 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -567,7 +567,8 @@ C2H_TEST("SegmentedSort works with custom types as values", "[segmented_sort][cu struct SegmentedSort_VariableSegments_Fixture_Tag; C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][variable_segments]", test_params_tuple) { - using T = c2h::get<0, TestType>; + using T = c2h::get<0, TestType>; + using key_t = typename T::KeyT; constexpr auto this_test_params = T(); constexpr bool is_descending = this_test_params.is_descending(); From e951835beb8e98b00e74c077df22848bd6f95f55 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 3 Nov 2025 16:41:03 +0000 Subject: [PATCH 091/100] Fix MSVC CI errors --- c/parallel/test/test_segmented_sort.cpp | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/c/parallel/test/test_segmented_sort.cpp b/c/parallel/test/test_segmented_sort.cpp index 7ef1dd36113..90b31ff88b1 100644 --- a/c/parallel/test/test_segmented_sort.cpp +++ b/c/parallel/test/test_segmented_sort.cpp @@ -197,7 +197,10 @@ C2H_TEST("segmented_sort can sort keys-only", "[segmented_sort][keys_only]", tes const std::size_t n_elems = n_segments * segment_size; std::vector host_keys_int = generate(n_elems); - std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); + std::vector host_keys(n_elems); + std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int val) { + return static_cast(val); + }); std::vector host_keys_out(n_elems); REQUIRE(host_keys.size() == n_elems); @@ -323,9 +326,15 @@ C2H_TEST("segmented_sort can sort key-value pairs", "[segmented_sort][key_value] const std::size_t n_elems = n_segments * segment_size; std::vector host_keys_int = generate(n_elems); - std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); + std::vector host_keys(n_elems); + std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int val) { + return static_cast(val); + }); std::vector host_values_int = generate(n_elems); - std::vector host_values(host_values_int.begin(), host_values_int.end()); + std::vector host_values(n_elems); + std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int val) { + return static_cast(val); + }); std::vector host_keys_out(n_elems); std::vector host_values_out(n_elems); @@ -593,11 +602,17 @@ C2H_TEST("SegmentedSort works with variable segment sizes", "[segmented_sort][va std::size_t n_elems = std::accumulate(segment_sizes.begin(), segment_sizes.end(), 0ULL); std::vector host_keys_int = generate(n_elems); - std::vector host_keys(host_keys_int.begin(), host_keys_int.end()); + std::vector host_keys(n_elems); + std::transform(host_keys_int.begin(), host_keys_int.end(), host_keys.begin(), [](int val) { + return static_cast(val); + }); // Generate float values by first generating ints and then transforming std::vector host_values_int = generate(n_elems); - std::vector host_values(host_values_int.begin(), host_values_int.end()); + std::vector host_values(n_elems); + std::transform(host_values_int.begin(), host_values_int.end(), host_values.begin(), [](int val) { + return static_cast(val); + }); std::vector host_keys_out(n_elems); std::vector host_values_out(n_elems); From 8fdb4fec520ab6c90bf245eaa55c4625a8f4a048 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 3 Nov 2025 23:06:49 +0000 Subject: [PATCH 092/100] Add comment explaining selector op compilation --- c/parallel/src/segmented_sort.cu | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index 3773d30c1c5..de68f0b78d8 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -590,6 +590,10 @@ CUresult cccl_device_segmented_sort_build_ex( constexpr size_t num_lto_args = 2; const char* lopts[num_lto_args] = {"-lto", arch.c_str()}; + // TODO: we currently compile each selector op separately from the main TU. + // We do this because we need to pass the selector ops to + // DispatchThreeWayPartition eventually. This causes increased compilation + // times, which might be avoidable. cccl_op_t large_selector_op = segmented_sort::make_segments_selector_op( 0, start_offset_it, From 5dedbd5ae461992b119c7460545f1319f5c77ed5 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 3 Nov 2025 23:29:05 +0000 Subject: [PATCH 093/100] Use dummy global variable instead of &op in indirect_arg_t constructor --- c/parallel/src/util/indirect_arg.h | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/c/parallel/src/util/indirect_arg.h b/c/parallel/src/util/indirect_arg.h index d74341a09f7..752ae3c719d 100644 --- a/c/parallel/src/util/indirect_arg.h +++ b/c/parallel/src/util/indirect_arg.h @@ -16,6 +16,20 @@ #include +// GCC emits a dangling pointer warning in the `LargeSegmentsSelector` and +// `SmallSegmentsSelector` functions in segmented_sort.cu. The warning occurs +// when `return indirect_arg_t()` invokes the `indirect_arg_t` constructor that +// accepts a `cccl_op_t`. Even though this is a stateless op, we must initialize +// the `ptr` member to a valid address. +// +// We cannot use `nullptr` because the pointer is passed to a driver API that +// requires the size of empty arguments to be 1 (not 0), meaning it will attempt +// to copy a byte from the address. +// +// Initially, we initialized `ptr` to `this`, but this triggered the dangling +// pointer warning in GCC. To avoid this, we use a global variable instead. +static inline char _global_storage = 0; + struct indirect_arg_t { void* ptr; @@ -25,7 +39,7 @@ struct indirect_arg_t {} indirect_arg_t(cccl_op_t& op) - : ptr(op.type == cccl_op_kind_t::CCCL_STATEFUL ? op.state : &op) + : ptr(op.type == cccl_op_kind_t::CCCL_STATEFUL ? op.state : &_global_storage) {} indirect_arg_t(cccl_value_t& val) From 169b5fc16d5f70ac8dd266c994baea23f923194a Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 4 Nov 2025 02:59:51 +0000 Subject: [PATCH 094/100] Revert change made to step counting iterator that was causing segmented reduce test to hang --- c/parallel/test/test_util.h | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/c/parallel/test/test_util.h b/c/parallel/test/test_util.h index 257faabe2ae..ed566e33837 100644 --- a/c/parallel/test/test_util.h +++ b/c/parallel/test/test_util.h @@ -1108,7 +1108,6 @@ inline std::tuple make_reverse_iterator_s return std::make_tuple(iterator_state_src, advance_fn_src, dereference_fn_src); } -// Common iterator helpers shared by segmented tests inline std::tuple make_step_counting_iterator_sources( std::string_view index_ty_name, std::string_view state_name, @@ -1135,9 +1134,9 @@ extern "C" __device__ void {0}({1}* state, {2} offset) std::format(it_def_src_tmpl, /*0*/ advance_fn_name, state_name, index_ty_name); static constexpr std::string_view it_deref_src_tmpl = R"XXX( -extern "C" __device__ {2} {0}({1}* state) +extern "C" __device__ void {0}({1}* state, {2}* result) {{ - return (state->linear_id) * (state->segment_size); + *result = (state->linear_id) * (state->segment_size); }} )XXX"; From 499e772cf22bc80385a5e99fe02e53afec1e2323 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 4 Nov 2025 05:42:15 +0000 Subject: [PATCH 095/100] Fix missing args in call --- .../tests/compute/examples/sort/segmented_sort_object.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py index cde29b8127e..50744aa642a 100644 --- a/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py +++ b/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py @@ -42,8 +42,8 @@ d_output_keys, d_input_vals, d_output_vals, - len(h_input_keys), - len(start_offsets), + h_input_keys.size, + start_offsets.size, cp.asarray(start_offsets), cp.asarray(end_offsets), ) @@ -56,7 +56,8 @@ d_output_keys, d_input_vals, d_output_vals, - len(h_input_keys), + h_input_keys.size, + start_offsets.size, cp.asarray(start_offsets), cp.asarray(end_offsets), ) From 27ba76515bf0726bfff1acfe37a709772a8ae713 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 4 Nov 2025 14:50:11 +0000 Subject: [PATCH 096/100] Use cccl_type_name_from_nvrtc to avoid windows errors --- c/parallel/src/segmented_sort.cu | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index de68f0b78d8..085894c6183 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -56,10 +56,10 @@ std::string get_device_segmented_sort_fallback_kernel_name( cccl_sort_order_t sort_order) { std::string chained_policy_t; - check(nvrtcGetTypeName(&chained_policy_t)); + check(cccl_type_name_from_nvrtc(&chained_policy_t)); std::string offset_t; - check(nvrtcGetTypeName(&offset_t)); + check(cccl_type_name_from_nvrtc(&offset_t)); /* template (&chained_policy_t)); + check(cccl_type_name_from_nvrtc(&chained_policy_t)); std::string offset_t; - check(nvrtcGetTypeName(&offset_t)); + check(cccl_type_name_from_nvrtc(&offset_t)); /* template (&chained_policy_t)); + check(cccl_type_name_from_nvrtc(&chained_policy_t)); std::string offset_t; - check(nvrtcGetTypeName(&offset_t)); + check(cccl_type_name_from_nvrtc(&offset_t)); /* template (); std::string offset_t; - check(nvrtcGetTypeName(&offset_t)); + check(cccl_type_name_from_nvrtc(&offset_t)); const std::string code = std::format( R"XXX( @@ -297,7 +297,7 @@ std::string get_three_way_partition_init_kernel_name() std::string get_three_way_partition_kernel_name(std::string_view large_selector_t, std::string_view small_selector_t) { std::string chained_policy_t; - check(nvrtcGetTypeName(&chained_policy_t)); + check(cccl_type_name_from_nvrtc(&chained_policy_t)); static constexpr std::string_view input_it_t = "thrust::counting_iterator"; @@ -308,7 +308,7 @@ std::string get_three_way_partition_kernel_name(std::string_view large_selector_ static constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*"; static constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT"; std::string offset_t; - check(nvrtcGetTypeName(&offset_t)); + check(cccl_type_name_from_nvrtc(&offset_t)); static constexpr std::string_view per_partition_offset_t = "cub::detail::three_way_partition::per_partition_offset_t"; static constexpr std::string_view streaming_context_t = From a588a5dab518cab410a4f04f166c3950ac2ab1f1 Mon Sep 17 00:00:00 2001 From: Michael Schellenberger Costa Date: Tue, 4 Nov 2025 16:39:00 +0100 Subject: [PATCH 097/100] Fix incorrect file name We had a merge conflict where a file was renamed and used in a new PR --- c/parallel/src/segmented_sort.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index de68f0b78d8..f1d451b8a96 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -12,7 +12,7 @@ #include // cub::detail::CudaDriverLauncherFactory #include #include // cub::DispatchSegmentedSort -#include // DeviceSegmentedSort kernels +#include // DeviceSegmentedSort kernels #include // policy_hub #include // cub::LoadModifier @@ -186,7 +186,7 @@ cccl_op_t make_segments_selector_op( const std::string code = std::format( R"XXX( -#include +#include extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* result_ptr) {{ @@ -638,7 +638,7 @@ CUresult cccl_device_segmented_sort_build_ex( const std::string final_src = std::format( R"XXX( -#include +#include #include #include #include From e8fcef4d3dd20f60a2b0831dbaf42476e4a460e6 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Tue, 4 Nov 2025 18:02:31 +0100 Subject: [PATCH 098/100] fix --- c/parallel/src/segmented_sort.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/c/parallel/src/segmented_sort.cu b/c/parallel/src/segmented_sort.cu index f1d451b8a96..48b0067fbf5 100644 --- a/c/parallel/src/segmented_sort.cu +++ b/c/parallel/src/segmented_sort.cu @@ -640,7 +640,7 @@ CUresult cccl_device_segmented_sort_build_ex( R"XXX( #include #include -#include +#include #include {0} From e83693126b49b1cd4341db5947c95636614e197f Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 4 Nov 2025 17:44:58 +0000 Subject: [PATCH 099/100] Address reviewer feedback --- .../cuda_cccl/cuda/compute/_bindings_impl.pyx | 7 +----- .../cuda/compute/algorithms/_sort/__init__.py | 23 +++++++++++++++++++ 2 files changed, 24 insertions(+), 6 deletions(-) create mode 100644 python/cuda_cccl/cuda/compute/algorithms/_sort/__init__.py diff --git a/python/cuda_cccl/cuda/compute/_bindings_impl.pyx b/python/cuda_cccl/cuda/compute/_bindings_impl.pyx index f6cccca4b2a..af8687d93e8 100644 --- a/python/cuda_cccl/cuda/compute/_bindings_impl.pyx +++ b/python/cuda_cccl/cuda/compute/_bindings_impl.pyx @@ -2289,12 +2289,7 @@ cdef extern from "cccl/c/segmented_sort.h": cccl_iterator_t d_keys_out, cccl_iterator_t begin_offset_in, cccl_iterator_t end_offset_in, - int, - int, - const char *, - const char *, - const char *, - const char * + int, int, const char *, const char *, const char *, const char * ) nogil cdef CUresult cccl_device_segmented_sort( diff --git a/python/cuda_cccl/cuda/compute/algorithms/_sort/__init__.py b/python/cuda_cccl/cuda/compute/algorithms/_sort/__init__.py new file mode 100644 index 00000000000..7f49aaebd51 --- /dev/null +++ b/python/cuda_cccl/cuda/compute/algorithms/_sort/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +from ._merge_sort import make_merge_sort as make_merge_sort +from ._merge_sort import merge_sort as merge_sort +from ._radix_sort import make_radix_sort as make_radix_sort +from ._radix_sort import radix_sort as radix_sort +from ._segmented_sort import make_segmented_sort as make_segmented_sort +from ._segmented_sort import segmented_sort as segmented_sort +from ._sort_common import DoubleBuffer, SortOrder + +__all__ = [ + "make_merge_sort", + "merge_sort", + "make_radix_sort", + "radix_sort", + "make_segmented_sort", + "segmented_sort", + "DoubleBuffer", + "SortOrder", +] From 0c03945c6deab06e3bee9eda541e47747257deb3 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 4 Nov 2025 17:48:12 +0000 Subject: [PATCH 100/100] Address reviewer feedback --- .../cuda_cccl/cuda/compute/algorithms/__init__.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/algorithms/__init__.py b/python/cuda_cccl/cuda/compute/algorithms/__init__.py index a4fe0b3c006..db746de0bf3 100644 --- a/python/cuda_cccl/cuda/compute/algorithms/__init__.py +++ b/python/cuda_cccl/cuda/compute/algorithms/__init__.py @@ -13,13 +13,13 @@ from ._scan import make_inclusive_scan as make_inclusive_scan from ._segmented_reduce import make_segmented_reduce as make_segmented_reduce from ._segmented_reduce import segmented_reduce -from ._sort._merge_sort import make_merge_sort as make_merge_sort -from ._sort._merge_sort import merge_sort as merge_sort -from ._sort._radix_sort import make_radix_sort as make_radix_sort -from ._sort._radix_sort import radix_sort as radix_sort -from ._sort._segmented_sort import make_segmented_sort as make_segmented_sort -from ._sort._segmented_sort import segmented_sort as segmented_sort -from ._sort._sort_common import DoubleBuffer, SortOrder +from ._sort import DoubleBuffer, SortOrder +from ._sort import make_merge_sort as make_merge_sort +from ._sort import make_radix_sort as make_radix_sort +from ._sort import make_segmented_sort as make_segmented_sort +from ._sort import merge_sort as merge_sort +from ._sort import radix_sort as radix_sort +from ._sort import segmented_sort as segmented_sort from ._three_way_partition import make_three_way_partition as make_three_way_partition from ._three_way_partition import three_way_partition as three_way_partition from ._transform import binary_transform, unary_transform