[FEA]: Follow up work to "Segmented TopK using Thread-Block-Clusters"
- 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
CUB
### Is your feature request related to a problem? Please describe.
There are things that we do not want to resolve in #9224 to get it merged ASAP.
### Describe the solution you'd like
- **Improve code quality with abstraction:**
- [ ] Use `BlockLoad` for non-bulk loads.
- [ ] Ring buffers for slots/stages/mbarriers?
- [ ] TMA abstraction (Work on `BlockLoadToShared` or a separate piepelining abstraction)
- **Docs:**
- [ ] Add API examples for determinism & tie-breaking requirements, now that clusters supports it.
- **Dispatch:**
- [ ] Logging
- [ ] Can we cache certain CUDA runtime queries in `static` variables?
- [ ] Add support for launching via driver API launcher factory.
- [ ] Fix `MaxPotentialDynamicSmemBytes()` and then use `launcher_factory.max_dynamic_smem_size_for`
- **Arg Fw:**
- [ ] Try to unify verification and traits in the argument annotation framework. Concretely, `segment_size_is_integral` and `segment_sizes_handle_ok` would come to mind.
- [ ] Should `bounded_offset_t` in `cub/detail/segmented_params.cuh` protect against negative upper bound instead of needing to catch this in each interface separately?
- [ ] Add a simple way to get element type from Arg Fw traits (currently value/element type depending on wrapper type)
- [ ] Make sure CUB primitives align in their Arg Fw usage.
- **Benchmarking/Tuning:**
- [ ] Make sure items-per-thread for scans is uneven to avoid smem bank conflicts.
- [ ] Add scan algo(s) to tuning policy.
- [ ] Make baseline tunable (currently its `policy_selector_from_types` is implemented in a way that assumes arch-independent tuning)
- [ ] Rename benchmark files for "baseline" backend to have the backend name in the file name once we have a proper backend name.
- [ ] Improve tuning for edge cases: #11358 and #11359
- **Testing:**
- [ ] Clean up tests (compile time and code size). See https://github.com/NVIDIA/cccl/pull/9224#pullrequestreview-4754059741
- **Feature:**
- [ ] Add a fallback through the cluster-based implementation (single block, but with streaming/ping-pong/determinism) for pre-Hopper.
- **Perf:**
- [ ] Have the histograms in a striped layout in smem so the load for the scan is cheaper (no bank conflicts). Can we do sth about the key loads for the selection scan as well?
- [ ] Try to improve the implementation of filter-pass-early exit by checking more rarely and reusing results from earlier atomics w/ 2x `__syncthreads_or()`.
- [ ] Use local `mbarrier`s instead of `cluster.barrier` for communication: #9994
- [ ] Try to use tensor memory to further expand the number of resident (not re-streamed) keys for large segments.
- [ ] Can we save on register pressure by not keeping `layout.block_keys_in = d_key_segments_it[segment_id]` around without regressing cases where that de-reference is expensive (needs new benchmark). Maybe even make the behavior depend on the type of iterator.
- [ ] Try using reusing atomic results and synchronizing them via `__syncthreads_or()` to avoid reads in stopping condition of filters.
- [ ] Try packing radix pass state (bucket id + early stop flag) into a single 32b DSMEM load/store.
- [ ] Re-assess 64-bit type usage https://github.com/NVIDIA/cccl/pull/9224#discussion_r3735663092
- [ ] Avoid using smem for edges when not using bulk copies https://github.com/NVIDIA/cccl/pull/9224#discussion_r3750044645
### Describe alternatives you've considered
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.