并行计算如何加速计算前缀和(cumsum)

前缀和(Prefix Sum)

前缀和是一种常用的算法技巧，它可以快速地求出一个数组的某个区间的和。前缀和的思想是，对于一个数组a，我们可以预先计算出一个数组b，使得b[i]等于a[0]到a[i]的和，即b[i] = a[0] + a[1] + … + a[i]。

下面用python代码讲述前缀和的实现流程

# 定义一个数组a
a = [1, 2, 3, 4, 5]
# 创建一个和a长度相同的数组b，用来存放前缀和
b = [0] * len(a)
# 用一个变量sum来记录当前的和
sum = 0
# 遍历a中的每个元素
for i in range(len(a)):
    # 把当前元素加到sum中
    sum += a[i]
    # 把sum赋值给b中对应的位置
    b[i] = sum
# 输出b
print(b)
# [1, 3, 6, 10, 15]

Cumsum

cumsum是AI框架中的一个函数，它可以实现前缀和的功能。它有一个参数axis，用来指定沿着哪个轴进行累加。如果axis为None，那么就对整个数组进行展平后累加；如果axis为0，那么就沿着第一个维度（行）累加；如果axis为1，那么就沿着第二个维度（列）累加，以此类推。cumsum的返回值是一个和原数组形状相同的数组，其中每个元素都是沿着指定轴的前缀和。

加速器

在CPU上，我们可以用一个简单的循环来实现前缀和的计算，如下所示：

void PrefixSum(const int32_t* input, size_t n, int32_t* output) {
  int32_t sum = 0;
  for (size_t i = 0; i < n; ++i) {
    sum += input[i];
    output[i] = sum;
  }
}

这种方法的时间复杂度是O(n)，即线性时间。但是，如果我们想要利用AI芯片的并行能力，就需要考虑如何解决数据之间的依赖关系，即每个元素的计算都需要之前的所有元素的和。这样的依赖关系似乎不利于并行化，但实际上，我们可以使用一种巧妙的算法，将时间复杂度缩减到 $O(log_2(n))$ ，即对数时间。

这种算法的思想是将数据分为多个部分，每个部分由一个线程处理，然后再将各个部分的结果合并起来。具体来说，我们可以将数据看成一个二叉树，每个节点的值是其左右子节点的和，根节点的值就是整个数组的和。这样，我们就可以用两个阶段来实现前缀和的计算：

第一阶段：从下往上，自底向上地计算每个节点的值，即每个节点的值等于其左右子节点的和。这个过程可以并行地进行，每个线程负责一个节点的计算。
第二阶段：从上往下，自顶向下地计算每个节点的前缀和，即每个节点的前缀和等于其父节点的前缀和加上其左兄弟节点的值。这个过程也可以并行地进行，每个线程负责一个节点的计算。

下图是一个示意图，横向的16个点代表16个数，时间轴从上往下，每个入度为2的节点会做加法，并将结果广播到其输出节点。

cumsum并行计算

如果不好理解，我在下图给出一个例子，其中，每次我都会取 $2^n+1$ 偏移的数据与原tensor进行相加，依次循环往复。在 $log2_(n)$ 次计算后，得到前缀和，也就是我们要的cumsum。

image.png

这种算法的优点是可以充分利用AI芯片的并行能力，将时间复杂度降低到 $log_2(n)$ ，而且不需要额外的空间复杂度。

[更多的细节和代码，您可以参考这些文章CUDA高性能计算经典问题（二）—— 前缀和（Prefix Sum） - 知乎 (zhihu.com)

最后编辑于：2023.12.23 11:25:57