继续优化基于树状数组的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的两个元素相加。最终的结果正好和树状数组的定义一模一样。

由于cuda编程的特殊性，比较难以实现跨block之间的同步，而为了构造完整的树状数组，我们必须要访问不同block之间的元素。在之前的方案中，我们只启动了一个block来实现跨block的树状数组构造，这样就严重限制了该步骤的并行性。事实上，我们还可以继续借助block之内的树状数组构造来完成block之间的树状数组构造。区别就在于block之内构造树状数组的步幅是1，而block之间的构造步幅不再是1，所以我们只需要传递一个步幅就可以将block之内和block之间的树状数组构造统一起来。按照该思路实现的reduce代码如下：

__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 9while(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一次遍历即可完成前缀和计算，但是这篇论文较难读懂暂时不理解详细的原理，待后续研究明白再仿照其进行实现。用树状数组实现的好处是代码较为简洁，速度也还凑合，可以作为面试中的实现来学习。