TopK选择算法

/**
 * TopK选择算法 - GPU实现
 * 
 * 方法: 分块选择 + 归并
 * 1. 每个Block找本Block的TopK
 * 2. 归并所有Block的TopK
 * 
 * 编译: nvcc -o topk topk.cu
 */

#include <stdio.h>
#include <cuda_runtime.h>
#include <thrust/sort.h>
#include <thrust/device_vector.h>

#define CUDA_CHECK(call) do { cudaError_t e = call; if(e != cudaSuccess) { \
    fprintf(stderr, "CUDA Error: %s\n", cudaGetErrorString(e)); exit(1); } } while(0)

// ==================== Block级TopK ====================
__global__ void block_topk(const float* input, float* block_topk_vals, 
                         int* block_topk_idxs, int n, int k) {
    extern __shared__ float smem[];
    float* s_vals = smem;
    int* s_idxs = (int*)(s_vals + blockDim.x);
    
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + tid;
    
    // 加载数据
    if (i < n) {
        s_vals[tid] = input[i];
        s_idxs[tid] = i;
    } else {
        s_vals[tid] = -INFINITY;
        s_idxs[tid] = -1;
    }
    __syncthreads();
    
    // 简单选择排序找TopK（仅前K个）
    for (int j = 0; j < k && j < blockDim.x; j++) {
        // 找最大值位置
        int max_idx = tid;
        for (int s = blockDim.x / 2; s > 0; s >>= 1) {
            if (tid < s && tid + s < blockDim.x) {
                if (s_vals[tid + s] > s_vals[max_idx]) {
                    max_idx = tid + s;
                }
            }
            __syncthreads();
        }
        
        // 只有线程0记录并交换
        if (tid == 0 && j < k) {
            block_topk_vals[blockIdx.x * k + j] = s_vals[0];
            block_topk_idxs[blockIdx.x * k + j] = s_idxs[0];
            s_vals[0] = -INFINITY;
        }
        __syncthreads();
    }
}

// ==================== 全局TopK（使用thrust） ====================
void topk_thrust(const float* d_input, float* d_topk_vals, int* d_topk_idxs, int n, int k) {
    // 创建索引数组
    thrust::device_vector<int> indices(n);
    thrust::device_vector<float> values(n);
    
    // 复制输入
    thrust::copy(thrust::device_pointer_cast(d_input), 
                 thrust::device_pointer_cast(d_input) + n, 
                 values.begin());
    
    // 生成索引
    thrust::sequence(indices.begin(), indices.end());
    
    // 按值排序（降序）
    thrust::sort_by_key(values.begin(), values.end(), indices.begin(), 
                        thrust::greater<float>());
    
    // 取前K个
    thrust::copy(values.begin(), values.begin() + k, 
                 thrust::device_pointer_cast(d_topk_vals));
    thrust::copy(indices.begin(), indices.begin() + k, 
                 thrust::device_pointer_cast(d_topk_idxs));
}

// ==================== 主函数 ====================
int main() {
    const int N = 1024;
    const int K = 10;
    
    printf("TopK选择: N=%d, K=%d\n", N, K);
    
    // 分配内存
    float* h_input = (float*)malloc(N * sizeof(float));
    float* h_topk = (float*)malloc(K * sizeof(float));
    int* h_idxs = (int*)malloc(K * sizeof(int));
    
    // 初始化（随机值）
    for (int i = 0; i < N; i++) {
        h_input[i] = (float)(rand() % 1000) / 1000.0f;
    }
    
    // GPU
    float *d_input, *d_topk;
    int *d_idxs;
    CUDA_CHECK(cudaMalloc(&d_input, N * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_topk, K * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_idxs, K * sizeof(int)));
    
    CUDA_CHECK(cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));
    
    // 使用thrust实现
    topk_thrust(d_input, d_topk, d_idxs, N, K);
    
    CUDA_CHECK(cudaMemcpy(h_topk, d_topk, K * sizeof(float), cudaMemcpyDeviceToHost));
    CUDA_CHECK(cudaMemcpy(h_idxs, d_idxs, K * sizeof(int), cudaMemcpyDeviceToHost));
    
    // 打印结果
    printf("Top %d values:\n", K);
    for (int i = 0; i < K; i++) {
        printf("  [%d] index=%d, value=%.4f\n", i, h_idxs[i], h_topk[i]);
    }
    
    // 清理
    cudaFree(d_input); cudaFree(d_topk); cudaFree(d_idxs);
    free(h_input); free(h_topk); free(h_idxs);
    
    return 0;
}