NVIDIA / NVIDIA/cccl

Refactor CUB device-scope APIs to take an environment instead of `cudaStream_t`

Open
#9,875 4 comments 0 reactions 1 assignee Claimed by @miscco View on GitHub
help wanted
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

We decided a while back on Slack that in order to solve the problem of querying the temporary storage size of CUB APIs that take an environment and have no equivalent non-env overload, we should just replace all `cudaStream_t` parameters in public device-scope APIs by a `const Env&` parameter and add the additional environment handling (excluding temporary storage allocation). This would result in every CUB API to have two flavors:

```c++
struct DeviceAlg {
// two-phase overload, user must handle temp storage allocation
template <..., typename Env>
static cudaError_t Algorithm(void* ts_ptr, size_t& ts_size, ..., const Env& env);

// single-phase overload, handles temp storage allocation itself
template <..., typename Env>
static cudaError_t Algorithm(..., const Env& env);
}
```
This change is transparent to users, since a passed `cudaStream_t` (or a type convertible to that) will be handled appropriately by the stream query of the environment.

This work has already begun, but was mostly driven by needs in PSTL to query the temporary storage to combine that storage with further allocations (like an output parameter). See also #9328 and the following PRs:

- #9403
- #9416
- #9418
- #9419
- #9455
- #9456
- #9457
- #9458
- #9552
- #9553
- #9554
- #9555
- #9454
- #9660
- #9318
- #8526
- #9459
- #9463
- #9464
- #9386
- #9387
- #9388
- #9389
- #9391
- #9393
- #9396
- #9462

However, we should complete this work. Here are the currently missing overloads that still have a `cudaStream_t` parameter:

| Struct | Public methods taking `cudaStream_t` |
|---|---|
| `DeviceFind` | `LowerBoundSortedValues`, `UpperBoundSortedValues` |
| `DeviceFor` | `Bulk`, `ForEach`, `ForEachN`, `ForEachCopy`, `ForEachCopyN`, `ForEachInExtents`, `ForEachInLayout` |
| `DeviceHistogram` | `MultiHistogramRange` *(deprecated raw-array overload only)* |
| `DeviceMemcpy` | `Batched` |
| `DeviceMerge` | `MergeKeys`, `MergePairs` |
| `DeviceMergeSort` | `SortKeys`, `SortKeysCopy`, `SortPairs`, `SortPairsCopy`, `StableSortKeys`, `StableSortKeysCopy`, `StableSortPairs`, `SortKeysNoNVTX`, `SortKeysCopyNoNVTX`, `SortPairsNoNVTX` |
| `DeviceRadixSort` | `SortKeys`, `SortKeysDescending`, `SortPairs`, `SortPairsDescending` |
| `DeviceReduce` | `Reduce`, `Sum`, `Min`, `Max`, `ArgMin`, `ArgMax`, `ReduceByKey`, `TransformReduce` |
| `DeviceRunLengthEncode` | `Encode`, `NonTrivialRuns` |
| `DeviceScan` | `ExclusiveScan`, `ExclusiveScanByKey`, `ExclusiveSum`, `ExclusiveSumByKey`, `InclusiveScan`, `InclusiveScanByKey`, `InclusiveScanInit`, `InclusiveSum`, `InclusiveSumByKey` |
| `DeviceSegmentedRadixSort` | `SortKeys`, `SortKeysDescending`, `SortPairs`, `SortPairsDescending` |
| `DeviceSegmentedReduce` | `Reduce`, `Sum`, `Min`, `Max`, `ArgMin`, `ArgMax` |
| `DeviceSegmentedScan` | `ExclusiveSegmentedScan`, `ExclusiveSegmentedSum`, `InclusiveSegmentedScan`, `InclusiveSegmentedScanInit`, `InclusiveSegmentedSum` |
| `DeviceSegmentedSort` | `SortKeys`, `SortKeysDescending`, `SortPairs`, `SortPairsDescending`, `StableSortKeys`, `StableSortKeysDescending`, `StableSortPairs`, `StableSortPairsDescending` |

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.