下面用 bit 二分,也就是 radix=2 的特例,说明 CUDA top-k selection 的基本思想。
目标是找出 top-k largest。算法分两步:
- 先找第 k 大元素的阈值
kth_value。 - 再收集所有
> kth_value的元素,并用== kth_value的元素补满 k 个。
这里为了突出算法本身,假设输入已经是 unsigned sortable key。对于 uint32 数据可以直接用;对于 float,真实 CUDA 实现会先把 float 转成保持排序关系的 unsigned bit key。