Top K Selection

原始题目：LeetGPU - Top K Selection

题目描述

编写一个 GPU 程序，给定一个长度为 $N$ 的 32 位浮点数一维数组 input，选出最大的 $k$ 个元素，按降序写入长度为 $k$ 的输出数组。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 数组中。

示例

示例 1

Input:  input = [1.0, 5.0, 3.0, 2.0, 4.0], N = 5, k = 3
Output: [5.0, 4.0, 3.0]

示例 2

Input:  input = [7.2, -1.0, 3.3, 8.8, 2.2], N = 5, k = 2
Output: [8.8, 7.2]

约束条件

• $1 \le N \le 100{,}000{,}000$，$1 \le k \le N$。
• 所有值为 32 位浮点数。
• 性能测试在 $N = 50{,}000{,}000,\ k = 100$ 的规模下进行。

解题思路

Top-K 是 GPU 上的经典问题。当 $k$ 很小时（如 100），不需要完整排序——使用基于堆的选择或基数选择（Radix Select）更高效。典型做法是：每个 block 用共享内存维护一个大小为 $k$ 的最小堆，扫描自己负责的数据段，最终合并各 block 的 top-k。也可以先用桶排序思想按最高几位分桶，只在候选桶中继续筛选。注意题目要求结果降序，最后一步需要对选出的 $k$ 个元素排序。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void topk_kernel(const float* input, float* output, int N, int k) {
    __shared__ float heap[128]; __shared__ int heap_size;
    if(threadIdx.x==0)heap_size=0; __syncthreads();
    for(int i=blockIdx.x*blockDim.x+threadIdx.x; i<N; i+=gridDim.x*blockDim.x) {
        float val=input[i]; int pos=atomicAdd(&heap_size,1);
        if(pos<k) heap[pos]=val;
        else { heap_size=k; int mi=0; for(int j=1;j<k;j++)if(heap[j]<heap[mi])mi=j; if(val>heap[mi])heap[mi]=val; }
    }
    __syncthreads();
    if(threadIdx.x<k) {
        for(int i=0;i<k;i++) for(int j=i+1;j<k;j++) if(heap[j]>heap[i]){float t=heap[i];heap[i]=heap[j];heap[j]=t;}
        output[threadIdx.x]=heap[threadIdx.x];
    }
}
extern "C" void solve(const float* input, float* output, int N, int k) {
    topk_kernel<<<min((N+255)/256,1024),256>>>(input,output,N,k);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def topk_kernel(input_ptr, output_ptr, N: tl.constexpr, k: tl.constexpr, BLOCK: tl.constexpr):
    idx=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
    mask=idx<N
    x=tl.load(input_ptr+idx,mask=mask,other=-float('inf'))
    sorted,_=tl.sort(x,descending=True)
    tl.store(output_ptr+tl.arange(0,k), tl.load(sorted.to(input_ptr.dtype),mask=tl.arange(0,k)<k))