现代GPU已从纯向量处理器发展为包含专用加速器的混合架构。标准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$。编译器必须对高级计算图进行分块,使得最内层循环严格匹配这些尺寸。如果逻辑张量大小不能被硬件支持的形状整除,编译器必须注入填充逻辑或生成标量回退循环(剥离)来处理边界,尽管填充是更好的选择,以避免线程发散。下图展示了喂入张量核单元所需的数据流。数据从设备内存经过层级结构移动,最终到达线程私有寄存器,硬件将其解释为集体矩阵片段。digraph G { rankdir=TB; node [shape=box, style=filled, fontname="Helvetica", fontsize=10]; edge [fontname="Helvetica", fontsize=9]; subgraph cluster_0 { label = "GPU内存层级"; style = filled; color = "#f8f9fa"; GlobalMem [label="全局内存 (DRAM)\n完整张量存储", fillcolor="#e9ecef", color="#adb5bd"]; SharedMem [label="共享内存 (SRAM)\n分块数据块", fillcolor="#d0bfff", color="#9775fa"]; subgraph cluster_warp { label = "Warp(32个线程)"; style = filled; color = "#e7f5ff"; Registers [label="线程寄存器\n(片段)", fillcolor="#a5d8ff", color="#4dabf7"]; TensorCore [label="张量核单元\n(MMA指令)", fillcolor="#ff8787", color="#fa5252", shape=Mdiamond]; } } GlobalMem -> SharedMem [label="合并加载\n(异步复制)"]; SharedMem -> Registers [label="同步加载矩阵\n(ldmatrix)"]; Registers -> TensorCore [label="执行 mma.sync"]; TensorCore -> Registers [label="累加结果"]; }使用张量核所需的数据移动层级。编译器必须明确管理数据从全局内存到共享内存的暂存,然后将其加载到适用于矩阵乘加 (MMA) 指令的寄存器片段中。片段抽象与布局在降低阶段,编译器通常针对接近汇编的中间表示,例如LLVM NVPTX或直接PTX ISA。然而,处理原始寄存器索引容易出错。MLIR(特别是NVGPU dialect)或TVM等编译器基础设施采用称为片段的逻辑抽象。片段是一种存储子矩阵的变量类型。重要的是,片段内部的数据布局是不透明的。线程0中 fragment[0] 处的float32值不一定对应于矩阵元素 $(0,0)$。硬件在线程间交错数据以优化通道使用率。因此,除非数据首先存回共享内存或通过特定片段迭代器访问,否则编译器不允许对片段进行标准逐元素操作。加载指令决定了内存的解释方式。对于FP16矩阵乘法,编译器必须发出 load_matrix_sync 指令。此指令要求源内存的步长在编译时已知或作为寄存器参数传递。编译器必须确保共享内存中的数据组织有序,以防止此群加载操作期间的存储体冲突。精度与累加张量核本质上是混合精度单元。常见模式涉及以半精度(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指令的构成检查深度学习编译器的输出时,你会遇到PTX(并行线程执行)汇编。mma 指令是公开张量核的基本操作。它的语法冗长,明确编码了形状、布局和数据类型。考虑指令 mma.sync.aligned.m16n8k16.row.col.f32.f16.f16。我们可以解析此操作码以理解编译器的意图:mma.sync:矩阵乘加,在warp中同步。aligned:假定内存地址与16字节边界对齐。编译器必须在内存分配期间强制执行此项。m16n8k16:操作的形状。$M=16, N=8, K=16$。row.col:输入$A$和$B$的布局。$A$是行主序,$B$是列主序。f32.f16.f16:数据类型。累加器是FP32;输入$A$和$B$是FP16。布局说明符 (row.col) 特别重要。如果内存中的输入张量是行主序但指令期望列主序,编译器有两种选择:在软件中明确转置矩阵(开销大)或选择匹配内存布局的不同内联函数变体。高级编译器执行布局传播分析,以确保数据以最有效MMA指令所需的格式到达共享内存。流水线同步在较新的架构(如NVIDIA Ampere和Hopper)中,张量核指令相对于内存操作是异步的。这允许编译器将计算与数据移动重叠。编译器生成一个流水线,其中:组0从全局内存加载数据到共享内存。组1从共享内存加载数据到寄存器。张量核在当前位于寄存器中的数据上执行。为了保持正确性,编译器必须注入同步屏障 (bar.sync 或 cp.async.wait_group),以确保数据在被使用前有效。缺少屏障会导致竞态条件,而过多屏障会使流水线停滞。自动调度算法尝试找到满足依赖关系所需的最少数量屏障,同时最大化复制和计算引擎的重叠。以下图表展示了应用最佳调度后,标准浮点单元 (FPU) 使用率与张量核使用率之间的吞吐量差异。{ "layout": { "title": "理论吞吐量:FPU与张量核对比 (FP16)", "xaxis": {"title": "操作类型", "showgrid": false}, "yaxis": {"title": "TFLOPS", "showgrid": true, "gridcolor": "#e9ecef"}, "plot_bgcolor": "#ffffff", "paper_bgcolor": "#ffffff", "font": {"family": "Helvetica", "color": "#495057"}, "barmode": "group" }, "data": [ { "type": "bar", "x": ["标准FPU (FP32)", "张量核 (FP16累加FP32)"], "y": [19.5, 312], "marker": {"color": ["#adb5bd", "#74c0fc"]}, "text": ["19.5 TFLOPS", "312 TFLOPS"], "textposition": "auto" } ] }峰值理论吞吐量比较。针对张量核的编译器可以实现数量级更高的性能,前提是循环调度和数据布局正确地与硬件约束对齐。MLIR中的内联函数降低在MLIR生态系统中,这些内联函数的生成由 NVGPU dialect 处理。一个高级的 linalg.matmul 操作首先被分块和缓冲。在转换为LLVM dialect期间,特定模式将分块循环匹配到 nvgpu.mma.sync 操作。该过程包括:矢量化: 将标量加载转换为矢量加载(例如,vector<4xf16>),以饱和内存带宽。Warp规约: 如果规约维度 $K$ 分布在多个线程上,则需要进行跨通道规约(蝴蝶混洗)来求和部分结果。内联函数发出: 将通用乘加主体替换为特定的LLVM内联函数调用 __nvvm_mma_sync_...。通过自动化此降低过程,编译器将模型定义与硬件的复杂特点解耦。然而,生成代码的效率完全依赖于上游pass正确预测内联函数所要求的块大小和布局。