Numba Python在简单操作上与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)

BLAS方法的性能对于大型数组（20M个元素）大约快了25％。这是在“热身”cuda.jit内核之后，先前调用它以使编译的PTX代码已经被缓存（不确定是否重要，但只是为了确保这不是问题）。

我可以理解这种性能差异适用于三级矩阵-矩阵运算，但这只是一个简单的加法。有什么方法可以挤出更多的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)

- user1554752

---

如果你要发布代码，能不能确保它至少可以编译通过呢？在像这样简单的内核中，执行参数对性能至关重要，但你没有展示它们。请你修复一下这个问题好吗？ - talonmies

这真的是你的执行参数吗？每个块64个线程和1024个块？ - talonmies

是的，但我尝试了其他TPB和区块的组合。 - user1554752

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中实现。另外，关于无法在2D数组上运行saxpy的评论，我不理解。如果您的数组在内存中是连续的（正如我怀疑它们必须是的那样）并且具有相同的布局（即不尝试添加转置），则您可以在2D数组上使用saxpy。

- talonmies