在Python CUDA内核中使用cublas GEMM

时间:2017-08-01 05:33:34

标签: python cuda cublas

我有一个简单的矩阵 - 矩阵乘法码如下:

TPB = 32
@cuda.jit('void(double[:, :], double[:,:], double[:, :])', device = True)
def GPU_Mat2(A, B, C):
    bx = cuda.blockIdx.x
    by = cuda.blockIdx.y
    tx = cuda.threadIdx.x
    ty = cuda.threadIdx.y
    ROW = bx * TPB + tx
    COL = by * TPB + ty
    res = 0

    for k in range(A.shape[1]):
        if ROW < A.shape[0] and COL < B.shape[1]:
           res += A[ROW, k] * B[k, COL]
    cuda.syncthreads()

    if ROW < A.shape[0] and COL < B.shape[1]:
       C[ROW, COL] = res
    cuda.syncthreads()

然后我在另一个内核中调用此函数两次。

@cuda.jit('void(double[:, :], double[:,:], double[:, :], double[:, :])')
def call_Mat2(A, B, C, D):
    for _ in range(200):
        GPU_Mat2(A, B, C)
        GPU_Mat2(C, B, D)        # Is this correct?

不幸的是,与主机中的相同计算相比,此过程无法给出正确的答案。即使我在每次GPU_Mat2调用后使用cuda.syncthreads(),答案仍然是错误的。我的问题是“是否可以在另一个内核中使用内核调用的输出(此处为C)作为输入?”

def main():
N = 300
A = np.asfortranarray(np.random.random_sample((N,N)))
B = np.asfortranarray(np.random.random_sample((N,N)))
C_GPU = np.zeros((N,N), dtype = np.double, order = 'F')
D_GPU = np.zeros((N,N), dtype = np.double, order = 'F')

numThreads = [TPB, TPB]
numBlocks =[(A.shape[0]+TPB-1)//TPB, (B.shape[1]+TPB-1)//TPB]

d_A = cuda.to_device(A)
d_B = cuda.to_device(B)
d_C = cuda.to_device(C_GPU)
d_D = cuda.to_device(D_GPU)

call_Mat2[numBlocks, numThreads](d_A, d_B, d_C, d_D)

其次,基于this,可以在内核中调用“blas GEMM”,但我在python脚本中找不到类似的例子。这种类型的调用是否受python支持? 非常感谢您的帮助。

1 个答案:

答案 0 :(得分:1)

根据documentation

  

注意:较新的CUDA设备支持设备端内核启动;此功能称为动态并行,但Numba目前不支持它。

所以不,你现在不能在numba编译的CUDA Python中调用其他设备库或@cuda.jit函数。