Warp bitonic-sort based top-k
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 487
- Avg merge
- 2d 7h
- Merged PRs (30d)
- 296
Description
## Summary
- A warp-wide top-k built on a register-resident bitonic sorting network, exchanging data
through warp shuffles (no shared memory for the array path), intended as the warp-per-segment building
block for small-segment `DeviceBatchedTopK`. Using
- **Scope**: warp
- **Comparison method**: comparison-based (`CompareOp` with strict less-than semantics)
- **Status**: in review / draft: `WarpBitonicSort` (#8391) in review; `WarpBitonicTopK` (#9281) draft.
## Problem coverage
### Segment / input size
- **Supported range**: small (tens to few hundreds of items).
- **Expected sweet spot**: small segments. Per #8391, the warp bitonic sort beats warp merge sort for
`N <= 224` and degrades beyond that due to the network's complexity.
### k
- **Supported range**: `1 .. max_k`, `k <= num_items`. **`max_k` must be a multiple of the warp size
(32).** Per-lane register cost scales with `max_k / 32`.
- **Expected sweet spot**: small k (small `max_k/32`), to keep register pressure bounded.
### Types
- **Supported**: any type with a comparator that is also warp-shuffleable (native `__shfl` types take the
fast path, others go through `cuda::device::warp_shuffle_idx`). Keys-only and key-value.
- **Expected sweet spot**: 32-bit wide keys (and values).
- **Custom types**: just a comparator needed, type needs to be trivially copyable.
## Output guarantees
### Determinism (set membership)
| `determinism` | Supported | Notes |
|---|---|---|
| `not_guaranteed` | ✅ | Trivially (weaker than what the network provides). |
| `run_to_run` | ✅ | The compare-exchange schedule is fixed and data-independent; no atomics. |
| `gpu_to_gpu` | ⬜ | TBC. Potentially reachable by augmenting the sort key with the source index |
### Tie-break (only meaningful with `gpu_to_gpu`)
| `tie_break` | Supported | Notes |
|---|---|---|
| `unspecified` | ✅ | With keys only, the boundary tie is resolved by lane position (a fixed function of the network), which is deterministic but not defined in terms of source index. |
| `prefer_smaller_index` | ⬜ | TBC. Not guaranteed by a keys-only network. Reachable by augmenting the sort key with the source index (compare by (key, index)), then smaller-index ties fall out naturally. |
| `prefer_larger_index` | ⬜ | TBC. Not guaranteed by a keys-only network. Reachable by augmenting the sort key with the source index (compare by (key, index)), then smaller-index ties fall out naturally. |
### Output ordering (result sequence)
| `output_ordering` | Supported | Notes |
|---|---|---|
| `unsorted` | ✅ | Trivially. |
| `sorted` | ✅ | The network yields keys in sorted (striped) order. |
| `stable_sorted` | ⬜ | TBC. Bitonic sort is not stable. Only reachable by augmenting the sort key with the source index (which then also pins the tie-break). |
## Other capabilities
- **Values (key-value / pairs)**: ✅ keys-only and pairs overloads, values follow their keys through
the network.
- **Multi-key efficiency**: ✅ comparison-based, so composite/multi-key inputs are handled by supplying
an appropriate comparator. Extra cost may incur from non-native shuffles from wider types.
- **Architecture constraints**: none special.
- **Temporary / shared memory**: array overloads are register-resident (no smem); the iterator overload
uses a small warp-private buffer (`WARP_THREADS` keys + values).
- **Occupancy / register pressure**: TBC.
## Performance notes
- **Benchmarks**: `cub/benchmarks/bench/bitonic_sort/` (warp keys and pairs, latency + throughput).
- **Known cross-over points**: warp bitonic sort wins for N ≲ 224 vs warp merge sort (#8391); for larger
retained sets a radix/selection strategy is expected to be preferable.
## Known limitations
- `MAX_K` must be a multiple of the warp size (32).
- Not stable. A source-index tie-break or `stable_sorted` requires augmenting the sort key with the index.
- Sequential cost grows quickly with input length. Keep N and the retained set small.
- All lanes in the warp must participate and must agree on `num_items` (undefined otherwise).
Contributor guide
Assessment
This issue has not been assessed yet.