CUDA 经典问题:前缀和 | CUDA

本文最后更新于:2022年1月31日

问题

对于数组 a,其前缀和为数组 b,a 和 b 的长度均为 n。对于任意 i < n 都满足 b[i] = a[0] + a[1] + ... + a[i]

基本思路

前缀和的思路如下:

  1. 将整个数据分成几个部分,每个部分分别计算前缀和,存入数组 output 中,然后将每个部分中最大的值存入一个数组 part 中
  2. 对上述数组 part 求前缀和
  3. 将 part 中的元素分别加到 output 中

代码实现

Baseline

__global__ void ScanPart(int *input, int *part, int *output, int n, int part_num) {
  for (int part_i = blockIdx.x; part_i < part_num; part_i += gridDim.x) {
    int part_begin = part_i * blockDim.x;
    int part_end = min((part_i + 1) * blockDim.x, n);
    if (threadIdx.x == 0) {
      int acc = 0;
      for (int i = part_begin; i < part_end; ++i) {
        acc += input[i];
        output[i] = acc;
      }
      part[part_i] = acc;
    }
  }
}

__global__ void ScanPartSum(int *part, int part_num) {
  int acc = 0;
  for (int i = 0; i < part_num; ++i) {
    acc += part[i];
    part[i] = acc;
  }
}

__global__ void AddPartSum(