Skip to content

5.8。示例

原文： http://numba.pydata.org/numba-doc/latest/roc/examples.html

5.8.1。矩阵乘法

以下是使用 HSA 内核的矩阵乘法的简单实现：

@roc.jit
def matmul(A, B, C):
    i = roc.get_global_id(0)
    j = roc.get_global_id(1)

    if i >= C.shape[0] or j >= C.shape[1]:
        return

    tmp = 0

    for k in range(A.shape[1]):
        tmp += A[i, k] * B[k, j]

    C[i, j] = tmp

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

如果我们使用阻塞算法来减少对设备内存的访问，则会更快。 HSA 为组中的工作项提供快速共享内存，以协同计算任务。以下实现了使用共享内存的方形矩阵乘法的更快版本：

import numpy as np
from numba import roc
from numba import float32
from time import time as timer

blocksize = 16
gridsize = 16

@roc.jit('(float32[:,:], float32[:,:], float32[:,:])')
def matmulfast(A, B, C):
    x = roc.get_global_id(0)
    y = roc.get_global_id(1)

    tx = roc.get_local_id(0)
    ty = roc.get_local_id(1)

    sA = roc.shared.array(shape=(blocksize, blocksize), dtype=float32)
    sB = roc.shared.array(shape=(blocksize, blocksize), dtype=float32)

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

    tmp = 0

    for i in range(gridsize):
        # preload
        sA[tx, ty] = A[x, ty + i * blocksize]
        sB[tx, ty] = B[tx + i * blocksize, y]
        # wait for preload to end
        roc.barrier(1)
        # compute loop
        for j in range(blocksize):
            tmp += sA[tx, j] * sB[j, ty]
        # wait for compute to end
        roc.barrier(1)

    C[x, y] = tmp

N = gridsize * blocksize
A = np.random.random((N, N)).astype(np.float32)
B = np.random.random((N, N)).astype(np.float32)
C = np.zeros_like(A)

griddim = gridsize, gridsize
blockdim = blocksize, blocksize

with roc.register(A, B, C):
    ts = timer()
    matmulfast[griddim, blockdim](A, B, C)
    te = timer()
    print("1st GPU time:", te - ts)

with roc.register(A, B, C):
    ts = timer()
    matmulfast[griddim, blockdim](A, B, C)
    te = timer()
    print("2nd GPU time:", te - ts)

ts = timer()
ans = np.dot(A, B)
te = timer()
print("CPU time:", te - ts)
np.testing.assert_allclose(ans, C, rtol=1e-5)

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

回到顶部