BlockRadixSort needs overloads that take the problem size and correctly sets the padding value for unused inputs
- 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
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