Add implementation that supports large segments in `cub::DeviceSegmentedTopK`
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
The goal is to have an implementation that can process a batch of segments, where each segment may require multiple thread blocks to collaborate to compute the result
**Tasks:**
- [x] Extract building blocks from `AgentTopK`
- [x] Put the building blocks back together to compare against `DeviceTopK`
- [x] Benchmark approach against `DeviceTopK`
- [x] Investigate performance regressions from refactoring
- [x] Mitigate register pressure
- [x] Introduce infrastructure for work assignment:
- [x] determining per segment parameters (counters, input size, k, ...)
- [x] determining per segment input and output iterators
- [x] Move utilities shared across single-problem and segmented to common header helpers
- [x] Add tests for variable-size large segments
- [x] Fix bugs
- [x] Prepare benchmarks that dispatch single-problem to the batched algorithm for performance comparison
- [x] Try to mitigate performance regressions / register pressure
- [ ] ⏳ Clean up and carve out stand-alone buildings blocks to pipeline for reviews:
- [ ] ⏳ First histogram kernel and agent
- [ ] ⏳ `TileDataSource` for direct data sources (ones not requiring smem/syncs)
- [ ] `BlockPartition`: atomic strategy, inlined classify, gather_one
- [ ] `BlockFilter`: atomic strategy, inlined classify, gather_one
- [ ] ⏳ Document approach: https://github.com/NVIDIA/cccl_private/issues/659
Improvements:
- [x] Skip back-inserter candidates once we filled up the k slots with the candidates.
- [ ] Put key-prefix member var of `IdentifyCandidatesOp` into on-chip memory
- [ ] Allow loading into striped arrangement
- [ ] Store back via TMA (questionable: expectedly low number of selected and candidates per block(?))
- [ ] Store twiddled keys to buffer and only twiddle-out when writing to output-iterator
- [ ] Load values part only for non-rejected items (instead of loading full values tiles ahead of time)
_Depends on:_
_No dependencies._
Contributor guide
Assessment
This issue has not been assessed yet.