[FEA]: Design Thread Group Abstraction for Safe Cooperative Algorithm Interfaces
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
### Is this a duplicate?
- [x] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this request and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)
### 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:
```c++
template
__device__ void foo(Hierarchy h) {
int rank = h.rank(cudax::thread, cudax::block);
int size = h.count(cudax::thread, cudax::block);
}
template
__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:
```c++
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](https://nvidia.github.io/cccl/libcudacxx/extended_api/thread_groups.html) 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:
```c++
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-cooperative
```
### Describe 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_
Contributor guide
Assessment
This issue has not been assessed yet.