美文网首页
并行计算如何加速计算前缀和(cumsum)

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

作者: 若隐爱读书 | 来源:发表于2023-12-22 11:25 被阅读0次

前缀和(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)

相关文章

网友评论

      本文标题:并行计算如何加速计算前缀和(cumsum)

      本文链接:https://www.haomeiwen.com/subject/jzofndtx.html