【发布时间】：2021-08-12 08:58:33
【问题描述】：
我正在发现 numba 的 cuda 扩展，并查看了 CUDA 上矩阵乘法的示例实现。代码在numba's web site。
然后我用我认为不太理想的实现对其进行了基准测试：numpy 的 dot 函数，将两个 1024x1024 矩阵相乘（使用 randn(1024,1024) 生成）
结果：
- CUDA 每次乘法 40 毫秒，
- numpy 每次乘法 5 毫秒。
如果numpy的算法是朴素矩阵乘法，那么它应该需要1024^3~1e9的乘法和加法。这是每 5ms/1e9 = 5 皮秒一次操作的平均吞吐量。我的 CPU 运行频率约为 3.4 GHz，因此每个周期需要 300 皮秒。
所以这是我的问题：numpy 的矩阵乘法如何比简单的矩阵乘法快 60 倍？
我听说过 Strassen 的算法，它的复杂度约为 N^2.8，因此每次乘法和加法需要 20 皮秒。仍然比 CPU 的速度快 30 倍。
编辑：
- cuda 方法的定义
from numba import cuda, float32
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
- Cuda 调用
N=1024
a=randn(N,N).astype(np.float32)
b=a.T.copy()
c=zeros((N,N),dtype=np.float32)
threadsperblock = (16, 16)
blockspergrid_x = math.ceil(a.shape[0] / threadsperblock[0])
blockspergrid_y = math.ceil(a.shape[1] / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
fast_matmul[blockspergrid, threadsperblock](a,b,c) # takes 40ms
同时，使用 numpy：
c=a.dot(b) # takes 5ms
我认为吞吐量不是瓶颈，因为矩阵的大小为数百万，所需的周期为数十亿。
正如一位评论者所问，这些数组是 32 位浮点数。
编辑 2：
我了解 3GHz CPU 无法在 5ps 内执行单个任务，所以显然我的意思是平均吞吐量。由于吞吐量比假设每个周期 1 次乘法和加法的估计值好 30 倍，这意味着实现被严重优化，使用
- 向量运算（我认为是 SSE 或 AVX）
- 也可能在多个内核/CPU 上进行操作。
对于 CUDA，我测试的基本类型是单精度浮点数，事先假设它们的最佳位置是单精度或半精度浮点数。
- CUDA 函数的编译时间没有实质性影响。基准测试宏 (%%timeit) 对此并不敏感。然而，
- 我认为主内存和 GPU 内存之间的传输不会显着影响性能数据，因为传输的大小远小于计算数量（小几个数量级）。
- 我尝试通过预先将阵列传输到 CUDA 设备来验证这一点，然后……哇哦，计算时间从 40 毫秒下降到 144 微秒。所以内存设备传输是大量昂贵的。感谢@talonmies 指出这一点。
但对我来说，最重要的信息是 CPU 的编译现在可以执行极其激进的优化（向量操作、多线程），在某些情况下，即使对于计算过程，它也很难用 GPU 击败它有一个非常规则的模式。即使您希望每个输入样本执行 1000 次左右的计算，与设备之间的传输也可能非常昂贵。
最后但同样重要的是，我要感谢@user2640045 他/她的基准和对数插值，这似乎表明计算在 O(N^3) 中，因此 numpy 似乎使用了最简单的矩阵乘积实施。
【问题讨论】：
-
您的时间安排无疑是错误的。在没有看到代码的情况下，您的 Numba 执行时间几乎肯定包含编译和主机设备数据传输。而且您对数据的有限描述表明您使用的是双精度，这可能会使 GPU 处于巨大的吞吐量劣势，具体取决于您使用的硬件
-
我们不将 GPU 用于一切的原因是因为在某些情况下它并不是最成功的。 talonmies 提到主机设备数据传输，这在许多 GPU 调用中往往是一个相当大的瓶颈，因为对于小数据，启动 GPU，放入数据，取出数据等可能比数学慢数百倍.
-
numpy 最肯定不使用幼稚的 matmul 实现，并且可能比 Strassen 更快（我前段时间读过这个，但不记得详细使用了哪种算法，但源代码应该可用），此外它将使用 SIMD 指令集，并且可能同时使用多个线程
-
@geebert：当前的答案表明它可以按 N^3 进行缩放，因此它可能没有使用 Strassen。所以可能算法上很简单，但在实现方面显然在 SIMD 和缓存阻塞方面进行了高度优化，并且可能还可以并行化到多个内核。 （这将使某些“5 皮秒”计算无效，这已经很奇怪，因为它们没有序列化，而是在 2x 256 位 SIMD 流水线 FMA 单元中进行的多个操作，假设是现代 x86）。同意“天真”不是一个很好的描述。也许是“小心使用蛮力”。
-
好的，现在我们看到了你的代码，我的假设是正确的。顾名思义，即时编译意味着函数调用包括编译以及与在设备上获取代码、为 GPU 分配内存以及将数据复制到 GPU 以及从 GPU 复制数据相关的所有 API 开销。使用预分配的 GPU 缓冲区运行代码并为 第二次调用 计时。您将观察到显着差异。
标签： python numpy cuda benchmarking numba