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