1、负载均衡概述：什么是GPU计算单元负载均衡、为什么需要负载均衡、负载不均的典型表现
大家好，我是你们这堂课的讲师。在GPU高性能计算这个领域摸爬滚打了十几年，我见过太多因为负载不均导致性能翻车的案例。今天咱们就来聊聊这个最基础、也最容易被忽视的话题——GPU计算单元的负载均衡。
说白了，负载均衡就是让GPU里所有的计算单元都忙起来，别让一部分累死，另一部分闲着。你想想看，一块GPU里有成百上千个计算核心，如果只有一半在工作，那性能不就打对折了吗？
什么是GPU计算单元负载均衡
先搞清楚一个概念。GPU里的计算单元，通常指的是SM（Streaming Multiprocessor）或者CU（Compute Unit）。每个SM里又包含多个CUDA核心、共享内存、寄存器等资源。
负载均衡，就是让这些SM之间的工作量尽可能平均。我习惯用一个比喻来解释：
想象一下码头卸货：你有10个工人（SM），每个工人一次能搬一箱货（线程块）。如果有的工人只搬了10箱，有的搬了100箱，那整个码头的工作效率就被最慢的工人拖累了。GPU也是这个道理。
在GPU编程中，负载均衡主要关注两个层面：
- 线程块级别的均衡：每个SM分配到的线程块数量要大致相当
- 线程级别的均衡：同一个线程块内的线程，工作量要尽量平均
嗯，这里要注意，线程块级别的均衡是硬件调度器自动做的，但线程级别的均衡，就得靠我们程序员自己把控了。
为什么需要负载均衡
这个问题其实很简单——为了性能。但我想从更深层次聊聊。
GPU的并行计算模型是SIMT（单指令多线程）。这意味着，同一个warp（32个线程）里的线程，必须执行同一条指令。如果某个线程提前完成了任务，它只能干等着其他线程跑完。
我在项目中遇到过这样一个场景：一个图像处理算法，需要对不同区域的像素做不同处理。边缘区域的像素计算量小，中心区域的计算量大。结果呢？整个warp的执行时间被最慢的那个线程拖长了。这就是典型的负载不均。
具体来说，负载均衡能带来三个好处：
- 提高硬件利用率：让所有SM都保持忙碌，避免资源浪费
- 减少执行时间：消除木桶效应，让整体执行时间不再被最慢的单元拖累
- 降低功耗：负载均衡后，GPU不需要频繁地在高负载和空闲状态之间切换，功耗更稳定
一个小技巧：我建议你在写kernel之前，先画一张数据分布图。看看你的数据是不是均匀的。如果数据本身就不均匀，那代码写得再好也没用。
负载不均的典型表现
怎么判断你的程序是不是负载不均？我总结了几个典型症状：
| 表现 | 现象 | 原因 |
|---|---|---|
| GPU利用率低 | nvidia-smi显示GPU利用率只有30%-50% | 部分SM空闲，等待其他SM完成 |
| warp发散严重 | 同一个warp内，线程执行路径不同 | 条件分支导致部分线程被屏蔽 |
| 内存访问不均衡 | 某些SM频繁访问全局内存，其他SM在等待 | 数据分布不均匀，导致访存模式不均衡 |
| 性能随数据变化剧烈 | 同样的代码，换一组数据性能差好几倍 | 数据分布影响负载均衡 |
我曾经接手过一个项目，客户说他们的程序在A100上跑，性能还不如V100。我一看nvidia-smi，GPU利用率只有25%。再一查代码，好家伙，一个for循环里，不同线程的迭代次数差了100倍。这就是典型的负载不均。
避坑指南：我曾经以为只要线程块数量足够多，硬件调度器就能自动做好负载均衡。后来发现，如果线程块的计算量差异太大，硬件调度器也无能为力。比如，有的线程块要处理1000个元素，有的只处理10个，那再怎么调度也没用。
再举一个更具体的例子。假设我们要对一个数组做归约求和：
// 负载不均的写法
__global__ void bad_reduce(int *data, int *sum, int n) {
int tid = threadIdx.x + blockIdx.x * blockDim.x;
if (tid < n) {
// 每个线程处理一个元素，但n可能不是blockDim.x的整数倍
atomicAdd(sum, data[tid]);
}
}
// 负载均衡的写法
__global__ void good_reduce(int *data, int *sum, int n) {
int tid = threadIdx.x + blockIdx.x * blockDim.x;
int stride = gridDim.x * blockDim.x;
int local_sum = 0;
for (int i = tid; i < n; i += stride) {
local_sum += data[i];
}
atomicAdd(sum, local_sum);
}
你看，第一种写法，如果n不是blockDim.x的整数倍，最后一个线程块就只有几个线程在工作。第二种写法，每个线程处理多个元素，而且通过stride的方式让所有线程的工作量更均衡。
为什么会这样？因为第一种写法里，线程块的数量和n强相关。n=1000，blockDim.x=256，那就需要4个线程块。前3个块满负荷，第4个块只有232个线程在工作。而第二种写法，每个线程都处理差不多数量的元素，负载自然就均衡了。
嗯，说到这里，我想强调一点：负载均衡不是银弹。有时候为了追求负载均衡，反而会增加额外的开销。比如上面那个例子，第二种写法虽然负载均衡了，但每个线程要处理多个元素，可能会增加寄存器压力。所以，具体怎么选，还是要看实际情况。
我个人习惯的做法是：先用profiler跑一遍，看看各个SM的占用率。如果发现某个SM的占用率明显低于其他SM，那就说明有负载不均的问题。然后再根据具体的数据分布和算法特点，选择合适的优化策略。
好了，这一章的内容就到这里。负载均衡是GPU性能优化的基石，理解了这个概念，后面的章节我们才能更好地讨论具体的优化技巧。下一章，我会聊聊如何通过线程块大小的选择来改善负载均衡。