最近，我一直在尝试使用 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 = …