继续优化基于树状数组的cuda前缀和
继续优化基于树状数组的cuda前缀和
在之前的博客《
》中，我们用三个kernel实现了基于树状数组的cuda版前缀和，但是在数据量较大时速度不如传统的reduce-then-scan方法，主要原因在于跨block的reduce阶段没有充分利用所有的cuda核心。在本博客中，我们尝试进一步优化，将三个kernel减少到两个kernel，并在跨block的reduce阶段尝试使用更多核心来提升性能。
基于树状数组的方法源于传统方法在reduce阶段其实就是构造了一个完备的树状数组，所以我们可以将scan阶段改成基于树状数组查询的方式，从而达到简化代码的目的。
1. Reduce阶段
Reduce阶段完成两个工作：1. block内构造树状数组；2.block之间构造树状数组。Block之内构造树状数组如下图所示，首先步幅为1的两个元素相加，然后步幅为2的两个元素相加，之后是步幅为4、8、16…直到步幅为block_size/2的两个元素相加。最终的结果正好和树状数组的定义一模一样。
__global__ static void gen_bit_in_one_block(int *input, long long n, long long offset) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
long long pos = (tid + 1) * offset - 1;
int id = threadIdx.x;
__shared__ int reduced_sum[512];
reduced_sum[id] = input[pos];
__syncthreads();
int offset = 2;
#pragma unroll 9
while(offset <= 512) {
if(((tid + 1) & (offset - 1)) == 0) {
reduced_sum[id] += reduced_sum[id - offset / 2];
}
__syncthreads();
offset <<= 1;
}
input[pos] = reduced_sum[id];
}
}上面的代码和之前的实现基本一致，主要是增加了一个offset变量用来获取正确的初始步幅。此外，为了加速显存访问，我们引入了共享内存，构造树状数组的过程都在共享内存中完成。
上述代码在构造一个warp内的树状数组时，我们还可以使用__shfl_up_sync来加速，但是收益不明显，感兴趣的可以继续尝试优化。
2. Scan阶段
构造好完整的树状数组之后，我们就可以利用其查询每个位置的前缀和了，代码和之前一样：
__global__ static void calc_sum_using_bit(int *input, int *output, int n) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < n) {
int sum = input[tid];
int idx = tid + 1;
idx -= (idx & -idx);
while (idx > 0) {
sum += input[idx - 1];
idx -= (idx & -idx);
}
output[tid] = sum;
}
}不过经过仔细分析我们就会发现上面的代码实现不是work-efficient的，我们详细解释一下原因。CPU串行实现计算前缀和只需要n次加法。Reduce阶段总共也只会有n个线程参与求和：n/2 + n/4 + n/8 + … = n，所以reduce阶段是work-efficient的。但是scan阶段，参与运算的线程总数为从1到n的二进制中1的个数，而这个数值是n*logn量级的，所以虽然整体复杂度是O(logn)的，但是参与运算的线程数是O(nlogn)的，所以就不属于work-efficient的实现。
3. 完整调用逻辑
我们借助上面两个kernel实现完整的前缀和计算。代码如下：
void calc_prefix_sum(int *input, int *output, int n) {
int *buffer1, *buffer2;
cudaMalloc(&buffer1, n * sizeof(int));
cudaMalloc(&buffer2, n * sizeof(int));
cudaMemcpy(buffer1, input, n * sizeof(int), cudaMemcpyHostToDevice);
long long offset = 1;
long long count = n;
dim3 dimBlock(512);
do {
dim3 dimGrid(get_block_size(count, 512));
gen_bit_in_one_block<<<dimGrid, dimBlock>>>(buffer1, count, offset);
count /= 512;
offset *= 512;
} while(offset < n);
dim3 dimGrid2(get_block_size(n, 512));
calc_sum_using_bit<<<dimGrid, dimBlock2>>>(buffer1, buffer2, n);
cudaMemcpy(output, buffer2, n * sizeof(int), cudaMemcpyDeviceToHost);
cudaFree(buffer1);
cudaFree(buffer2);
}相比之前的实现，我们需要循环调用gen_bit_in_one_block来构造完整的树状数组。所以，虽然只有两个kernel，但是调用kernel的次数不止两次。
4. 性能对比
我们用RTX4090来验证性能。这次我们引入cub实现，号称最快的前缀和实现，然后再加上传统的reduce-then-scan实现，看看三种不同实现在不同长度下的性能如何，对比如下（单位毫秒）：
| 长度 | 100 | 1000 | 1万 | 10万 | 100万 | 1000万 | 1亿 | 10亿 |
|---|---|---|---|---|---|---|---|---|
| BIT | 0.19 | 0.19 | 0.19 | 0.19 | 0.22 | 0.39 | 2.11 | 19.46 |
| BCAO | 0.02 | 0.02 | 0.03 | 0.03 | 0.08 | 0.31 | 2.23 | 19.59 |
| CUB | 0.03 | 0.03 | 0.03 | 0.04 | 0.04 | 0.07 | 0.87 | 8.69 |
可以看出，基于树状数组的实现在长度较长时可以与BCAO类似，但是两者都远差于CUB的实现。CUB的实现基于论文《Single-pass Parallel Prefix Scan with Decoupled Look-back》实现，只需要一个kernel一次遍历即可完成前缀和计算，但是这篇论文较难读懂暂时不理解详细的原理，待后续研究明白再仿照其进行实现。用树状数组实现的好处是代码较为简洁，速度也还凑合，可以作为面试中的实现来学习。