适用于CUDA GPU的Numba例子

• 适用于CUDA GPU的Numba例子

矩阵乘法

这是使用CUDA内核的矩阵乘法的简单实现：

@cuda.jit

def matmul(A, B, C):

    """Perform square matrix multiplication of C = A * B

    """

    i, j = cuda.grid(2)

    if i < C.shape[0] and j < C.shape[1]:

        tmp = 0.

        for k in range(A.shape[1]):

            tmp += A[i, k] * B[k, j]

        C[i, j] = tmp

这种实现方式简单直观，但性能不佳，因为相同的矩阵元素将从设备内存中多次加载，这很慢（某些设备可能具有透明的数据缓存，但它们可能不足以一次容纳整个输入）。

如果使用阻塞算法来减少对设备内存的访问，它将更快。CUDA为 块中的线程提供快速共享内存，以协作执行任务。以下使用共享内存实现了方阵乘法的更快版本：

from numba import cuda, float32

 

# Controls threads per block and shared memory usage.

# The computation will be done on blocks of TPBxTPB elements.

TPB = 16

 

@cuda.jit

def fast_matmul(A, B, C):

    # Define an array in the shared memory

    # The size and type of the arrays must be known at compile time

    sA = cuda.shared.array(shape=(TPB, TPB), dtype=float32)

    sB = cuda.shared.array(shape=(TPB, TPB), dtype=float32)

 

    x, y = cuda.grid(2)

 

    tx = cuda.threadIdx.x

    ty = cuda.threadIdx.y

    bpg = cuda.gridDim.x    # blocks per grid

 

    if x >= C.shape[0] and y >= C.shape[1]:

        # Quit if (x, y) is outside of valid C boundary

        return

 

    # Each thread computes one element in the result matrix.

    # The dot product is chunked into dot products of TPB-long vectors.

    tmp = 0.

    for i in range(bpg):

        # Preload data into shared memory

        sA[tx, ty] = A[x, ty + i * TPB]

        sB[tx, ty] = B[tx + i * TPB, y]

 

        # Wait until all threads finish preloading

        cuda.syncthreads()

 

        # Computes partial product on the shared memory

        for j in range(TPB):

            tmp += sA[tx, j] * sB[j, ty]

 

        # Wait until all threads finish computing

        cuda.syncthreads()

 

    C[x, y] = tmp

因为共享内存是有限的资源，所以代码一次从输入数组中预加载小块。然后，调用 syncthreads()以等待所有线程完成预加载，再对共享内存进行计算。计算之后，再次同步，以确保所有线程在共享内存中的数据均已完成之后，在下一个循环迭代中将其覆盖。