NVIDIA / NVIDIA/cccl

BlockRadixSort needs overloads that take the problem size and correctly sets the padding value for unused inputs

Open
#869 15 comments 0 reactions 0 assignees View on GitHub
cub
Dominant language
C++
Stars
2.5k
Forks
487
Avg merge
2d 7h
Merged PRs (30d)
296

Description

I am trying `cub::BlockRadixSort` with PyTorch, it is getting good performance, but I find it is hard to use:

For example, if I want to sort 1023 elements, then I would use 256 threads and 4 elements for each threads to get a total of 1024 elements. So I have to add one padding element.

Assuming that `cub::BlockRadixSort` is stable (please confirm and document it !!!), if the key type is `int`, then I can just pad with `INT_MAX` and after sort, the padding value will remain in the last position and I can just discard it.

But if the key type is `float`, it becomes more complicated. I have to pad with `nan` in order to make the padding value to remain at the end, because after bit transformation, the largest number is `nan`. According to IEEE 754, there are multiple different `nan`s, and I need to make sure that my padding `nan` is no smaller than the `nan`s in user's input. Doing so requires me to assume on the implementation detail of how cub transform bits, and it might not be forward compatible if cub changes how bits are transformed.

A PyTorch program to demonstrate this problem is:
```python
>>> a = torch.full((4096,), math.nan, device='cuda')
>>> a.sort()
torch.return_types.sort(
values=tensor([nan, nan, nan, ..., nan, nan, nan], device='cuda:0'),
indices=tensor([ 0, 1, 2, ..., 4093, 4094, 4095], device='cuda:0'))
>>> a.view(torch.int) # reinterpret bits as int
tensor([2143289344, 2143289344, 2143289344, ..., 2143289344, 2143289344,
2143289344], device='cuda:0', dtype=torch.int32)
>>> a.view(torch.int)[0].fill_(2147483647)
tensor(2147483647, device='cuda:0', dtype=torch.int32)
>>> a.view(torch.int)
tensor([2147483647, 2143289344, 2143289344, ..., 2143289344, 2143289344,
2143289344], device='cuda:0', dtype=torch.int32)
>>> a.sort()
torch.return_types.sort(
values=tensor([nan, nan, nan, ..., nan, nan, nan], device='cuda:0'),
indices=tensor([ 1, 2, 3, ..., 4094, 4095, 0], device='cuda:0'))
```
Both `2143289344` and `2147483647` are `nan`s when reinterpreting bits as float, but `2147483647` is a larger nan.

To make it easier to find the correct padding values, I suggest the following change:

We can either:
1. Add an API that provides the largest and lowest padding values for radix sort and radix sort descending. Something like `cub::limits::largest_radix()`. Or,
2. Modify the behavior of cub sorting to make sure that it is stable with respect to different `nan`s, that is: `sort([2143289344, 2147483647]) == [2143289344, 2147483647] != sort([2147483647, 2143289344]) == [2147483647, 2143289344]`

Contributor guide

Open the contributing guide

Research direction

Start with the cub::BlockRadixSort API and the radix ordering behavior described in the issue, especially stability and NaN handling. The issue names no files or tests; done requires a settled approach for problem-size-aware padding or radix padding limits, with the resulting behavior documented and validated.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
hpc
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.