Repository navigation
[FEA]: Design Thread Group Abstraction for Safe Cooperative Algorithm Interfaces #6181
Description
Is this a duplicate?
- I confirmed there appear to be no duplicate issues for this request and that I agree to the Code of Conduct
Area
CUDA Experimental (cudax)
Is your feature request related to a problem? Please describe.
cudax provides abstractions to express thread hierarchies, creating a single source of truth for dimensions that flow from kernel launch through to algorithm implementations:
template <class Hierarchy>
__device__ void foo(Hierarchy h) {
int rank = h.rank(cudax::thread, cudax::block);
int size = h.count(cudax::thread, cudax::block);
}
template <class Config>
__global__ void kernel_taking_config(Config config) {
foo(config.dims);
}
cudax::launch(dst, config, kernel_instance);This is crucial for cooperative interfaces. Say, CUB block-level algorithms are often misused because users provide different block sizes at launch and cooperative algorithms instantiation. To address this, we'll design new cooperative interface based on free functions accepting a thread group description as the first parameter.
That said, hierarchies are not sufficient. Cooperative algorithms often have to synchronize. We could provide a free function cudax::sync(dims) to address this need, but this would be dangerous. If synchronization is hidden on the implementation side, there's no way of knowing on caller end if the interface is cooperative or not:
auto dims = config.dims;
foo(dims); // Does it sync internally? We can't tell!
bar(dims); // What about this one?Instead, we could wrap dimensions in a thread group abstraction, similar to this one. This way, if an interface requires synchronization - it's immediately visible on caller end and user knows that all threads have to invoke the function in order to avoid deadlock:
auto dims = config.dims;
foo(dims); // might be a query based on the threads number
auto block = dims.group_by(cudax::block);
bar(block); // syncs on block level, hense block-cooperative
auto warp = dims.group_by(cudax::warp);
baz(warp); // syncs on warp level, hense warp-cooperativeDescribe the solution you'd like
This issue can be closed with a cudax hierarchies-based thread group design that'd make it possible to implement next revision of cooperative algorithms. On the algorithm side, it should be possible to statically distinguish dynamic and static dimensions, query thread rank and group size, as well as synchronize threads within a group.
Describe alternatives you've considered
No response
Additional context
No response
Activity
reducefree function that takes new group #7580bernhardmgruber commentedon Feb 10, 2026
I suggest to also consider the current design of squads in the warpspeed library by @ahendriksen, which we quite successfully used to implement the new scan in #6811. A squad is a group of warps, thus a group of threads with a thread count that is a multiple of 32. Here is a small excerpt of the API:
This design facilities warp specialized kernels and the generation of uniform data paths (as commonly used in conjunction with TMA and as a reduction of register usage), which are essential tools for SOL kernel authoring on Hopper+.
davebayer commentedon Feb 10, 2026
@bernhardmgruber I'm trying to integrate the warpspeed into this design. I'd like to make the groups generic enough, so squad is just a specific group of warps.
I had in mind something like:
25 remaining items
cudax::level_synchronizer#9870cudax::group_view#10754cudax::coop::namespace #10755group's copy constructor #10821arrive/waitAPIs for groups #11421cudax::coop::namespace #10755