AMD 矩阵核心

AMD matrix cores — ROCm Blogs

注意： 本文博客之前是 AMD lab notes 博客系列的一部分。

矩阵乘法是线性代数的一个基本方面，它在高性能计算（HPC）应用中是一个普遍的计算。自从 AMD 推出 CDNA 架构以来，广义矩阵乘法（GEMM）计算现在通过矩阵核心处理单元实现了硬件加速。矩阵核心加速的 GEMM 内核位于像 rocBLAS 这样的 BLAS 库的核心，但开发人员也可以直接编程。通过利用矩阵核心，可以使在 GEMM 计算受限的吞吐量的应用程序获得额外的加速。

AMD 的矩阵核心技术支持全范围的混合精度操作，使我们能够处理大型模型并增强任何 AI 和机器学习工作负载的内存受限操作性能。各种数值格式在不同的应用中有其用途。例如，8 位整数（INT8）用于机器学习推理，32 位浮点数（FP32）用于机器学习训练和高性能计算应用，16 位浮点数（FP16）用于图形工作负载，以及 16 位脑浮点（BF16）用于在训练中减少收敛问题的机器学习。

要了解使用矩阵核心相比于 SIMD 向量单元所能实现的理论加速，请参考下表。表格列出了上一代（MI100）和当前一代（MI250X）CDNA 加速器的向量单元（即融合乘加（FMA））和矩阵核心单元的性能。

MI100 和 MI250X 的矩阵核心性能：

Data format	MI100 Flops/Clock/CU	MI250X Flops/Clock/CU
FP64	N/A	256
FP32	256	256
FP16	1024	1024
BF16	512	1024
INT8	1024	1024

矢量（FMA）单元在 MI100 和 MI250X 上的性能：

Data format	MI100 Flops/Clock/CU	MI250X Flops/Clock/CU
FP64	64	128
FP32	128	128

与矢量单元性能相比，MI100 和 MI250X 上的矩阵核心速度提升。_注意，MI250X 还支持打包 FP32 指令，这也会使 FP32 吞吐量加倍_：

Data format	MI100 Matrix/Vector Speedup	MI250X Matrix/Vector Speedup
FP64	N/A	2x
FP32	2x	2x

使用 AMD 矩阵核心

AMD CDNA GPU 中的矩阵融合乘加（MFMA）指令在每个波前（wavefront）上操作，而不是在每个车道（线程）上操作：输入和输出矩阵的条目分布在波前的矢量寄存器的车道上。

可以通过多种方式利用 AMD 矩阵核心。在高层次上，可以使用诸如 rocBLAS 或 rocWMMA 等库在 GPU 上进行矩阵操作。例如，rocBLAS 可以在有利于当前计算时选择使用 MFMA 指令。对于更接近底层的方法，可以选择：
- 完全用汇编语言编写 GPU 内核（这可能有些具有挑战性且不太实用）
- 在 HIP 内核中夹杂内联汇编（不推荐，因为编译器不会查看内联指令的语义，可能不会解决数据危害，例如在使用 MFMA 指令结果之前所需的强制性周期数）

- 使用编译器内置函数：这些函数表示汇编指令，以便编译器了解其语义和要求。

本文中的编码示例使用了一些可用的 MFMA 指令的编译器内置函数，并展示了如何将输入和输出矩阵的条目映射到波前的矢量寄存器车道上。所有示例都使用单个波前来计算一个小的矩阵乘法。这些示例并非旨在展示如何从 MFMA 操作中获得高性能。

MFMA编译器内部函数语法

考虑以下矩阵乘法 MFMA 操作，其中所有操作数 $A$ 、 $B$ 、 $C$ 、 $D$ 均为矩阵：

$D = AB + C$

要在 AMD GPU 上执行 MFMA 操作，LLVM 内置了函数。回想一下，这些内置函数是在整个波阵面宽度（wavefront-wide）上执行的，输入和输出矩阵的部分内容会加载到波阵面中每条通道的寄存器中。MFMA 编译器内部函数的语法如下所示：
$d = \_\_builtin\_amdgcn\_mfma\_CDFmt\_MxNxKABFmt (a, b, c, cbsz, abid, blgp)$

其中，
- CDFmt 是 C 和 D 矩阵的数据格式
- ABFmt 是 A 和 B 矩阵的数据格式
- M、`N` 和 K 是矩阵的维度：
  - mA[M][K] 源矩阵 A
  - mB[K][N] 源矩阵 B
  - mC[M][N] 累加输入矩阵 C
  - mD[M][N] 累加结果矩阵 D
- a 是存储源矩阵 A 的值的向量寄存器集合
- b 是存储源矩阵 B 的值的向量寄存器集合
- c 是存储累加输入矩阵 C 的值的向量寄存器集合
- d 是存储累加结果矩阵 D 的值的向量寄存器集合
- cbsz，控制广播大小修饰符，用于更改输入值馈送到矩阵核心的方式，仅受到具有多个输入块的 A 矩阵指令的支持。设置 cbsz 会通知指令将一个选定的输入块的值广播到 A 矩阵中的 2^cbsz 个其他邻近块。使用 abid 参数来确定选择哪个输入块进行广播。默认值 0 表示不广播值。例如，对于 16 块的 A 矩阵，设置 cbsz=1 将导致块 0 和 1 接收相同的输入值，块 2 和 3 接收相同的输入值，块 4 和 5 接收相同的输入值，等等。

- abid，A 矩阵广播标识符，支持具有多个输入块的 A 矩阵指令。它与 cbsz 一起使用，并指示选择哪个输入块广播到 A 矩阵中的其他邻近块。例如，对于 16 块的 A 矩阵，设置 cbsz=2 且 abid=1 将导致块 1 的值被广播到块 0-3，块 5 的值被广播到块 4-7，块 9 的值被广播到块 8-11，依此类推。
- blgp，B 矩阵通道组模式修饰符，允许对通道之间的 B 矩阵数据进行一组限制的变换操作。对于支持此修饰符的指令，可以使用以下值：
  - blgp=0 正常的 B 矩阵布局
  - blgp=1 从通道 0-31 的 B 矩阵数据也会被广播到通道 32-63
  - blgp=2 从通道 32-63 的 B 矩阵数据会被广播到通道 0-31
  - blgp=3 所有通道的 B 矩阵数据向下旋转 16 位（例如，通道 0 的数据会被放入通道 48，通道 16 的数据会被放入通道 0）
  - blgp=4 从通道 0-15 的 B 矩阵数据会被广播到通道 16-31、32-47 和 48-63
  - blgp=5 从通道 16-31 的 B 矩阵数据会被广播到通道 0-15、32-47 和 48-63
  - blgp=6 从通道 32-47 的 B 矩阵数据会被广播到通道 0-15、16-31 和 48-63
  - blgp=7 从通道 48-63 的 B 矩阵数据会被广播到通道 0-15、16-31 和 32-47

在 CDNA2 GPU 上支持的矩阵维度和块数量列在下表中。

A/B Data Format	C/D Data Format	M	N	K	Blocks	Cycles	Flops/cycle/CU
FP32	FP32
		32	32	2	1	64	256
		32	32	1	2	64	256
		16	16	4	1	32	256
		16	16	1	4	32	256
		4	4	1	16	8	256
FP16	FP32
		32	32	8	1	64	1024
		32	32	4	2	64	1024
		16	16	16	1	32	1024
		16	16	4	4	32	1024
		4	4	4	16	8	1024
INT8	INT32
		32	32	8	1	64	1024
		32	32	4	2	64	1024
		16	16	16	1	32	1024
		16	16	4	4	32	1024
		4	4	4	16	8	1024
BF16	FP32
		32	32	8	1	64	1024
		32	32	4	2	64	1024
		16	16	16	1	32	1024
		16	16	4	4	32	1024
		4	4	4	16	8	1024

		32	32	4	1	64	512
		32	32	2	2	64	512
		16	16	8	1	32	512
		16	16	2	4	32	512
		4	4	2	16	8	512
FP64	FP64
		16	16	4	1	32	256
		4	4	4	4	16	128

完成的 CDNA2 架构支持的所有指令列表可以在 AMD Instinct MI200 Instruction Set Architecture Reference Guide 中找到。AMD 的 Matrix Instruction Calculator 工具允许生成关于 AMD Radeon™ 和 AMD Instinct™ 加速器上 MFMA 指令的计算吞吐量和寄存器使用等更多信息。

示例 1 - V_MFMA_F32_16x16x4F32

考虑矩阵乘法运算 $D = AB$ ，其中 $M = N = 16$ 和 $K = 4$ ，且元素类型为 FP32。为简化计算，我们假设输入矩阵 $C$ 含有零元素。我们将演示使用内建函数 __builtin_amdgcn_mfma_f32_16x16x4f32 计算一次调用中四个外积的和。此函数操作单个块的矩阵。

输入矩阵 $A$ 和 $B$ 的尺寸分别为 $16 \times 4$ 和 $4 \times 16$ ，矩阵 $C$ 和 $D$ 的尺寸为 $16 \times 16$ 。将一个 $16 \times 4$ 线程块映射到两个输入矩阵的元素是方便的。在此，每个线程块有一个波阵面，x 维上有 16 个线程，y 维上有 4 个线程。我们采用行主序格式来表示矩阵： A[i][j] = j + i * N，其中 $i$ 是行索引， $j$ 是列索引。使用此表示方法，位置 $x, y$ 的线程会加载条目 A[x][y] 和 B[y][x]。输出矩阵有 $16 \times 16$ 个元素，因此每个线程都有 4 个元素要存储，如下图和代码片段所示。

以下两张图显示了 1) A 和 B 输入的形状和大小；2) A 和 B 的元素如何在波阵面所属的寄存器中映射到不同的通道中。

通过这样的描述和图示，您能更直观地理解 MFMA 指令在高性能计算任务中的应用及其实现方式。

下面的两幅图显示了：1) 输出矩阵 D 的形状和大小；2) D 矩阵的元素如何映射到波前拥有的寄存器中的通道中。

下面给出了一个执行此 MFMA 操作的示例内核。

#define M 16
#define N 16
#define K 4using float4 = __attribute__( (__vector_size__(K * sizeof(float)) )) float;__global__ void sgemm_16x16x4(const float *A, const float *B, float *D)
{float4 dmn = {0};int mk = threadIdx.y + K * threadIdx.x;int kn = threadIdx.x + N * threadIdx.y;float amk = A[mk];float bkn = B[kn];dmn = __builtin_amdgcn_mfma_f32_16x16x4f32(amk, bkn, dmn, 0, 0, 0);for (int i = 0; i < 4; ++i) {const int idx = threadIdx.x + i * N + threadIdx.y * 4 * N;D[idx] = dmn[i];}
}

该内核的启动方式如下。

dim3 grid (1, 1, 1);
dim3 block(16, 4, 1);sgemm_16x16x4 <<< grid, block >>> (d_A, d_B, d_D);

如前所述，输入 C 矩阵假定包含零。

例子 2 - V_MFMA_F32_16x16x1F32

考虑使用编译器内建函数 __builtin_amdgcn_mfma_f32_16x16x1f32 进行矩阵乘法，矩阵的尺寸为 M=N=16 和 K=1。在这种情况下，输入值可以仅由波阵面（wavefront）的16个通道（lanes）持有。实际上，这条指令可以同时乘以4个这样的矩阵，因此每个通道持有其中一个矩阵的值。
我们可以重新使用上一个例子的图来说明该操作的数据布局。在这种情况下，输入矩阵 A 不是16×4的矩阵，而是四个16×1的矩阵。但它们的布局方式，以及每个通道在波阵面（wavefront）中拥有的元素是相同的。A矩阵的“列”是不同的16×1矩阵。输入矩阵 B 也是类似的。

给定矩阵乘法的输出数据布局与前一个例子完全相同。不同之处在于，现在有四个独立的输出，每个乘法对应一个输出。
下面的代码示例展示了对4个尺寸为M=N=16和K=1的矩阵进行批量打包乘法运算的内核。

#define M 16
#define N 16
#define K 1using float16 = __attribute__( (__vector_size__(16 * sizeof(float)) )) float;__global__ void sgemm_16x16x1(const float *A, const float *B, float *D)
{float16 dmnl = {0};int mkl = K * threadIdx.x + M * K * threadIdx.y;int knl = threadIdx.x + N * K * threadIdx.y;float amkl = A[mkl];float bknl = B[knl];dmnl = __builtin_amdgcn_mfma_f32_16x16x1f32(amkl, bknl, dnml, 0, 0, 0);for (int l = 0; l < 4; ++l) {for (int i = 0; i < 4; ++i) {const int idx = threadIdx.x + i * N  + threadIdx.y * 4 * N + l * M * N;D[idx] = dmnl[i];}}
}

此内核使用以下方式启动：

dim3 grid (1, 1, 1);
dim3 block(16, 4, 1);sgemm_16x16x1 <<< grid, block >>> (d_A, d_B, d_D);

示例 3 - V_MFMA_F64_4x4x4F64

考虑 V_MFMA_F64_4x4x4F64 指令，它计算四个独立的大小为 4×4 的矩阵块的 MFMA。执行的操作是 $Z_{N}=W_{N}X_{N}+Y_{N}$ ，其中， $W_{N}$ 、 $X_{N}$ 、 $Y_{N }$ 和 $Z_{N}$ 都是大小为 4×4 元素的矩阵，且 N=0,1,2,3。
下图显示了 1) 输入参数 A 和 B 的四个组成部分的大小和形状，以及 2) 这些组成部分如何映射到波前持有的寄存器中的通道。该指令的参数包括 A、B、C 并返回 D，因此我们理解为每个参数和输出都包含 4 个矩阵。

输出D和输入C的布局与输入B的布局相同。

关于rocWMMA的一点说明

我们仅介绍了三个使用编译器内建函数来利用AMD矩阵核心的示例。更多示例可以在rocm-blogs/blogs/software-tools-optimization/matrix-cores at release · ROCm/rocm-blogs · GitHub找到。请注意，内建函数可能会在未来发生变化，因此最好使用AMD的rocWMMA C++库来加速混合精度MFMA操作。rocWMMA API有助于将矩阵乘累加问题分解为片段，并在波阵列内并行分布进行块状操作。该API是GPU设备代码的头文件库，可以将矩阵核心加速直接编译到你的内核设备代码中。这可以在生成内核汇编时受益于编译器优化。更多详情请参考rocWMMA仓库。