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)
解决方法:
非常简短的答案是“否”. 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.
标签:cuda,numba,python 来源: https://codeday.me/bug/20191025/1932489.html