Numba python CUDA 与 cuBLAS 简单操作的速度差异

Question

我正在分析一些代码，但无法找出性能差异。我正在尝试在两个数组之间（就地）进行简单的逐元素加法。这是使用 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%。这是在 "warming up" cuda.jit 内核之前调用它之后，因此编译的 PTX 代码已经被缓存（不确定这是否重要，但这样做只是为了确保这不是问题）。

我可以理解 3 级矩阵-矩阵运算的这种性能差异，但这是一个简单的补充。我可以做些什么来从 cuda.jit 代码中获得更多性能吗？我问是因为我想要优化的真正代码是一个二维数组，它不能传递给 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)

Answer 1

非常简短的回答是否定的。 CUBLAS 利用许多东西（纹理、矢量类型）来提高像这样的内存绑定代码的性能，numba CUDA 方言目前不支持。

我在 CUDA 中完成了这个：

__device__ float4 add(float4 x, float4 y) 
{
    x.x += y.x; x.y += y.y; x.z += y.z; x.w += y.w; 
    return x;
} 

__global__ void mykern(float* x, float* y, int N)
{
    float4* x4 = reinterpret_cast<float4*>(x);
    float4* y4 = reinterpret_cast<float4*>(y);

    int strid = gridDim.x * blockDim.x;
    int tid = threadIdx.x + blockDim.x * blockIdx.x;

    for(; tid < N/4; tid += strid) {
        float4 valx = x4[tid];
        float4 valy = y4[tid];
        y4[tid] = add(valx, valy);
    }       
}

我的基准测试显示它在 CUBLAS 的大约 5% 以内，但我不相信你现在可以在 numba 中做到这一点。

顺便说一句，我不明白你关于不能运行 saxpy 二维数组的说法。如果你的数组在内存中是连续的（我怀疑它们必须是）并且具有相同的布局（即不尝试添加转置），那么你 can 使用 saxpy on二维数组。

Numba python CUDA 与 cuBLAS 简单操作的速度差异

Numba python CUDA vs. cuBLAS speed difference on simple operations

python

cuda

numba