在之前的博客《借助樹狀數組的思想實現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一次遍歷即可完成前綴和計算，但是這篇論文較難讀懂暫時不理解詳細的原理，待后續研究明白再仿照其進行實現。用樹狀數組實現的好處是代碼較為簡潔，速度也還湊合，可以作為面試中的實現來學習。
))
