张量核内联函数

现代GPU已从纯向量 (vector)处理器发展为包含专用加速器的混合架构。标准CUDA核心使用单指令多线程 (SIMT) 模型处理标量运算，而张量核则同时操作整个矩阵。对于编译器工程师来说，针对这些单元需要代码生成策略的根本性转变。你不再是为单个线程发出指令来执行标量融合乘加 (FMA)。相反，你必须协调一组线程，通常是一个32个线程的warp，以协作加载数据段并触发同步矩阵指令。

Warp级矩阵运算

标准指令操作线程私有寄存器。张量核指令操作分布在warp中所有32个线程的寄存器上的片段。单个指令，例如 mma.sync，计算矩阵乘积 $D = A \times B + C$。这里，$A$、$B$、$C$和$D$是位于warp集体寄存器文件中的小矩阵块。

硬件规定了这些块尺寸的严格限制。形状 $(M, N, K)$ 的常见配置包括 $16 \times 16 \times 16$、$32 \times 8 \times 16$ 和 $8 \times 32 \times 16$。编译器必须对高级计算图进行分块，使得最内层循环严格匹配这些尺寸。如果逻辑张量大小不能被硬件支持的形状整除，编译器必须注入填充逻辑或生成标量回退循环（剥离）来处理边界，尽管填充是更好的选择，以避免线程发散。

下图展示了喂入张量核单元所需的数据流。数据从设备内存经过层级结构移动，最终到达线程私有寄存器，硬件将其解释为集体矩阵片段。

使用张量核所需的数据移动层级。编译器必须明确管理数据从全局内存到共享内存的暂存，然后将其加载到适用于矩阵乘加 (MMA) 指令的寄存器片段中。

片段抽象与布局

在降低阶段，编译器通常针对接近汇编的中间表示，例如LLVM NVPTX或直接PTX ISA。然而，处理原始寄存器索引容易出错。MLIR（特别是NVGPU dialect）或TVM等编译器基础设施采用称为片段的逻辑抽象。

片段是一种存储子矩阵的变量类型。重要的是，片段内部的数据布局是不透明的。线程0中 fragment[0] 处的float32值不一定对应于矩阵元素 $(0,0)$。硬件在线程间交错数据以优化通道使用率。因此，除非数据首先存回共享内存或通过特定片段迭代器访问，否则编译器不允许对片段进行标准逐元素操作。

加载指令决定了内存的解释方式。对于FP16矩阵乘法，编译器必须发出 load_matrix_sync 指令。此指令要求源内存的步长在编译时已知或作为寄存器参数 (parameter)传递。编译器必须确保共享内存中的数据组织有序，以防止此群加载操作期间的存储体冲突。

精度与累加

张量核本质上是混合精度单元。常见模式涉及以半精度（FP16或BF16）加载输入 ($A$ 和 $B$)，同时以单精度（FP32）累加结果 ($C$ 和 $D$)。这可以防止部分积求和期间的数值溢出。

硬件执行的数学运算是：

$D_{m,n} = \sum_{k=0}^{K-1} A_{m,k} \times B_{k,n} + C_{m,n}$

如果输入张量是FP32，编译器负责在操作前转换数据类型。这通常涉及在片段加载前发出转换指令 (cvt.f16.f32)。未能匹配特定MMA内联函数预期的数据类型将导致运行时出现非法指令错误。

PTX MMA指令的构成

检查深度学习 (deep learning)编译器的输出时，你会遇到PTX（并行线程执行）汇编。mma 指令是公开张量核的基本操作。它的语法冗长，明确编码了形状、布局和数据类型。

考虑指令 mma.sync.aligned.m16n8k16.row.col.f32.f16.f16。我们可以解析此操作码以理解编译器的意图：

1. mma.sync：矩阵乘加，在warp中同步。
2. aligned：假定内存地址与16字节边界对齐 (alignment)。编译器必须在内存分配期间强制执行此项。
3. m16n8k16：操作的形状。$M=16, N=8, K=16$。
4. row.col：输入$A$和$B$的布局。$A$是行主序，$B$是列主序。
5. f32.f16.f16：数据类型。累加器是FP32；输入$A$和$B$是FP16。

布局说明符 (row.col) 特别重要。如果内存中的输入张量是行主序但指令期望列主序，编译器有两种选择：在软件中明确转置矩阵（开销大）或选择匹配内存布局的不同内联函数变体。高级编译器执行布局传播分析，以确保数据以最有效MMA指令所需的格式到达共享内存。

流水线同步

在较新的架构（如NVIDIA Ampere和Hopper）中，张量核指令相对于内存操作是异步的。这允许编译器将计算与数据移动重叠。编译器生成一个流水线，其中：

1. 组0从全局内存加载数据到共享内存。
2. 组1从共享内存加载数据到寄存器。
3. 张量核在当前位于寄存器中的数据上执行。

为了保持正确性，编译器必须注入同步屏障 (bar.sync 或 cp.async.wait_group)，以确保数据在被使用前有效。缺少屏障会导致竞态条件，而过多屏障会使流水线停滞。自动调度算法尝试找到满足依赖关系所需的最少数量屏障，同时最大化复制和计算引擎的重叠。

以下图表展示了应用最佳调度后，标准浮点单元 (FPU) 使用率与张量核使用率之间的吞吐量 (throughput)差异。

峰值理论吞吐量比较。针对张量核的编译器可以实现数量级更高的性能，前提是循环调度和数据布局正确地与硬件约束对齐 (alignment)。

MLIR中的内联函数降低

在MLIR生态系统中，这些内联函数的生成由 NVGPU dialect 处理。一个高级的 linalg.matmul 操作首先被分块和缓冲。在转换为LLVM dialect期间，特定模式将分块循环匹配到 nvgpu.mma.sync 操作。

该过程包括：

1. 矢量化 (quantization)： 将标量加载转换为矢量加载（例如，vector<4xf16>），以饱和内存带宽。
2. Warp规约： 如果规约维度 $K$ 分布在多个线程上，则需要进行跨通道规约（蝴蝶混洗）来求和部分结果。
3. 内联函数发出： 将通用乘加主体替换为特定的LLVM内联函数调用 __nvvm_mma_sync_...。

通过自动化此降低过程，编译器将模型定义与硬件的复杂特点解耦。然而，生成代码的效率完全依赖于上游pass正确预测内联函数所要求的块大小和布局。