CUDA分块矩阵乘法解释

Question

我试图理解CUDA SDK 8.0中这个示例代码的工作方式：

template <int BLOCK_SIZE> __global__ void
matrixMulCUDA(float *C, float *A, float *B, int wA, int wB)
{
// Block index
int bx = blockIdx.x;
int by = blockIdx.y;

// Thread index
int tx = threadIdx.x;
int ty = threadIdx.y;

// Index of the first sub-matrix of A processed by the block
int aBegin = wA * BLOCK_SIZE * by;

// Index of the last sub-matrix of A processed by the block
int aEnd   = aBegin + wA - 1;

// Step size used to iterate through the sub-matrices of A
int aStep  = BLOCK_SIZE;

// Index of the first sub-matrix of B processed by the block
int bBegin = BLOCK_SIZE * bx;

// Step size used to iterate through the sub-matrices of B
int bStep  = BLOCK_SIZE * wB;

....
....

对于我来说，内核的这一部分相当棘手。我知道矩阵A和B表示为数组(*float)，我也知道使用共享内存以计算点积的概念。

我的问题是，我不理解代码的开头，特别是3个特定变量（aBegin、aEnd和bBegin）。能否给我一个可能执行的示例图，帮助我理解这个特定情况下索引的工作方式？谢谢

- docdev

Answer 1

这是一幅图示，用来理解CUDA内核中第一个变量设置的值和整个计算流程： 矩阵使用行主序存储。CUDA代码假定矩阵大小可以被BLOCK_SIZE整除。

矩阵A、B和C根据kernel CUDA网格被虚拟分成块。所有C的块都可以并行计算。对于给定的C的深灰色块，主循环遍历几个浅灰色块A和B(同步进行)。每个块使用BLOCK_SIZE * BLOCK_SIZE线程并行计算。 bx和by是当前块在CUDA网格内的基于块的位置。 tx和ty是由当前线程计算的单元格在CUDA网格的当前计算块内的基于单元格的位置。

这里是关于aBegin变量的详细分析： aBegin指的是矩阵A的第一个计算块的第一个单元格的内存位置。它被设置为wA * BLOCK_SIZE * by，因为每个块包含BLOCK_SIZE * BLOCK_SIZE个单元格，并且水平方向上有wA / BLOCK_SIZE块，在当前计算块的上方还有by个块。因此，(BLOCK_SIZE * BLOCK_SIZE) * (wA / BLOCK_SIZE) * by = BLOCK_SIZE * wA * by。

同样的逻辑也适用于bBegin： 它被设置为BLOCK_SIZE * bx，因为在矩阵B的第一个计算块的第一个单元格之前，内存中有bx个大小为BLOCK_SIZE的块。 a在循环中按aStep = BLOCK_SIZE递增，这样下一个计算块就是当前A计算块右边（在图示中）。b在同一循环中按bStep = BLOCK_SIZE * wB递增，这样下一个计算块就是当前B计算块底部（在图示中）。

- Jérôme Richard