NVIDIA / NVIDIA/cccl

[FEA]: Make `DeviceMergeSort` also consider `ValueT` for scaling down `ITEMS_PER_THREAD` in its policy

Open
#1,141 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

### Is this a duplicate?

- [X] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this request and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)

### Area

CUB

### Is your feature request related to a problem? Please describe.

In most of our tuning policies we scale down the `ITEMS_PER_THREAD` the larger the type is that the algorithm works on. The main motivation is to keep shared memory and/or register usage in our kernels somewhat constant despite varying data type sizes.

In `DeviceMergeSort`, we scale down `ITEMS_PER_THREAD` the larger `sizeof(KeyT)` is. However, when sorting pairs, we do not consider `ValueT` for reducing `ITEMS_PER_THREAD`.

As a result, when, for instance, using a 128-bit value type along with a 32-bit key type, we need to use virtual shared memory (in our current implementation) or use the fallback policy (once, https://github.com/NVIDIA/cccl/pull/1117 is merged).

```
AgentMergeSortPolicy<256,
Nominal4BItemsToItems(11),
cub::BLOCK_LOAD_WARP_TRANSPOSE,
cub::LOAD_LDG,
cub::BLOCK_STORE_WARP_TRANSPOSE>;
```

### Describe the solution you'd like

I think should pursue a similar approach to `DeviceRadixSort`, where we consider the larger of the two types (excerpt from our radix sort policy) for scaling down `ITEMS_PER_THREAD`. This would be reflective of the effective shared memory requirements in our kernels.

```
// Dominant-sized key/value type
using DominantT = cub::detail::conditional_t<(sizeof(ValueT) > sizeof(KeyT)), ValueT, KeyT>;

AgentMergeSortPolicy<256,
Nominal4BItemsToItems(11),
cub::BLOCK_LOAD_WARP_TRANSPOSE,
cub::LOAD_LDG,
cub::BLOCK_STORE_WARP_TRANSPOSE>;
```

### Describe alternatives you've considered

_No response_

### Additional context

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.