4. 线程束发散（Warp Divergence）：分支语句对Warp的影响、如何避免Warp发散、实战代码重构

好，咱们进入一个非常实战的话题——线程束发散。说实话，很多新手写CUDA代码，性能上不去，十有八九是栽在这个坑里。我自己刚入行那会儿，写了个并行归约，自我感觉良好，结果一跑profiling，发现SM利用率低得可怜。排查了半天，罪魁祸首就是warp发散。

4.1 什么是线程束发散？

先搞清楚基本概念。GPU执行指令的最小单位不是线程，而是线程束（Warp）。在NVIDIA的架构里，一个warp包含32个线程。这32个线程必须执行同一条指令。你想想看，如果这32个线程走的是不同的分支路径，会发生什么？

举个例子：

__global__ void divergent_kernel(float* data, int N) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        if (data[tid] > 0.0f) {
            data[tid] = sqrtf(data[tid]);  // 分支A
        } else {
            data[tid] = 0.0f;              // 分支B
        }
    }
}

这段代码看起来没问题吧？逻辑上完全正确。但在GPU上，它的执行效率可能很差。为什么？因为warp里的32个线程，有的走分支A，有的走分支B。硬件没办法同时执行两条不同的路径，只能先执行分支A，把走分支B的线程屏蔽掉；然后再执行分支B，把走分支A的线程屏蔽掉。

说白了，就是串行化了。一个warp本来应该一次搞定的事，现在要分两次执行。如果分支更多，性能损失就更严重。

4.2 分支语句对Warp的影响

我遇到过不少项目，代码里到处都是if-else，觉得反正线程多，分支无所谓。嗯，这里要注意，分支对warp的影响不是线性的，而是阶梯式的。

咱们来看一个更典型的场景——基于数据值的分支：

__global__ void process_by_value(float* input, float* output, int N) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        float val = input[tid];
        // 根据数值范围走不同分支
        if (val < 0.1f) {
            output[tid] = val * 2.0f;      // 分支1
        } else if (val < 0.5f) {
            output[tid] = val * 1.5f;      // 分支2
        } else {
            output[tid] = val * 0.5f;      // 分支3
        }
    }
}

这个kernel有三个分支。如果数据分布是随机的，那么每个warp里的32个线程可能分布在三个分支上。硬件会怎么做？它会依次执行分支1、分支2、分支3。每个分支执行时，只有部分线程真正在工作，其他线程都在“空转”。

我曾经在一个图像处理项目里做过测试：一个简单的三路分支，让整体性能下降了将近40%。你想想看，40%的算力就这么白白浪费了。

分支数量	理想执行次数	实际执行次数	性能损失
2路分支	1	2	~50%
3路分支	1	3	~67%
4路分支	1	4	~75%

注意，这个表格是理想情况下的估算。实际中，如果某个分支的线程数很少，性能损失可能没那么夸张。但无论如何，分支越多，损失越大。

4.3 如何避免Warp发散

好，问题摆在这了，怎么解决？我总结了几个实战中常用的方法。

4.3.1 数据重排——让同一个warp的线程走相同分支

这是最直接的方法。既然发散的原因是同一个warp里的线程走了不同分支，那我们把走相同分支的数据聚到一起不就行了？

举个例子，假设我们要处理一组数据，正数做开方，负数取绝对值：

// 发散版本
__global__ void divergent(float* data, int N) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        if (data[tid] > 0) {
            data[tid] = sqrtf(data[tid]);
        } else {
            data[tid] = fabsf(data[tid]);
        }
    }
}

// 优化版本——先排序再处理
__global__ void optimized(float* data, int* flags, int N) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        // flags数组已经标记了正负，且按正负排序
        if (flags[tid] == 1) {
            data[tid] = sqrtf(data[tid]);
        } else {
            data[tid] = fabsf(data[tid]);
        }
    }
}

优化版本的关键在于：我们在CPU端或者用一个单独的kernel，把数据按照分支条件排序。这样，同一个warp里的32个线程，要么全是正数，要么全是负数，就不会发散了。

4.3.2 用三元运算符替代if-else

有些简单的分支，可以用三元运算符（?:）替代。编译器对三元运算符的优化通常更好，有时候能生成无分支的指令序列。

// 不推荐
if (data[tid] > 0) {
    result[tid] = data[tid] * 2;
} else {
    result[tid] = data[tid] * 0.5;
}

// 推荐
result[tid] = (data[tid] > 0) ? (data[tid] * 2) : (data[tid] * 0.5);

不过要注意，三元运算符不是万能的。如果分支里的计算很复杂，编译器还是会生成分支指令。我个人习惯是：简单的赋值或算术运算用三元，复杂的逻辑还是用if-else，但配合数据重排。

4.3.3 将分支条件提到循环外部

这个技巧很实用。如果分支条件在整个循环中不变，那就提到外面去。

// 不好的写法
for (int i = 0; i < N; i++) {
    if (mode == 0) {
        process_mode0(data[i]);
    } else {
        process_mode1(data[i]);
    }
}

// 好的写法
if (mode == 0) {
    for (int i = 0; i < N; i++) {
        process_mode0(data[i]);
    }
} else {
    for (int i = 0; i < N; i++) {
        process_mode1(data[i]);
    }
}

这个优化在CPU上也有用，但在GPU上效果更明显。因为GPU的warp调度器最怕的就是频繁切换分支。

4.4 实战代码重构

说了这么多理论，咱们来一个完整的实战案例。假设我们要实现一个函数，对数组中的每个元素做如下处理：

• 如果元素值大于阈值THRESHOLD，做三次方运算
• 否则，做平方运算

这是原始版本：

__global__ void process_v1(float* data, float* output, int N, float THRESHOLD) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        if (data[tid] > THRESHOLD) {
            output[tid] = data[tid] * data[tid] * data[tid];
        } else {
            output[tid] = data[tid] * data[tid];
        }
    }
}

这个版本的问题很明显——warp发散。我们来做重构。

第一步：数据重排

我们先用一个辅助kernel，把数据按照是否大于阈值分成两组：

__global__ void classify(float* data, int* indices, int* count_high, int N, float THRESHOLD) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        if (data[tid] > THRESHOLD) {
            indices[tid] = 1;  // 标记为高值组
            atomicAdd(count_high, 1);
        } else {
            indices[tid] = 0;  // 标记为低值组
        }
    }
}

第二步：重排数据

根据indices数组，把数据重新排列，让高值组和低值组连续存放：

__global__ void reorder(float* data, float* reordered, int* indices, int* prefix_sum, int N) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        int new_pos;
        if (indices[tid] == 1) {
            // 高值组放在前面
            new_pos = prefix_sum[tid];
        } else {
            // 低值组放在高值组后面
            new_pos = count_high + (tid - prefix_sum[tid]);
        }
        reordered[new_pos] = data[tid];
    }
}

第三步：无发散的计算kernel

现在数据已经排好了，前count_high个元素是高值，后面都是低值。我们可以用两个独立的kernel来处理：

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

__global__ void process_low(float* data, float* output, int count_high, int N) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N - count_high) {
        output[count_high + tid] = data[count_high + tid] * data[count_high + tid];
    }
}

你看，这两个kernel里都没有分支了。每个warp里的32个线程执行完全相同的指令，效率拉满。

当然，这个重构方案也有代价——多了分类和重排的开销。我个人建议：如果分支逻辑简单，或者数据量不大，直接用三元运算符或者编译器优化就够了。但如果分支复杂、数据量大，数据重排绝对是值得的。

嗯，关于线程束发散，今天就聊到这。记住一句话：让同一个warp里的线程做同样的事。这是GPU优化的黄金法则之一。下一节咱们聊聊共享内存的bank conflict，那也是个大坑。