线程束发散优化：从根源上解决性能杀手

说到GPU优化，线程束发散绝对是个绕不开的话题。我刚开始接触CUDA时，写过一段看似完美的代码——逻辑正确、算法清晰，结果一跑性能惨不忍睹。后来才发现，问题就出在线程束发散上。

说白了，线程束发散就是同一个warp里的线程走了不同的分支路径。GPU的SIMT架构要求warp内所有线程执行相同的指令，一旦出现分支，那些走不同路径的线程就得等待，这就造成了资源浪费。

线程束执行模型：32个线程绑在一起

先聊聊warp的基本概念。在NVIDIA GPU中，32个线程组成一个warp，这是硬件调度的基本单位。你可以把warp想象成一列火车，32个线程就是32节车厢。火车要么一起前进，要么一起停下。

每个warp里的线程共享程序计数器，这意味着它们在同一时刻执行同一条指令。我习惯把这个模型叫做「锁步执行」——步调一致，谁也别想单独行动。

分支发散的原因与影响

为什么会发散？看个简单例子：

__global__ void divergent_kernel(float* data, int N) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid < N) {
        if (data[tid] > 0.5f) {
            data[tid] = data[tid] * 2.0f;
        } else {
            data[tid] = data[tid] * 0.5f;
        }
    }
}

这段代码里，同一个warp内的线程可能有的走if分支，有的走else分支。硬件会怎么做？它会先执行所有走if的线程，让走else的线程等待；然后再执行else分支，让if的线程等待。你想想看，本来一次能搞定的事，现在要分两次执行。

我在项目中遇到过这样一个案例：一个图像处理核函数，根据像素值做不同处理。原本以为逻辑很清晰，结果profiling一看，warp执行效率只有40%左右。说白了，60%的周期都浪费在等待上了。

分支情况	warp执行次数	效率损失
所有线程走同一分支	1次	0%
两个分支各占一半	2次	50%
四个分支均匀分布	4次	75%

通过数据重排减少发散

既然发散是因为数据分布不均匀导致的，那我们就想办法让数据「物以类聚」。数据重排的核心思想是：把需要相同处理的数据集中到一起，让同一个warp内的线程处理相似的数据。

嗯，这里要注意，数据重排不是万能的。它需要额外的预处理开销，你得权衡一下是否值得。

方法一：基于条件的排序

最简单的做法就是先对数据排序，把满足条件的和不满足条件的分开。比如上面的例子，我们可以先根据data[tid] > 0.5f这个条件做一次排序或分区：

// 先做一次数据分区
// 把大于0.5的数据放到前面，小于等于0.5的放到后面
// 然后启动两个核函数分别处理

__global__ void process_high(float* data, int start, int end) {
    int tid = start + threadIdx.x + blockIdx.x * blockDim.x;
    if (tid < end) {
        data[tid] = data[tid] * 2.0f;
    }
}

__global__ void process_low(float* data, int start, int end) {
    int tid = start + threadIdx.x + blockIdx.x * blockDim.x;
    if (tid < end) {
        data[tid] = data[tid] * 0.5f;
    }
}

这样每个核函数内部就没有分支了，warp执行效率接近100%。代价是多了排序或分区的开销。

方法二：warp级别的手动控制

如果你对性能要求极高，可以尝试在warp级别手动控制。利用__ballot_sync和__shfl_sync等warp原语，让同一个warp内的线程协同工作：

__global__ void warp_optimized_kernel(float* data, int N) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= N) return;
    
    unsigned int mask = __ballot_sync(0xFFFFFFFF, data[tid] > 0.5f);
    // mask的每一位表示对应线程是否满足条件
    
    if (__popc(mask) == 32 || __popc(mask) == 0) {
        // 整个warp都走同一分支，直接处理
        if (data[tid] > 0.5f) {
            data[tid] = data[tid] * 2.0f;
        } else {
            data[tid] = data[tid] * 0.5f;
        }
    } else {
        // 有发散，需要特殊处理
        // 这里可以用shuffle指令做数据交换
        // 或者干脆走回普通分支
        if (data[tid] > 0.5f) {
            data[tid] = data[tid] * 2.0f;
        } else {
            data[tid] = data[tid] * 0.5f;
        }
    }
}

实战中的选择策略

在实际项目中，我一般按这个思路来决策：

• 分支简单且数据随机：直接接受发散，因为重排开销可能更大
• 分支复杂但数据有规律：用数据重排，把相似数据集中处理
• 分支嵌套深：考虑重构算法，用数学方式替代分支
• warp内大部分线程走同一条路：用warp原语做快速判断

我记得有一次做金融风险计算，数据分布极其不均匀——90%的数据走A路径，10%走B路径。这种情况下，我选择不做重排，因为那10%的线程造成的损失可以接受。但如果反过来，50%对50%，那就必须重排了。

最后说一句，线程束发散优化是个需要经验积累的活。别指望看一遍文档就能搞定，多跑profiling、多观察数据分布，慢慢你就会有感觉了。我个人觉得，与其追求理论上的完美，不如在实际项目中找到那个「够用就好」的平衡点。