如何在单Warp中按谓词分区数据，使保留项内存连续？

针对单Warp（32个数据项）的筛选需求，最优实现是利用CUDA的Warp级内置指令——这类指令能在Warp内高效完成同步、掩码收集和索引计算，完全避免全局同步开销，是最快且代码简洁的方案。

核心步骤： #

• 生成保留掩码
  每个线程判断自身数据是否需要保留（比如判断是否为质数），生成1位掩码（保留为1，丢弃为0）。通过__ballot_sync(0xFFFFFFFF, keep_flag)收集整个Warp的掩码，得到一个32位整数，每一位对应一个线程的保留状态。

• 计算目标存储索引
  对需要保留数据的线程，用__popc(warp_mask & ((1 << threadIdx.x) - 1))计算当前线程之前（不含自身）的保留项数量，这个值就是当前数据在输出数组中的连续索引。

• 写入连续内存
  若当前线程的保留掩码为1，直接将数据写入输出数组的对应索引位置。

代码示例（以筛选质数为例）： #

__global__ void warpFilter(int* input, int* output) {
    const int tid = threadIdx.x;
    int val = input[tid];
    
    // 判断是否为质数（可替换为自定义筛选逻辑）
    bool is_prime = true;
    if (val <= 1) is_prime = false;
    else if (val == 2) is_prime = true;
    else if (val % 2 == 0) is_prime = false;
    else {
        for (int i = 3; i*i <= val; i += 2) {
            if (val % i == 0) {
                is_prime = false;
                break;
            }
        }
    }
    
    // 收集Warp内的保留掩码
    unsigned int warp_mask = __ballot_sync(0xFFFFFFFF, is_prime);
    
    // 计算当前保留数据的目标索引
    int dst_idx = __popc(warp_mask & ((1 << tid) - 1));
    
    // 写入连续内存
    if (is_prime) {
        output[dst_idx] = val;
    }
}

方案优势： #

• 无额外同步开销：Warp级指令__ballot_sync和__popc都是单周期指令，Warp内线程天然同步，无需__syncthreads()。
• 计算效率高：索引计算通过位操作完成，比传统前缀和方法更快。
• 代码简洁：核心筛选和索引映射逻辑仅需3行指令，无需复杂的共享内存操作。

内容的提问来源于stack exchange，提问作者Johan