前缀和(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芯片的并行能力,就需要考虑如何解决数据之间的依赖关系,即每个元素的计算都需要之前的所有元素的和。这样的依赖关系似乎不利于并行化,但实际上,我们可以使用一种巧妙的算法,将时间复杂度缩减到,即对数时间。
这种算法的思想是将数据分为多个部分,每个部分由一个线程处理,然后再将各个部分的结果合并起来。具体来说,我们可以将数据看成一个二叉树,每个节点的值是其左右子节点的和,根节点的值就是整个数组的和。这样,我们就可以用两个阶段来实现前缀和的计算:
- 第一阶段:从下往上,自底向上地计算每个节点的值,即每个节点的值等于其左右子节点的和。这个过程可以并行地进行,每个线程负责一个节点的计算。
- 第二阶段:从上往下,自顶向下地计算每个节点的前缀和,即每个节点的前缀和等于其父节点的前缀和加上其左兄弟节点的值。这个过程也可以并行地进行,每个线程负责一个节点的计算。
下图是一个示意图,横向的16个点代表16个数,时间轴从上往下,每个入度为2的节点会做加法,并将结果广播到其输出节点。
cumsum并行计算
如果不好理解,我在下图给出一个例子,其中,每次我都会取偏移的数据与原tensor进行相加,依次循环往复。在次计算后,得到前缀和,也就是我们要的cumsum。
这种算法的优点是可以充分利用AI芯片的并行能力,将时间复杂度降低到,而且不需要额外的空间复杂度。
[更多的细节和代码,您可以参考这些文章CUDA高性能计算经典问题(二)—— 前缀和(Prefix Sum) - 知乎 (zhihu.com)
网友评论