1. Reduce算子优化入门:从基础实现到性能翻倍
在GPU编程中,Reduce算子是最基础也是最常用的操作之一。简单来说,Reduce就是对数组中的元素进行归约计算,比如求和(sum)、求最大值(max)、求最小值(min)等。想象一下你有一长串数字,需要计算它们的总和——这就是Reduce操作的典型场景。
为什么Reduce在GPU上如此重要?因为GPU擅长并行计算,而Reduce操作虽然简单,但在深度学习、科学计算等领域无处不在。一个优化良好的Reduce算子可以显著提升整体性能。我刚开始接触CUDA时,第一个优化的算子就是Reduce,当时性能提升了近200%,这种成就感至今难忘。
我们先来看最基础的实现版本(Kernel 0):
__global__ void reduce_v0(float *g_idata, float *g_odata) { __shared__ float sdata[BLOCK_SIZE]; unsigned int tid = threadIdx.x; unsigned int i = blockIdx.x*blockDim.x + threadIdx.x; sdata[tid] = g_idata[i]; __syncthreads(); for(unsigned int s=1; s < blockDim.x; s *= 2) { if (tid % (2*s) == 0) { sdata[tid] += sdata[tid + s]; } __syncthreads(); } if (tid == 0) g_odata[blockIdx.x] = sdata[0]; }这个实现分为三个步骤:
- 每个线程加载一个数据到共享内存
- 在共享内存中进行树形归约
- 线程0将结果写回全局内存
但实测下来,这个基础版本在V100上的带宽利用率只有40.97%,性能瓶颈非常明显。问题主要出在两个方面:一是取模运算(%)性能很差,二是条件判断导致warp divergence(线程束分化)。
2. 性能优化第一战:解决warp divergence和bank conflict
2.1 消除warp divergence
在Kernel 1中,我们通过改变线程的工作方式来解决这两个问题:
__global__ void reduce_v1(float *g_idata, float *g_odata) { __shared__ float sdata[BLOCK_SIZE]; unsigned int tid = threadIdx.x; unsigned int i = blockIdx.x*blockDim.x + threadIdx.x; sdata[tid] = g_idata[i]; __syncthreads(); for(unsigned int s=1; s < blockDim.x; s *= 2) { int index = 2 * s * tid; if (index < blockDim.x) { sdata[index] += sdata[index + s]; } __syncthreads(); } if (tid == 0) g_odata[blockIdx.x] = sdata[0]; }这个改进的关键在于:不再是每个线程固定处理一个元素,而是让活跃线程处理更多工作。这样在早期迭代中,整个warp的线程要么都工作,要么都空闲,避免了warp divergence。实测性能从788.29us提升到502.43us,加速比达到1.56倍。
2.2 解决bank conflict
Kernel 1虽然解决了warp divergence,但引入了新的问题——bank conflict。在共享内存中,相邻地址位于不同的bank,当线程访问间隔为2*s时,可能导致多个线程访问同一个bank。这时候Kernel 2采用了"顺序寻址"的策略:
__global__ void reduce_v2(float *g_idata, float *g_odata) { __shared__ float sdata[BLOCK_SIZE]; unsigned int tid = threadIdx.x; unsigned int i = blockIdx.x*blockDim.x + threadIdx.x; sdata[tid] = g_idata[i]; __syncthreads(); for(unsigned int s=blockDim.x/2; s>0; s >>= 1) { if (tid < s) { sdata[tid] += sdata[tid + s]; } __syncthreads(); } if (tid == 0) g_odata[blockIdx.x] = sdata[0]; }这个版本让相邻线程访问相邻的内存地址,避免了bank conflict。性能进一步提升到375.90us,带宽利用率达到85.79%。从性能数据可以看出,解决bank conflict带来的提升比解决warp divergence更显著。
3. 高阶优化技巧:充分利用硬件资源
3.1 解决线程闲置问题
观察前面的kernel会发现一个问题:在归约过程中,有一半线程会逐渐闲置。Kernel 3通过让每个线程处理更多数据来解决这个问题:
__global__ void reduce_v3(float *g_idata, float *g_odata) { __shared__ float sdata[BLOCK_SIZE]; unsigned int tid = threadIdx.x; unsigned int i = blockIdx.x*(blockDim.x*2) + threadIdx.x; sdata[tid] = g_idata[i] + g_idata[i + blockDim.x]; __syncthreads(); for(unsigned int s=blockDim.x/2; s>0; s >>= 1) { if (tid < s) { sdata[tid] += sdata[tid + s]; } __syncthreads(); } if (tid == 0) g_odata[blockIdx.x] = sdata[0]; }这个改动让每个线程在加载数据时就进行一次加法操作,减少了闲置线程的数量。性能直接翻倍,达到205.89us,带宽利用率提升到81.72%。
3.2 展开最后一个warp
当活跃线程数小于等于32时(即1个warp),我们可以省略同步操作,因为warp内的线程是天然同步的。Kernel 4实现了这个优化:
__device__ void warpReduce(volatile float* cache, unsigned int tid) { cache[tid] += cache[tid+32]; cache[tid] += cache[tid+16]; cache[tid] += cache[tid+8]; cache[tid] += cache[tid+4]; cache[tid] += cache[tid+2]; cache[tid] += cache[tid+1]; } __global__ void reduce_v4(float *g_idata, float *g_odata) { __shared__ float sdata[BLOCK_SIZE]; unsigned int tid = threadIdx.x; unsigned int i = blockIdx.x*(blockDim.x*2) + threadIdx.x; sdata[tid] = g_idata[i] + g_idata[i + blockDim.x]; __syncthreads(); for(unsigned int s=blockDim.x/2; s>32; s >>= 1) { if (tid < s) { sdata[tid] += sdata[tid + s]; } __syncthreads(); } if (tid < 32) warpReduce(sdata, tid); if (tid == 0) g_odata[blockIdx.x] = sdata[0]; }这个版本性能提升到176.86us。需要注意的是,对于计算能力7.0以上的GPU(如V100),需要使用__syncwarp()来保证正确性,这就是Kernel 4.1的改进。
4. 终极优化:完全展开与向量化访存
4.1 完全展开循环
Kernel 5通过模板参数将循环完全展开,让编译器生成更优化的指令:
template <unsigned int blockSize> __device__ void warpReduce(volatile float* cache, int tid) { if(blockSize >= 64) cache[tid] += cache[tid+32]; if(blockSize >= 32) cache[tid] += cache[tid+16]; if(blockSize >= 16) cache[tid] += cache[tid+8]; if(blockSize >= 8) cache[tid] += cache[tid+4]; if(blockSize >= 4) cache[tid] += cache[tid+2]; if(blockSize >= 2) cache[tid] += cache[tid+1]; } template <unsigned blockSize> __global__ void reduce_v5(float *g_idata, float *g_odata) { __shared__ float sdata[BLOCK_SIZE]; unsigned int tid = threadIdx.x; unsigned int i = blockIdx.x*(blockDim.x*2) + threadIdx.x; sdata[tid] = g_idata[i] + g_idata[i + blockDim.x]; __syncthreads(); if (blockSize >= 1024) { if (tid < 512) sdata[tid] += sdata[tid+512]; __syncthreads(); } if (blockSize >= 512) { if (tid < 256) sdata[tid] += sdata[tid+256]; __syncthreads(); } if (blockSize >= 256) { if(tid < 128) sdata[tid] += sdata[tid+128]; __syncthreads(); } if (blockSize >= 128) { if (tid < 64) sdata[tid] += sdata[tid+64]; __syncthreads(); } if (tid < 32) warpReduce<blockSize>(sdata, tid); if (tid == 0) g_odata[blockIdx.x] = sdata[0]; }4.2 向量化访存
最后的性能杀手锏是向量化访存,一次性读取多个数据:
template <typename T, int pack_size> struct alignas(sizeof(T) * pack_size) Packed { T elem[pack_size]; __device__ void operator+=(Packed<T, pack_size> packA) { #pragma unroll for (int i = 0; i < pack_size; i++) { elem[i] += packA.elem[i]; } } }; __global__ void reduce_v8(float *g_idata, float *g_odata, unsigned int n) { __shared__ float warpLevelSums[kWarpSize]; unsigned int i = blockDim.x * blockIdx.x + threadIdx.x; Packed<float, 4> sum_pack(0.0); const auto *pack_ptr = reinterpret_cast<const Packed<float, 4>*>(g_idata); for (int linear_index = i; linear_index < n/4; linear_index += blockDim.x*gridDim.x) { sum_pack += pack_ptr[linear_index]; } float sum = sum_pack.elem[0] + sum_pack.elem[1] + sum_pack.elem[2] + sum_pack.elem[3]; // ... 后续reduce操作 }经过所有这些优化,最终性能从最初的788.29us提升到162.21us,加速比达到4.86倍,带宽利用率提升到34.3%。虽然看起来带宽利用率不高,但这已经接近Reduce这种低计算强度算子的理论极限了。