python cuda_Numba python CUDA与cuBLAS在简单操作上的速度差...

我正在分析一些代码,无法找出性能差异.我正在尝试在两个数组(就地)之间做一个简单的元素方式加法.这是使用numba的CUDA内核：

from numba import cuda

@cuda.jit('void(float32[:], float32[:])')

def cuda_add(x, y):

ix = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x

stepSize = cuda.gridDim.x * cuda.blockDim.x

while ix < v0.shape[0]:

y[ix] += x[ix]

ix += stepSize

我以为性能不错,但随后将其与cuBLAS方法进行了比较：

from accelerate.cuda.blas import Blas

blas = Blas()

blas.axpy(1.0, X, Y)

对于大型阵列(20M元素),BLAS方法的性能大约快25％.这是在通过预先调用cuda.jit内核“预热”了它之后,因此已经缓存了已编译的PTX代码(不确定是否很重要,但只是为了确保不是问题).

我可以理解3级矩阵矩阵操作的性能差异,但这是一个简单的加法.我有什么办法可以从cuda.jit代码中榨取更多性能？我问是因为我要优化的真实代码是2d数组,无法传递给blas.axpy.

编辑执行代码和其他所需的软件包：

import numpy as np

def main():

n = 20 * 128 * 128 * 64

x = np.random.rand(n).astype(np.float32)

y = np.random.rand(n).astype(np.float32)

## Create necessary GPU arrays

d_x = cuda.to_device(x)

d_y = cuda.to_device(y)

## My function

cuda_add[1024, 64](d_x , d_y)

## cuBLAS function

blas = Blas()

blas.axpy(1.0, d_x , d_y)