最近，我一直在尝试使用 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和tx + i * TPB上尝试了一些 if 语句，但这些都不起作用。
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 = …我有这段代码通过将这个 AU 矩阵的每字节 8 个元素打包到 A 中来减少内存消耗，从而使 100k*200k 矩阵占用更少的空间。如您所料，此代码需要永远运行，我还计划将行数增加到 200k。我在一个非常强大的实例（CPU 和 GPU）上运行代码并且可以扩展它，所以任何人都可以帮助并行化此代码，使其更快。
import numpy as np
colm = int(2000000/8)
rows = 1000000
cols = int(colm*8)
AU = np.random.randint(2,size=(rows, cols),dtype=np.int8)
start_time = time.time()
A = np.empty((rows,colm), dtype=np.uint8)
for i in range(A.shape[0]):
for j in range(A.shape[1]):
A[i,j] = 0
for k in range(8):
if AU[i,(j*8)+k] == 1:
A[i,j] = A[i,j] | (1<<(7-k))