为张量核心和矩阵单元生成代码

现代硬件加速器，特别是图形处理器（GPU），逐渐集成专门的执行单元，旨在大幅加速密集矩阵乘法和卷积运算，这些运算是深度学习工作负载的基础。英伟达（NVIDIA）从Volta架构开始引入了张量核心（Tensor Cores），而AMD则通过其CDNA/RDNA架构中的矩阵核心（Matrix Core）或矩阵融合乘加（MFMA）指令提供类似功能。与标准浮点单元相比，这些单元提供了显著更高的理论吞吐量，但发挥这一潜力需要复杂的编译器支持。

矩阵单元的架构原理

张量核心和矩阵核心通常在一个周期或少量周期内执行类似于小型、固定大小的矩阵乘法和累加运算。一种常见形式是：

$$D = A \times B + C$$

其中 $A$、$B$、$C$ 和 $D$ 是小型矩阵（例如 $4 \times 4$、$8 \times 8$、$16 \times 16$，具体取决于硬件代次和特定指令）。一个重要方面是它们对混合精度的原生支持。输入 $A$ 和 $B$ 通常要求采用较低精度格式，如FP16（半精度浮点）、BF16（bfloat16），甚至INT8/INT4（8位/4位整数），而累加（$C$ 和 $D$）经常以更高精度进行，如FP32（单精度浮点），以在多次累加中保持数值精度。

这种在矩阵片段上的混合精度FMA运算是它们性能优势的主要原因。与使用一系列标量或向量FMA指令执行等效运算相比，它们将更多的计算密度集成到硬件中。

编译器针对矩阵单元的挑战

为这些单元生成高效代码给编译器后端带来了若干挑战：

1. 指令映射与抽象： 编译器需要识别计算模式，主要是中间表示（IR）中的矩阵乘法或卷积，这些模式可以映射到这些专门的矩阵指令（例如，英伟达张量核心的PTX中的mma.sync，AMD的mfma内联函数）。这通常涉及通过多个阶段降低高级线性代数运算（如MLIR中的linalg.matmul）。
2. 数据布局与寄存器要求： 矩阵单元通常要求其输入操作数（$A$ 和 $B$ 片段）位于特定寄存器中，这些寄存器通常以不直接映射到标准标量或向量寄存器文件的方式组织。编译器必须管理将数据从内存（全局或共享）加载到这些寄存器中，并采用正确的布局。例如，一个张量核心指令可能要求矩阵A的片段以行主序格式加载到寄存器中，而矩阵B的片段以列主序格式加载。
3. 分块与循环转换： 大型矩阵乘法必须分解（分块）成与硬件矩阵指令支持的维度相匹配的小型矩阵乘法。编译器需要生成嵌套循环来迭代这些块，并协调输入块的加载和输出块的累加。
4. 共享内存管理： 为隐藏内存延迟并提供数据重用，输入矩阵块通常放置在GPU的快速共享内存中。编译器必须生成代码以高效地将块从全局内存加载到共享内存，同时遵守共享内存 bank 约束，然后将片段从共享内存加载到矩阵单元的寄存器中。
5. 调度与同步： 编译器必须仔细调度内存加载指令、矩阵乘加指令以及任何必要的同步（特别是在使用mma.sync等指令的GPU warp内），以在确保正确性的同时最大化指令级并行性和流水线利用率。

矩阵单元代码生成的编译器技术

编译器采用若干技术来应对这些挑战：

• 目标特定IR方言： MLIR等框架采用方言（例如，英伟达的nvgpu，AMD的rocdl）来暴露更接近硬件能力的抽象。这些方言可能包含直接表示矩阵FMA指令的操作或内联函数。
• 渐进式降低： 编译过程涉及多个降低步骤。例如，linalg.matmul可能首先分块以在适合共享内存的块上操作。然后，在每个块内，会发生进一步的分块（寄存器分块）以匹配矩阵单元的维度。这个内部循环随后被降低到方言特定的内联函数或直接到汇编指令，如PTX mma.sync或ROCDL mfma。
• 仿射循环分析与转换： 可以使用多面体建模或其他仿射分析技术来分析分块生成的循环嵌套，并优化数据移动和计算调度。
• 矩阵片段的寄存器分配： 需要专门的寄存器分配策略来管理逻辑矩阵片段到矩阵单元使用的物理寄存器文件的映射。这涉及仔细跟踪寄存器使用情况，如果压力过高，可能涉及溢出/重新加载。
• 共享内存布局优化： 编译器分析内存访问模式，以确定共享内存中矩阵块的最佳布局，旨在避免 bank 冲突，并在 warp 内线程加载数据时最大化带宽。在共享内存中填充或交织数据是常见技术。

矩阵乘法的分块策略。大型矩阵A、B和C被分割成更小的块（例如A11、B11、C11），由硬件矩阵单元迭代处理。输入块被加载、乘法运算并累加到相应的输出块中。

• 内联函数调用： 编译器通常不直接生成原始指令，而是降低到目标特定的内联函数（例如，通过CUDA C++或HIP头文件暴露）。这些内联函数为矩阵指令提供了略高层次的接口，简化了代码生成，同时仍允许细粒度控制。
• 使用供应商库： 对于非常常见的矩阵大小和配置，编译器可能会决定调用供应商库中（如英伟达的cuBLAS或AMD的rocBLAS，其内部使用张量/矩阵核心）高度优化的函数，比生成自定义代码性能更好。这通常涉及模式匹配特定的操作序列并用库调用替换它们。

为专门的矩阵单元生成高效代码是一个复杂的优化问题，涉及计算、数据移动和资源管理的仔细协调。它需要编译器IR、转换遍和目标特定后端之间的深度集成。在此取得成功对于在现代加速器上实现深度学习训练和推理的最先进性能非常重要。