用CUDA写出比Numpy更快的规约求和函数

技术背景

在前面的几篇博客中我们介绍了在Python中使用Numba来写CUDA程序的一些基本操作和方法，并且展示了GPU加速的实际效果。在可并行化的算法中，比如计算两个矢量的加和，或者是在分子动力学模拟领域中的查找近邻表等等，都是可以直接并行的算法，而且实现起来难度不大。而有一种情况是，如果我们要计算的内容的线程之间互相存在依赖，比方说最常见的，计算一个矩阵所有元素的和。

CUDA的atomic运算

正如前面所提到的问题，如何去计算一个矩阵所有元素之和呢？具体问题可以表述为：

\[S=\sum_{i,j}A_{i,j}
\]

对于此类的问题，如果我们像普通的CUDA并行操作一样，直接创建一个S变量，然后直接在线程和分块上直接把每一个矩阵元素加到这个S变量中，那么会出现一种情况：在线程同步时，存在冲突的线程是无法同时加和成功的，也就是说，这种情况下虽然程序不会报错，但是得到的结果是完全错误的。对于此类情况，CUDA官方给出了atomic运算这样的方案，可以保障线程之间不被干扰：

import numpy as np
from numba import cuda
from numba import vectorize
cuda.select_device(1)
@cuda.jit
def ReducedSum(arr, result):
    i, j = cuda.grid(2)
    cuda.atomic.add(result, 0, arr[i][j])
if __name__ == '__main__':
    import time
    np.random.seed(2)
    data_length = 2**10
    arr = np.random.random((data_length,data_length)).astype(np.float32)
    print (arr)
    arr_cuda = cuda.to_device(arr)
    np_time = 0.0
    nb_time = 0.0
    for i in range(100):
        res = np.array([0],dtype=np.float32)
        res_cuda = cuda.to_device(res)
        time0 = time.time()
        ReducedSum[(data_length,data_length),(1,1)](arr_cuda,res_cuda)
        time1 = time.time()
        res = res_cuda.copy_to_host()[0]
        time2 = time.time()
        np_res = np.sum(arr)
        time3 = time.time()
        if i == 0:
            print ('The error rate is: ', abs(np_res-res)/res)
            continue
        np_time += time3 - time2
        nb_time += time1 - time0
    print ('The time cost of numpy is: {}s'.format(np_time))
    print ('The time cost of numba is: {}s'.format(nb_time))

这里需要重点关注的就是用CUDA实现的简单函数ReducedSum，这个函数中调用了CUDA的atomic.add方法，用这个方法直接替代系统内置的加法，就完成了所有的操作。我们将这个函数的运行时间去跟np.sum函数做一个对比，结果如下：

$ python3 cuda_reduced_sum.py
[[0.4359949  0.02592623 0.5496625  ... 0.3810055  0.6834749  0.5225032 ]
 [0.62763107 0.3184925  0.5822277  ... 0.89322233 0.7845663  0.4595605 ]
 [0.9666947  0.16615923 0.6931703  ... 0.29497907 0.63724256 0.06265242]
 ...
 [0.96224505 0.36741972 0.6673239  ... 0.3115176  0.7561843  0.9396167 ]
 [0.781736   0.28829736 0.38047555 ... 0.15837361 0.00392629 0.6236886 ]
 [0.03247315 0.3664344  0.00369871 ... 0.0205253  0.15924706 0.8655231 ]]
The error rate is:  4.177044e-06
The time cost of numpy is: 0.027491092681884766s
The time cost of numba is: 0.01042938232421875s

在GPU的计算中，会有一定的精度损失，比如这里的误差率就在1e-06级别，但是运行的速度要比numpy的实现快上2倍！

总结概要

我们知道GPU加速在可并行化程度比较高的算法中，能够发挥出比较大的作用，展示出明显的加速效果，而对于一些线程之间存在依赖这样的场景就不一定能够起到很大的加速作用。CUDA官方针对此类问题，提供了atomic的内置函数解决方案，包含有求和、求最大值等常用函数。而这些函数的特点就在于，线程与线程之间需要有一个时序的依赖关系。就比如说求最大值的函数，它会涉及到不同线程之间的轮询。经过测试，CUDA的这种atomic的方案，实现起来非常方便，性能也很乐观，相比于自己动手实现一个不断切割、递归的规约函数，还是要容易快捷的多。

版权声明

本文首发链接为：https://www.cnblogs.com/dechinphy/p/gpu-sum.html

作者ID：DechinPhy

更多原著文章请参考：https://www.cnblogs.com/dechinphy/

打赏专用链接：https://www.cnblogs.com/dechinphy/gallery/image/379634.html

腾讯云专栏同步：https://cloud.tencent.com/developer/column/91958