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

[bernhardmgruber]

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:

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:

• [CUB] Refactor DeviceReduce::Arg{Min, Max} to always take an environment #9403
• [CUB] Refactor DeviceCopy to always take an environment #9416
• [CUB] Refactor DeviceAdjacentDifference::SubtractLeft to always take an environment #9418
• [CUB] Refactor DeviceAdjacentDifference::SubtractRight to always take an environment #9419
• [CUB] Refactor DeviceSelect::Flagged to always take an environment #9455
• [CUB] Refactor cub::DeviceSelect::If to always take an environment #9456
• [CUB] Refactor DeviceSelect::FlaggedIf to always take an environment #9457
• [CUB] Refactor DeviceSelect::Unique to always take an environment #9458
• [CUB] Refactor DeviceHistogram::MultiHistogramEven to always take an environment #9552
• [CUB] Refactor DeviceHistogram::HistogramEven to always take an environment #9553
• [CUB] Refactor DeviceHistogram::MultiHistogramRange to always take an environment #9554
• [CUB] Refactor DeviceHistogram::HistogramRange to always take an environment #9555
• Refactor device_adjacent_difference::dispatch to take a tuning #9454
• [CUB] Refactor dispatch_streaming_arg_reduce to take a tuning environment #9660
• [PSTL] Use env based overload for DeviceFindIf #9318
• Use environment based CUB APIs when possible in PSTL #8526
• [CUB] Refactor DeviceSelect::UniqueByKey to always take an environment #9459
• [CUB] Refactor DevicePartition::Flagged to always take an environment #9463
• [CUB] Refactor DevicePartition::If to always take an environment #9464
• [CUB] Refactor DeviceSelect::Unique to always take an environment #9386
• [CUB] Refactor DeviceSelect::UniqueByKey to always take an environment #9387
• [CUB] Refactor DeviceSelect::Flagged to always take an environment #9388
• [CUB] Refactor DeviceSelect::FlaggedIf to always take an environment #9389
• [CUB] Refactor DeviceSelect::If to always take an environment #9391
• [CUB] Refactor DevicePartition::If to always take an environment #9393
• [CUB] Refactor DevicePartition::Flagged to always take an environment #9396
• [CUB] Refactor DevicePartition::Flagged to always take an environment #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

[bernhardmgruber]

@miscco I just assigned this to you, since you already started, but anyone can work on this. If you already have a good prompt, consider posting it here.