python - 与 numpy 相比，cupy 代码不够快

Question

我是cupy的新手，我一直在玩numpy和cupy来比较执行时间。

import cupy 
import cupyx
import numpy as np
import time

def test_cupy(a0,a1,ind_i,ind_j):

        t1=time.time()
        for i in range(300):
            cupyx.scatter_add(a0,(ind_i,ind_j),cupy.exp(-(ind_i+ind_j)))
            cupyx.scatter_add(a1,(ind_i,ind_j),cupy.exp(-(ind_i+ind_j)))
        t2=time.time()
        print("cupy",t2-t1)
        return t2-t1

def test_numpy(a0,a1,ind_i,ind_j):

        t1=time.time()
        for i in range(300):
            np.add.at(a0,(ind_i,ind_j),np.exp(-(ind_i+ind_j)))
            np.add.at(a1,(ind_i,ind_j),np.exp(-(ind_i+ind_j)))
        t2=time.time()
        print("numpy",t2-t1)
        return t2-t1

a_np=np.zeros((20000,20000)) + 1
a_cp=cupy.zeros((20000,20000)) + 1

ind_np=np.arange(np.shape(a_np)[0])
ind_cp=cupy.arange(np.shape(a_np)[0])

t_cupy=test_cupy(a_cp,a_cp,ind_cp,ind_cp)
t_numpy=test_numpy(a_np,a_np,ind_np,ind_np)

print(t_numpy/t_cupy)

我有一个具有 2560 个内核的 Geforce gtx 1080，时间比 t_numpy/t_cupy 是 18.. 比我预期的要慢得多。我怎样才能达到更好的性能？

Answer 1

（我建议您从使用nvvp或 Nsight 分析代码开始。）

代码看起来计算量不大，因此内存读/写或内核启动成本将成为瓶颈。一种方法是将元素操作融合到一个内核中，使用cupy.fuse.

@cupy.fuse
def exp_add_neg(ind_i, ind_j):
    return cupy.exp(-(ind_i + ind_j))

...
cupyx.scatter_add(a0, (ind_i, ind_j), exp_neg_add(ind_i, ind_j))
cupyx.scatter_add(a1, (ind_i, ind_j), exp_neg_add(ind_i, ind_j))

它在我的环境中提高了约 30% 的速度。