Repository navigation
[EPIC]: Env-based overloads implementation for all CUB device primitives #5635
Description
For every existing CUB device overload add an env-based counterpart that uses the new environment facilities from libcu++.
The goal is that we provide CUB device overloads that do not require explicit memory allocation which requires a 2-phase call. The existing CUB API looks like this:
void *d_temp_storage = nullptr;
size_t temp_storage_bytes = 0;
// First call: determine storage requirements
cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items, stream);
// Allocate storage
cudaMalloc(&d_temp_storage, temp_storage_bytes);
// Second call: actual algorithm execution
cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items, stream);Working PRs for reference: #4877 #4985
When the new APIs are added CUB should allow for the users to call the primitives like this:
auto env = cuda::execution::require(cuda::execution::determinism::run_to_run);
// Single call: allocation and execution in one call
cub::DeviceReduce::Sum(input.begin(), output.begin(), input.size(), env);The new environment argument should enable:
- access to
deterministicrequirements of the algorithm execution - flexibility on the
memory resourceselection - tuning without reaching for dispatch layer
- encapsulate the existing
streamselection option
List of primitives that have been provided with env-based overloads
Reacted by Elias Stehle, Srinivas Yadav and Bernhard Manfred Gruber
Activity
[-][EPIC]: Tracking env-based overloads implementation for all CUB primitives[/-][+][EPIC]: Track env-based overloads implementation for all CUB device primitives[/+]bdice commentedon Dec 11, 2025
@gonidelis asked for help prioritizing this. I did a quick analysis of usage in cuDF. Here are our priorities:
CUB Device Algorithm Usage in libcudf
Total: 97 usages across 23 distinct algorithms
16 remaining items
gonidelis commentedon Apr 29, 2026
Full API is covered by env overloads. Follow-up iterations, improvements and missing bits will be taking place as part of #8175