Refactor CUB device-scope APIs to take an environment instead of `cudaStream_t`
- 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
Assessment
This issue has not been assessed yet.