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

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))

$ 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