NVIDIA / NVIDIA/cccl

Add implementation that supports large segments in `cub::DeviceSegmentedTopK`

Open
#8,363 1 comment 0 reactions 1 assignee Claimed by @elstehle View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.