【项目复盘】CUDA编程——Reduce sum cuda核函数优化

【项目复盘】CUDA编程——Reduce sum cuda核函数优化

在高性能计算和深度学习领域，CUDA编程已经成为不可或缺的技能。其中，reduce sum操作作为常见的并行计算模式，其优化技巧更是值得深入研究。本文将从基础实现开始，逐步介绍如何通过多种优化手段提升reduce sum操作的性能。

reduce baseline实现

CUDA核函数的实现基本流程如下：

1. 分配内存

• 使用malloc在CPU上申请内存
• 使用cudaMalloc在GPU上申请显存
• 使用cudaMemcpy在CPU和GPU之间拷贝数据

2. 分配block和thread数量

• 通过dim3定义Grid和Block
• dim3有三个成员变量x、y、z，用于定位三维的size
• gridDim.x、gridDim.y、gridDim.z表示x、y、z方向上的block数量
• blockDim.x、blockDim.y、blockDim.z表示x、y、z方向上的thread数量

3. 定义CUDA核函数

• 使用__global__关键字表示核函数，由CPU调用但在GPU上执行
• __host__表示由CPU调用且在CPU上执行的函数
• __device__表示在GPU上执行的函数

reduce_v0普通并行版本

在设计核函数时，需要遵循以下规则：

1. 按照Block设计算法，但按Thread编写程序
2. 为每个block设计更进一步的索引
3. 绘制清晰的block算法图

核函数内部需要两个索引：

• tid：表示线程在block内的ID，即threadIdx.x
• gtid：表示线程在整个block范围内的全局ID，即blockIdx.x * blockDim.x + threadIdx.x

在v0版本中，gtid表示为blockIdx.x * blockSize + threadIdx.x，其中blockSize是模板参数，用于静态shared memory的申请。

在核函数中，使用__shared__关键字分配共享内存。需要注意的是，在对共享内存进行读写操作时，应该使用__syncthreads()确保所有线程都完成赋值后再进行后续计算。

上图展示了reduce_v0 block内部的算法图。两两相加后，block内所有值的sum最终会集中在第0个线程上。但需要注意，使用共享内存仅在block层面有效，如果需要在整个block范围内进行reduce，还需要对所有block的0号元素进行额外的reduce操作。

reduce_v1消除warp divergence版本

warp divergence是指在一个warp中，不同线程执行不同的代码路径，导致部分线程等待其他线程完成的情况。例如，当只有if语句而没有else语句时，可能会出现这种情况。

从上图可以看出，在第一轮迭代中，工作的线程（红色）是间隔分布的，这会导致warp内部出现工作和等待两种状态。为了避免这种情况，可以调整代码，使工作的线程集中在一起。

for(int index = 1; index < blockDim.x; index *= 2) {
    int index_s = 2*index*tid;
    if(tid < blockDim.x / (2* index))
        smem[index_s] += smem[index_s + index];
    __syncthreads();
}

在这个修改后的代码中，满足条件的线程（即if语句中的条件）是会工作的线程，而其他线程则不会参与运算。每次循环迭代后，工作的线程数量会减半，并且总是位于前blockDim.x / (2* index)个位置。

reduce_v1取余替换为位与版本

由于取余操作是一个非常耗时的操作，因此可以考虑使用位与操作来替代。具体实现如下：

for(int index = 1; index < blockDim.x; index *= 2) {
    if ((tid & (2 * index - 1)) == 0){
        smem[tid] += smem[tid + index];
    }
    __syncthreads();
}

reduce_v2解决bank_conflict版本

bank conflict发生在CUDA为了更好地并行，将共享内存分割成32个等大小的内存块（bank）。当一个warp中的不同线程访问相同bank的不同字段时，就会发生bank冲突。

共享内存中，连续的4-bytes被分配到连续的32个bank中（每个bank存放一个32-bits的数据）。上图展示了bank conflict的示意图：一个warp中的thread 0和thread 16都会访问bank 0和bank 1，从而产生两路冲突。需要注意的是，如果不同线程访问的是同一个bank中的同一个字段，会发生broadcast，但不会产生bank conflict。

为了解决bank conflict，可以调整代码，使每个线程都在同一个bank中活动。具体实现如下：

for (unsigned int index = blockDim.x / 2; index > 0; index >>= 1) {
    if (tid < index) {
        smem[tid] += smem[tid + index];
    }
    __syncthreads();
}

reduce_v3让空闲线程干活版本

在每次迭代中，最多只有一半的线程处于工作状态，这显然是一种资源浪费。因此，可以考虑让每个线程处理更多的数据。例如，在将数据加载到共享内存时，可以同时进行一次加法运算。

代码实现如下：

__shared__ float smem[blockSize];
unsigned int tid = threadIdx.x;
unsigned int gtid = blockIdx.x * (blockSize * 2) + threadIdx.x;
smem[tid] = d_in[gtid] + d_in[gtid + blockSize];
__syncthreads();

reduce_v4展开最后一维减少同步

在共享内存的读写操作中，需要特别注意同步问题。__syncthreads()用于保证一个block中所有线程的数据都正确写入后再进行后续计算。然而，在同一个warp内部，线程是并行执行的，不需要额外的同步。因此，当迭代到工作线程数量小于32时，可以避免使用__syncthreads()。

代码实现如下：

__device__ void WarpSharedMemReduce(volatile float* smem, int tid){
    float x = smem[tid];
    if (blockDim.x >= 64) {
      x += smem[tid + 32]; __syncwarp();
      smem[tid] = x; __syncwarp();
    }
    x += smem[tid + 16]; __syncwarp();
    smem[tid] = x; __syncwarp();
    x += smem[tid + 8]; __syncwarp();
    smem[tid] = x; __syncwarp();
    x += smem[tid + 4]; __syncwarp();
    smem[tid] = x; __syncwarp();
    x += smem[tid + 2]; __syncwarp();
    smem[tid] = x; __syncwarp();
    x += smem[tid + 1]; __syncwarp();
    smem[tid] = x; __syncwarp();
}

template<int blockSize>
__global__ void reduce_v4(float *d_in,float *d_out){
    __shared__ float smem[blockSize];
    int tid = threadIdx.x;
    int i = blockIdx.x * (blockSize * 2) + threadIdx.x;
    smem[tid] = d_in[i] + d_in[i + blockSize];
    __syncthreads();

    for (int s = blockDim.x / 2; s > 32; s >>= 1) {
        if (tid < s) {
            smem[tid] += smem[tid + s];
        }
        __syncthreads();
    }

    if (tid < 32) {
        WarpSharedMemReduce(smem, tid);
    }

    if (tid == 0) {
        d_out[blockIdx.x] = smem[0];
    }
}

需要注意的是，在warpReduce中也要注意读写依赖的问题。

reduce_v5完全展开

为了进一步优化性能，可以对循环进行完全展开，减少循环中的加法指令。代码实现如下：

template <int blockSize>
__device__ void BlockSharedMemReduce(float* smem) {
    if (blockSize >= 1024) {
        if (threadIdx.x < 512) {
            smem[threadIdx.x] += smem[threadIdx.x + 512];
        }
        __syncthreads();
    }
    if (blockSize >= 512) {
        if (threadIdx.x < 256) {
            smem[threadIdx.x] += smem[threadIdx.x + 256];
        }
        __syncthreads();
    }
    if (blockSize >= 256) {
        if (threadIdx.x < 128) {
            smem[threadIdx.x] += smem[threadIdx.x + 128];
        }
        __syncthreads();
    }
    if (blockSize >= 128) {
        if (threadIdx.x < 64) {
            smem[threadIdx.x] += smem[threadIdx.x + 64];
        }
        __syncthreads();
    }
    if (threadIdx.x < 32) {
        volatile float* vshm = smem;
        if (blockDim.x >= 64) {
            vshm[threadIdx.x] += vshm[threadIdx.x + 32];
        }
        vshm[threadIdx.x] += vshm[threadIdx.x + 16];
        vshm[threadIdx.x] += vshm[threadIdx.x + 8];
        vshm[threadIdx.x] += vshm[threadIdx.x + 4];
        vshm[threadIdx.x] += vshm[threadIdx.x + 2]; 
        vshm[threadIdx.x] += vshm[threadIdx.x + 1];
    }
}

reduce_v6在合适选择block数量

为了进一步压榨thread的计算能力，可以让每个线程在第一次迭代时处理更多的数据。代码实现如下：

float sum = 0.0f;
for (int32_t i = gtid; i < nums; i += total_thread_num) {
    sum += d_in[i];
}
smem[tid] = sum;
__syncthreads();

reduce_v7用shuffle做warp reduce

NV推出了Shuffle指令，对于reduce操作的优化效果显著。Shuffle指令允许warp内的寄存器相互访问，从而减少访存延迟。代码实现如下：

template <int blockSize>
__device__ float WarpShuffle(float sum) {
    sum += __shfl_down_sync(0xffffffff, sum, 16);
    sum += __shfl_down_sync(0xffffffff, sum, 8);
    sum += __shfl_down_sync(0xffffffff, sum, 4);
    sum += __shfl_down_sync(0xffffffff, sum, 2);
    sum += __shfl_down_sync(0xffffffff, sum, 1);
    return sum;
}

需要注意的是，shuffle之后超出的部分可以直接忽略。

参考链接

• 深入浅出GPU优化系列：reduce优化
• 【CUDA】Reduce规约求和（已完结~）