CUDA 平铺矩阵乘法解释

Posted 2023-04-15

【中文标题】CUDA 平铺矩阵乘法解释

【英文标题】：CUDA tiled matrix multiplication explanation

【发布时间】：2021-06-02 04:54:28

【问题描述】：

我正在尝试了解 CUDA SDK 8.0 中的 this sample 代码的工作原理：

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）。有人可以为我制作一个可能执行的示例图，以帮助我了解索引在这种特定情况下是如何工作的吗？谢谢

【参考方案1】：

这是一张图，用于了解设置为 CUDA 内核的第一个变量的值以及执行的整体计算：

矩阵使用行优先顺序存储。 CUDA 代码假设矩阵大小可以除以BLOCK_SIZE。

矩阵 A、B 和 C 根据内核 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 单元，并且在当前计算块A 上方有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 当前计算块的底部（在图形上）的下面。