[EPIC]: CUB large input support

[jrhemstad]

As a lower-level interface, CUB should optimize for flexibility and performance. As a result, CUB will not guarantee a large input will work by default. However, it should enable users to specify their desired offset type.

This means CUB should not perform any dynamic dispatch based on the input size. Instead, users should have a way to statically specify the offset type. In previous discussion we favored making the type of num_items a template and infer the offset type from the type of num_items.

Design-related research

• Identify CUB algorithms that do/do not support large inputs #1408
• Determine and finalize design for large input support in CUB #1454

Testing large number of items

• Add tests for large number of items & large number of segments to DeviceSegmentedRadixSort #2139
• Add tests for large number of items & large number of segments to DeviceSegmentedSort #2140
• Add tests for large number of items in DeviceScan::*ByKey

Enable large num_items in CUB algorithms that are sensitive to the choice of offset_t

• Add support for large num_items to device_scan.cuh #2062
• Add support for large num_items to DeviceScan::*ByKey #2458
• Add support for large num_items to device_select.cuh #1422
• Add support for large num_items to device_partition.cuh #1437
• Add support for large num_items to DevicePartition::ThreeWayPartition #2442
• Add support for large num_items to DeviceReduce::{ArgMin,ArgMax} #2515
• Add support for large num_items to cub::DeviceTransform #2714
• Add support for large number of items and large number of segments to device_segmented_sort.cuh #3132
• Add support for large num_items to device_merge.cuh #3134
• Add support for large num_items to device_segmented_radix_sort.cuh #3133
• Add support for large num_items to DeviceSegmentedReduce::{ArgMin,ArgMax} #2522
• Add support for large num_items to DeviceRunLengthEncode::Encode #2521
• Add support for large num_items to DeviceRunLengthEncode::NonTrivialRuns #2520
• Add support for large number of segments to DeviceSegmentedReduce #3242
• Add support for large num_items to device_spmv.cuh

Limit the number of kernel template instantiations by reducing the set of offset types

• Reduce the number of offset types for DeviceMergeSort #3312

Clean up interim testing infrastructure

• Switch tests for large number of items to use Device* interface in DeviceSelect
• Switch tests for large number of items to use Device* interface in DeviceScan

Documentation

• We want to be explicit about supported and tested offset types in the Dispatch interface (see Port thrust::merge[_by_key] to CUB #1817 (comment))

[elstehle]

legend for offset type: ✅ considered done | 🟡 considered lower priority | 🟠 considered higher priority, as it prevents usage for larger-than-INT_MAX number of items | ⏳ in progress

legend for testing columns: ✅ considered done | 🟡 to be done | 🟠 needs to support wider offset types first

algorithm	offset type	tests larger-thanINT_MAX	tests close to [U]INT_MAX
device_adjacent_difference.cuh	✅ choose_offset_t	✅ 2^33, sanity check, iterators	🟡
device_radix_sort.cuh	✅ choose_offset_t	✅ extensive check	✅ extensive check
device_reduce.cuh	✅ choose_offset_t: Reduce, Sum, Min, Max, ReduceByKey, TransformReduce ✅ streaming: ArgMin, ArgMax (#2647) ⚠️ (note: new interface only)	✅ sanity, 2^{30,31,33)	✅ sanity, 2^32-1
device_scan.cuh	✅ choose_offset_t: DeviceScan ✅ choose_offset_t: DeviceScanByKey	✅	✅
device_select.cuh	✅ choose_offset_t: UniqueByKey ✅ streaming (for any user-provided offset type): Flagged, If, Unique	✅	✅
device_partition.cuh	✅ streaming (for larger-than-int user-provided offset types): Flagged, If ✅ streaming (for larger-than-int user-provided offset types): ThreeWayPartition #2506	✅ ✅	✅ ✅
device_copy.cuh	🟡num_ranges: uint32_t 🟡buffer sizes: iterator_traits<SizeIteratorT>::value_type	🟡	🟡
device_memcpy.cuh	🟡 num_ranges: uint32_t 🟡 buffer sizes: iterator_traits<SizeIteratorT>::value_type	🟡	🟡
device_for.cuh	🟡NumItemsT: ForEachN, ForEachCopyN, Bulk difference_type: ForEach, ForEachCopy	🟡	🟡
device_histogram.cuh	🟡 dynamic dispatch: int for (num_rows * row_stride_bytes)<INT_MAX; OffsetT otherwise	🟡	🟡
device_merge_sort.cuh	🟡 NumItemsT ⏳choose_offset_t (#3328)	✅ extensive check	✅ extensive check
device_segmented_reduce.cuh	🟡 common_iterator_value_t({Begin,End}OffsetIteratorT): Reduce, Sum, Min, Max ⚠️ (note) int: ArgMin, ArgMax num_segments: int	✅ sanity, rnd [2^31; 2^33]	🟡
device_run_length_encode.cuh	🟠 int	🟠	🟠
device_segmented_sort.cuh	🟠 num_items & num_segments:int	🟠	🟠
device_segmented_radix_sort.cuh	🟠 num_items & num_segments:int	🟠	🟠
device_merge.cuh	🟠 int	🟡	🟡
device_transform.cuh	✅ choose_signed_offset_t (#3172)	✅	✅
device_spmv.cuh	🟠 int (not planned)	🟠	🟠