如何使用 numba 在 GPU 上推广快速矩阵乘法
How to generalize fast matrix multiplication on GPU using numba
最近我一直在尝试使用 Numba 库在 Python 中进行 GPU 编程。我一直在使用那里的教程在他们的网站上阅读它，目前我停留在他们的示例上，可以在这里找到：https://numba.pydata.org/numba-doc/latest/cuda/examples.html。我试图概括一下快速矩阵乘法的示例（其形式为 A*B=C）。在测试时，我注意到维度不能被每个块的线程数 (TPB) 完全整除的矩阵不会产生正确的答案。
我从 https://numba.pydata.org/numba-doc/latest/cuda/examples.html 的示例中复制了下面的代码，并创建了一个非常小的 4 x 4 矩阵测试用例。如果我选择 TPB=2 一切都很好，但是当我设置 TPB=3 时就出错了。我知道代码超出了矩阵的范围，但我无法阻止这种情况的发生（我在 ty + i * TPB 和 [=19= 上尝试了一些 if 语句]tx + i * TPB，但这些都不起作用。
from numba import cuda, float32
import numpy as np
import math
@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
#%%
x_h = np.arange(16).reshape([4,4])
y_h = np.ones([4,4])
z_h = np.zeros([4,4])
x_d = cuda.to_device(x_h)
y_d = cuda.to_device(y_h)
z_d = cuda.to_device(z_h)
TPB = 3
threadsperblock = (TPB, TPB)
blockspergrid_x = math.ceil(z_h.shape[0] / threadsperblock[0])
blockspergrid_y = math.ceil(z_h.shape[1] / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
z_h = z_d.copy_to_host()
print(z_h)
我想编写一些不依赖于维度可被 TPB 完全整除的矩阵 A、B 和 C 的代码，因为这些有时是我无法控制的。我知道 GPU 只有在对非常大的矩阵进行矩阵乘法时才会更快，但我想使用小示例来检查答案是否正确，然后再将其应用于实际数据。
可以说 that posted code 中至少有两个错误：
这不可能是正确的范围检查：
if x >= C.shape[0] and y >= C.shape[1]:
为了让我们确定网格中的特定线程不进行任何加载 activity，我们需要 x范围或 y 超出范围。 and 应该是 or.
如果块中的所有线程都不能参与该语句，则在条件代码中使用cuda.syncthreads()是illegal。上面第 1 项中的前一个 return 语句（即使从 and 更正为 or）几乎可以保证问题大小的这种非法行为不是 whole-number-divisible 线程块大小。
因此，要解决这些问题，我们不能只对 out-of-bounds 线程使用简单的 return 语句。相反，在加载点，如果计算出的全局加载索引（对于 A 或 B）是 in-bounds（共享内存），我们必须只允许线程从全局加载到共享内存根据定义，索引是 in-bounds。此外，在写入结果时，我们必须只写入 in-bounds for C.
的计算结果
以下代码修复了这些项目。它似乎适用于您给定的测试用例：
$ cat t49.py
from numba import cuda, float32
import numpy as np
import math
@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
# Each thread computes one element in the result matrix.
# The dot product is chunked into dot products of TPB-long vectors.
tmp = float32(0.)
for i in range(bpg):
# Preload data into shared memory
sA[tx, ty] = 0
sB[tx, ty] = 0
if x < A.shape[0] and (ty+i*TPB) < A.shape[1]:
sA[tx, ty] = A[x, ty + i * TPB]
if y < B.shape[1] and (tx+i*TPB) < B.shape[0]:
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()
if x < C.shape[0] and y < C.shape[1]:
C[x, y] = tmp
#%%
x_h = np.arange(16).reshape([4,4])
y_h = np.ones([4,4])
z_h = np.zeros([4,4])
x_d = cuda.to_device(x_h)
y_d = cuda.to_device(y_h)
z_d = cuda.to_device(z_h)
TPB = 3
threadsperblock = (TPB, TPB)
blockspergrid_x = math.ceil(z_h.shape[0] / threadsperblock[0])
blockspergrid_y = math.ceil(z_h.shape[1] / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
z_h = z_d.copy_to_host()
print(z_h)
print(x_h@y_h)
$ cuda-memcheck python t49.py
========= CUDA-MEMCHECK
[[ 6. 6. 6. 6.]
[22. 22. 22. 22.]
[38. 38. 38. 38.]
[54. 54. 54. 54.]]
[[ 6. 6. 6. 6.]
[22. 22. 22. 22.]
[38. 38. 38. 38.]
[54. 54. 54. 54.]]
========= ERROR SUMMARY: 0 errors
$
（请注意，在边界测试中使用 and 是正确的。测试一组索引是否为 in-bound 在布尔意义上与测试一组索引是否为 in-bound 是不同的out-of-bounds。在 in-bounds 测试中，我们要求两者都是 in-bounds。在 out-of-bounds 测试中，任何一个索引 out-of-bounds 都是不合格的）。
我并不是在暗示以上代码 defect-free 或适合任何特定目的。提供它是为了演示我发现的问题的可能修复方法。正如您所发现的，让 shared-memory 平铺矩阵乘法在每个可以想象的配置中工作是 non-trivial，并且我没有在此处显示的范围之外对其进行测试。 （例如，如果你决定让 TPB 大于 32，你会 运行 进入其他问题。另外，原来发布的代码只宣传方阵乘法，这在一般情况下不起作用 non-square 案例。)
如上所述，发布的代码和上面带有“修复”的代码将无法正确处理一般的 non-square 情况。我相信一些简单的修改将使我们能够处理 non-square 案例。简而言之，我们必须将网格的大小设置得足够大以处理两个输入矩阵的维度，同时仍然只为输出矩阵的 in-bounds 值写入结果。这是一个简单测试的例子：
$ cat t49.py
from numba import cuda, float32
import numpy as np
import math
@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
# Each thread computes one element in the result matrix.
# The dot product is chunked into dot products of TPB-long vectors.
tmp = float32(0.)
for i in range(bpg):
# Preload data into shared memory
sA[ty, tx] = 0
sB[ty, tx] = 0
if y < A.shape[0] and (tx+i*TPB) < A.shape[1]:
sA[ty, tx] = A[y, tx + i * TPB]
if x < B.shape[1] and (ty+i*TPB) < B.shape[0]:
sB[ty, tx] = B[ty + i * TPB, x]
# Wait until all threads finish preloading
cuda.syncthreads()
# Computes partial product on the shared memory
for j in range(TPB):
tmp += sA[ty, j] * sB[j, tx]
# Wait until all threads finish computing
cuda.syncthreads()
if y < C.shape[0] and x < C.shape[1]:
C[y, x] = tmp
#%%
x_h = np.arange(115).reshape([5,23])
y_h = np.ones([23,7])
z_h = np.zeros([5,7])
x_d = cuda.to_device(x_h)
y_d = cuda.to_device(y_h)
z_d = cuda.to_device(z_h)
#TPB must be an integer between 1 and 32
TPB = 32
threadsperblock = (TPB, TPB)
grid_y_max = max(x_h.shape[0],y_h.shape[0])
grid_x_max = max(x_h.shape[1],y_h.shape[1])
blockspergrid_x = math.ceil(grid_x_max / threadsperblock[0])
blockspergrid_y = math.ceil(grid_y_max / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
z_h = z_d.copy_to_host()
print(z_h)
print(x_h@y_h)
$ cuda-memcheck python t49.py
========= CUDA-MEMCHECK
[[ 253. 253. 253. 253. 253. 253. 253.]
[ 782. 782. 782. 782. 782. 782. 782.]
[1311. 1311. 1311. 1311. 1311. 1311. 1311.]
[1840. 1840. 1840. 1840. 1840. 1840. 1840.]
[2369. 2369. 2369. 2369. 2369. 2369. 2369.]]
[[ 253. 253. 253. 253. 253. 253. 253.]
[ 782. 782. 782. 782. 782. 782. 782.]
[1311. 1311. 1311. 1311. 1311. 1311. 1311.]
[1840. 1840. 1840. 1840. 1840. 1840. 1840.]
[2369. 2369. 2369. 2369. 2369. 2369. 2369.]]
========= ERROR SUMMARY: 0 errors
$
我还重新排序了 x 和 y 的含义（以及 tx 和 ty 的用法）以解决上述代码中的性能问题。同样的性能问题也存在于原始发布的文档代码中。
同样，没有无缺陷的索赔。此外，我确信可以得出“更优化”的代码。然而，优化矩阵乘法是一项应该很快导致使用库实现的练习。在 GPU 方法中使用 cupy 应该是在 GPU 上利用 high-quality 矩阵乘法例程的一种相当直接的方法。
编辑：正如所讨论的 here OP's code (and, it seems, the doc example) 在 tmp 变量的设置方面也存在性能问题。将其更改为适当的 32 位浮点变量会产生重要的性能差异。
最近我一直在尝试使用 Numba 库在 Python 中进行 GPU 编程。我一直在使用那里的教程在他们的网站上阅读它，目前我停留在他们的示例上，可以在这里找到：https://numba.pydata.org/numba-doc/latest/cuda/examples.html。我试图概括一下快速矩阵乘法的示例（其形式为 A*B=C）。在测试时，我注意到维度不能被每个块的线程数 (TPB) 完全整除的矩阵不会产生正确的答案。
我从 https://numba.pydata.org/numba-doc/latest/cuda/examples.html 的示例中复制了下面的代码，并创建了一个非常小的 4 x 4 矩阵测试用例。如果我选择 TPB=2 一切都很好，但是当我设置 TPB=3 时就出错了。我知道代码超出了矩阵的范围，但我无法阻止这种情况的发生（我在 ty + i * TPB 和 [=19= 上尝试了一些 if 语句]tx + i * TPB，但这些都不起作用。
from numba import cuda, float32
import numpy as np
import math
@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
#%%
x_h = np.arange(16).reshape([4,4])
y_h = np.ones([4,4])
z_h = np.zeros([4,4])
x_d = cuda.to_device(x_h)
y_d = cuda.to_device(y_h)
z_d = cuda.to_device(z_h)
TPB = 3
threadsperblock = (TPB, TPB)
blockspergrid_x = math.ceil(z_h.shape[0] / threadsperblock[0])
blockspergrid_y = math.ceil(z_h.shape[1] / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
z_h = z_d.copy_to_host()
print(z_h)
我想编写一些不依赖于维度可被 TPB 完全整除的矩阵 A、B 和 C 的代码，因为这些有时是我无法控制的。我知道 GPU 只有在对非常大的矩阵进行矩阵乘法时才会更快，但我想使用小示例来检查答案是否正确，然后再将其应用于实际数据。
可以说 that posted code 中至少有两个错误：
这不可能是正确的范围检查：
if x >= C.shape[0] and y >= C.shape[1]:为了让我们确定网格中的特定线程不进行任何加载 activity，我们需要
x范围或y超出范围。and应该是or.如果块中的所有线程都不能参与该语句，则在条件代码中使用
cuda.syncthreads()是illegal。上面第 1 项中的前一个return语句（即使从and更正为or）几乎可以保证问题大小的这种非法行为不是 whole-number-divisible 线程块大小。
因此，要解决这些问题，我们不能只对 out-of-bounds 线程使用简单的 return 语句。相反，在加载点，如果计算出的全局加载索引（对于 A 或 B）是 in-bounds（共享内存），我们必须只允许线程从全局加载到共享内存根据定义，索引是 in-bounds。此外，在写入结果时，我们必须只写入 in-bounds for C.
以下代码修复了这些项目。它似乎适用于您给定的测试用例：
$ cat t49.py
from numba import cuda, float32
import numpy as np
import math
@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
# Each thread computes one element in the result matrix.
# The dot product is chunked into dot products of TPB-long vectors.
tmp = float32(0.)
for i in range(bpg):
# Preload data into shared memory
sA[tx, ty] = 0
sB[tx, ty] = 0
if x < A.shape[0] and (ty+i*TPB) < A.shape[1]:
sA[tx, ty] = A[x, ty + i * TPB]
if y < B.shape[1] and (tx+i*TPB) < B.shape[0]:
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()
if x < C.shape[0] and y < C.shape[1]:
C[x, y] = tmp
#%%
x_h = np.arange(16).reshape([4,4])
y_h = np.ones([4,4])
z_h = np.zeros([4,4])
x_d = cuda.to_device(x_h)
y_d = cuda.to_device(y_h)
z_d = cuda.to_device(z_h)
TPB = 3
threadsperblock = (TPB, TPB)
blockspergrid_x = math.ceil(z_h.shape[0] / threadsperblock[0])
blockspergrid_y = math.ceil(z_h.shape[1] / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
z_h = z_d.copy_to_host()
print(z_h)
print(x_h@y_h)
$ cuda-memcheck python t49.py
========= CUDA-MEMCHECK
[[ 6. 6. 6. 6.]
[22. 22. 22. 22.]
[38. 38. 38. 38.]
[54. 54. 54. 54.]]
[[ 6. 6. 6. 6.]
[22. 22. 22. 22.]
[38. 38. 38. 38.]
[54. 54. 54. 54.]]
========= ERROR SUMMARY: 0 errors
$
（请注意，在边界测试中使用 and 是正确的。测试一组索引是否为 in-bound 在布尔意义上与测试一组索引是否为 in-bound 是不同的out-of-bounds。在 in-bounds 测试中，我们要求两者都是 in-bounds。在 out-of-bounds 测试中，任何一个索引 out-of-bounds 都是不合格的）。
我并不是在暗示以上代码 defect-free 或适合任何特定目的。提供它是为了演示我发现的问题的可能修复方法。正如您所发现的，让 shared-memory 平铺矩阵乘法在每个可以想象的配置中工作是 non-trivial，并且我没有在此处显示的范围之外对其进行测试。 （例如，如果你决定让 TPB 大于 32，你会 运行 进入其他问题。另外，原来发布的代码只宣传方阵乘法，这在一般情况下不起作用 non-square 案例。)
如上所述，发布的代码和上面带有“修复”的代码将无法正确处理一般的 non-square 情况。我相信一些简单的修改将使我们能够处理 non-square 案例。简而言之，我们必须将网格的大小设置得足够大以处理两个输入矩阵的维度，同时仍然只为输出矩阵的 in-bounds 值写入结果。这是一个简单测试的例子：
$ cat t49.py
from numba import cuda, float32
import numpy as np
import math
@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
# Each thread computes one element in the result matrix.
# The dot product is chunked into dot products of TPB-long vectors.
tmp = float32(0.)
for i in range(bpg):
# Preload data into shared memory
sA[ty, tx] = 0
sB[ty, tx] = 0
if y < A.shape[0] and (tx+i*TPB) < A.shape[1]:
sA[ty, tx] = A[y, tx + i * TPB]
if x < B.shape[1] and (ty+i*TPB) < B.shape[0]:
sB[ty, tx] = B[ty + i * TPB, x]
# Wait until all threads finish preloading
cuda.syncthreads()
# Computes partial product on the shared memory
for j in range(TPB):
tmp += sA[ty, j] * sB[j, tx]
# Wait until all threads finish computing
cuda.syncthreads()
if y < C.shape[0] and x < C.shape[1]:
C[y, x] = tmp
#%%
x_h = np.arange(115).reshape([5,23])
y_h = np.ones([23,7])
z_h = np.zeros([5,7])
x_d = cuda.to_device(x_h)
y_d = cuda.to_device(y_h)
z_d = cuda.to_device(z_h)
#TPB must be an integer between 1 and 32
TPB = 32
threadsperblock = (TPB, TPB)
grid_y_max = max(x_h.shape[0],y_h.shape[0])
grid_x_max = max(x_h.shape[1],y_h.shape[1])
blockspergrid_x = math.ceil(grid_x_max / threadsperblock[0])
blockspergrid_y = math.ceil(grid_y_max / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
z_h = z_d.copy_to_host()
print(z_h)
print(x_h@y_h)
$ cuda-memcheck python t49.py
========= CUDA-MEMCHECK
[[ 253. 253. 253. 253. 253. 253. 253.]
[ 782. 782. 782. 782. 782. 782. 782.]
[1311. 1311. 1311. 1311. 1311. 1311. 1311.]
[1840. 1840. 1840. 1840. 1840. 1840. 1840.]
[2369. 2369. 2369. 2369. 2369. 2369. 2369.]]
[[ 253. 253. 253. 253. 253. 253. 253.]
[ 782. 782. 782. 782. 782. 782. 782.]
[1311. 1311. 1311. 1311. 1311. 1311. 1311.]
[1840. 1840. 1840. 1840. 1840. 1840. 1840.]
[2369. 2369. 2369. 2369. 2369. 2369. 2369.]]
========= ERROR SUMMARY: 0 errors
$
我还重新排序了 x 和 y 的含义（以及 tx 和 ty 的用法）以解决上述代码中的性能问题。同样的性能问题也存在于原始发布的文档代码中。
同样，没有无缺陷的索赔。此外，我确信可以得出“更优化”的代码。然而，优化矩阵乘法是一项应该很快导致使用库实现的练习。在 GPU 方法中使用 cupy 应该是在 GPU 上利用 high-quality 矩阵乘法例程的一种相当直接的方法。
编辑：正如所讨论的 here OP's code (and, it seems, the doc example) 在 tmp 变量的设置方面也存在性能问题。将其更改为适当的 32 位浮点变量会产生重要的性能差异。