Add support for large `num_items` to `device_segmented_radix_sort.cuh`

[elstehle]

• Improve DeviceSegmentedRadixSort test run time for large number of items and large number of segments #3245
• Fix temp storage computation when using non-double buffer invocation
• Fix test failures for large segments
• Add support for large number of segments
• Add benchmarks for DeviceSegmentedRadixSort

[Enigmatisms]

Hi, I wonder whether DeviceSegmentedRadixSort supports num_items greater than INT_MAX? I encountered failed cases when sorting a tensor with, like [228170138, 10] as shape and the temp_storage_bytes calculated is overflowing. I am using an older version, so the input param is still int, not std::cuda::int64_t (but I manually specified int64_t as OffsetT).

[elstehle]

That's not supported by DeviceSegmentedRadixSort. But you can use DeviceSegmentedSort.

DeviceSegmentedRadixSort in it's current implementation DeviceSegmentedRadixSort isn't well-suited for your use case, because it uses just one-thread block-per-segment, which performs suboptimally for very large segments. DeviceSegmentedSort does load-balancing.

[Enigmatisms]

Thank you! @elstehle