向量/矩阵单元的寄存器分配

管理处理器寄存器是生成高性能代码的主要步骤之一，它发生在目标特定指令选择之后。寄存器分配是编译器的主要优化之一，但现代CPU、GPU和加速器中广泛存在的向量单元（SIMD）和专用矩阵乘法单元，给传统的标量寄存器分配带来了显著增加的复杂性。高效使用这些庞大且通常专用的寄存器文件，对于实现这些单元的全部吞吐量必不可少。

向量和矩阵寄存器的特殊要求

机器学习任务中常见的向量和矩阵操作，能够同时处理大量数据。硬件通过相应的庞大寄存器文件体现了这一点：

• 大小和位宽： CPU SIMD扩展，如AVX-512，提供32个512位向量寄存器（ZMM），而Arm的SVE提供可伸缩的向量长度。GPU每个处理块通常拥有更大的向量寄存器文件（例如，NVIDIA GPU中每个SM可能达到数百KB，逻辑上在线程之间划分）。矩阵单元（如NVIDIA Tensor Cores或AMD Matrix Cores）操作的块状结构会映射到寄存器组或专用累加器。
• 结构和别名： 向量寄存器通常允许访问子部分（例如，AVX允许在512位（ZMM）寄存器中以128位（XMM）和256位（YMM）视图访问）。分配器必须正确建立这些别名关系的模型，以确保正确性并避免不必要的限制。
• 专用性： 矩阵单元可能拥有专用累加器寄存器，其行为与通用向量寄存器不同。某些ISA施加配对限制或要求矩阵指令的特定操作数使用特定寄存器类型（例如，PTX中的mma指令）。
• 高压力： 张量操作，特别是通过im2col/GEMM实现的密集矩阵乘法（GEMM）和卷积，涉及复杂的循环结构和大量的中间值。旨在暴露出指令级并行性的优化，如循环展开和软件流水线，进一步增加了对寄存器的需求（寄存器压力），尤其是在计算密集的内循环中。
• 谓词/掩码寄存器： AVX-512和Arm SVE等架构使用掩码寄存器来控制向量内的元素级操作。这些掩码本身表示活跃值，必须分配到单独的、较小的谓词寄存器文件。

标准分配方法的局限

经典的图着色寄存器分配器（基于Chaitin或Briggs算法）构成了许多编译器的依据。它们建立一个干扰图，其中节点表示活跃范围，边连接相互干扰的范围，然后尝试使用等于可用物理寄存器数量的颜色对图进行着色。然而，将这些方法直接应用于大型向量/矩阵寄存器文件会遇到问题：

• 可伸缩性： 潜在活跃范围的数量之多以及寄存器的大小会导致庞大、密集的干扰图，使得着色问题在计算上代价高昂。
• 建模复杂性： 在标准干扰图中表示子寄存器别名、寄存器配对或专用累加器约束可能很繁琐或低效。
• 溢出成本： 溢出一个512位向量寄存器比溢出一个64位标量寄存器的代价高得多，因为需要更大的内存带宽。天真的溢出选择（例如，溢出干扰图中度数最高的活跃范围）可能不是最优的，它忽略了实际成本或溢出频率。

向量/矩阵单元的高级分配策略

为了应对这些挑战，编译器采用更复杂的、为向量和矩阵寄存器量身定制的技术：

1. 再具体化： 对于一个值（特别是常数或容易从其他值推导出的值，例如生成一个全零向量），分配器可以选择稍后重新计算（再具体化），而不是溢出和重新加载。这避免了对于易于重新生成的值产生的昂贵内存流量。编译器会识别其结果可以再具体化的指令，并权衡重新计算的成本与溢出/重新加载的成本。

2. 活跃范围拆分和寄存器打包： 当向量寄存器包含多个独立的较小值，或当一个值仅在向量通道的子集中活跃时，分配器可能会拆分该活跃范围。这允许原始活跃范围的不同部分分配给不同的物理寄存器或独立溢出。反之，如果多个小的、不相互干扰的值可以放入单个向量寄存器中，它们可以打包在一起，从而降低整体寄存器需求。

3. 优化溢出代码： 当溢出不可避免时，分配器必须生成高效的溢出代码。

   • 部分溢出： 如果在潜在溢出点之后只需要向量的某些通道，编译器可能会生成代码仅溢出这些特定通道，从而减少内存带宽使用。
   • 放置： 优先将数据溢出到最近/最快的缓存级别（例如L1），而非主内存。分配器可以使用启发式方法或分析数据来指导溢出槽位放置。
   • 调度： 溢出和填充指令（向量加载/存储）被视为内存操作，需要仔细调度以隐藏其延迟，并可能将其与计算重叠。

4. 寄存器分块： 这种技术将寄存器分配与循环分块优化（在第4章中讨论）紧密结合。内循环的结构使得一个计算块（例如，矩阵乘法的子块）的工作集能够放入可用的向量/矩阵寄存器中。对于GEMM ($C += A * B$)，这通常意味着将C矩阵的一个块（$C_{sub}$）保存在寄存器中（通常是累加器），并通过其他寄存器流式传输A和B的块。分配器的目的是最小化$C_{sub}$块在迭代间的重新加载。

5. 处理矩阵累加器： 针对矩阵单元的分配器需要特定的策略。在这些单元内累积的部分和极其有价值，且溢出代价高昂。分配器必须优先保持这些部分和驻留，通常通过仔细调度遍历矩阵块的外循环来实现。具体的指令（例如PTX mma、HLSL波矩阵内在函数）通常决定了操作数和累加器如何映射到寄存器文件。

6. 阶段排序考量： 寄存器分配是在指令调度之前还是之后执行的经典困境，在向量/矩阵单元上变得更加突出。早期分配会限制调度器，而后期分配如果调度创建了高寄存器压力，可能会强制产生更多溢出。现代编译器通常使用迭代方法或集成调度和分配阶段，特别是对于性能关键的循环。

GPU寄存器分配：占用率权衡

在GPU上，寄存器分配对占用率有直接而显著的影响。占用率指的是可以在流多处理器（SM）上并行驻留的活跃warp（线程组）的数量。每个SM都有一个庞大的物理寄存器文件，但它由在该SM上运行的所有线程共享。

• 每线程限制： GPU硬件和运行时对每个线程可使用的寄存器数量施加限制。
• 占用率计算： SM上可用的物理寄存器总数除以每个线程分配的寄存器数量，决定了（连同共享内存使用等其他因素一起）有多少线程以及多少warp可以并行运行。 $\text{每个SM的最大Warp数} \approx \frac{\text{每个SM的总寄存器数}}{\text{每个线程的寄存器数} \times \text{每个Warp的线程数}}$
• 性能影响： 更高的占用率允许SM通过在就绪的warp之间切换来更有效地隐藏内存延迟。然而，为增加占用率而减少每个线程的寄存器分配，可能需要线程内部更多的溢出代码，从而可能减慢单个线程的执行速度。

编译器必须应对这种权衡。积极地分配寄存器可能会在线程内部实现更好的指令级并行，但会降低线程级并行（占用率）。反之，最小化寄存器使用会增加占用率，但可能由于溢出或展开减少而导致性能损失。GPU编译器通常使用启发式方法、分析数据或允许程序员提示（如CUDA中的__launch_bounds__）来指导这种平衡。

每线程分配的寄存器数量与在SM上可以并发运行的最大warp数量之间的关系，假设寄存器是限制因素。

示例：分块GEMM中的寄存器压力

考虑一个简化的矩阵乘法内循环（$C_{ij} += A_{ik} \times B_{kj}$），我们目标是保持一个$4 \times 4$的C矩阵块在寄存器中。这需要16个累加器寄存器（标量或向量，取决于目标）。为了计算这个块，我们可能需要在最内层（$k$）循环的每次迭代中，为A的一个面板加载例如4个向量寄存器，为B的一个面板加载4个向量寄存器。

• 最小寄存器需求： 16（用于C块）+ 4（用于A面板）+ 4（用于B面板）= 24个向量寄存器，外加用于中间结果和地址计算的临时寄存器。
• 挑战： 如果可用的向量寄存器文件较小（例如，只有16个向量寄存器），编译器无法同时保持整个C块、A面板和B面板驻留。
• 策略： 在循环转换的指导下，分配器可能会优先保持C累加器驻留，并溢出/重新加载A或B的一部分，或者使用较小的C块大小。或者，如果ISA支持，它可能会采用寄存器旋转或其他复杂技术，从而可能将下一次迭代的A/B面板加载与当前计算重叠。

有效管理向量和矩阵寄存器不仅仅是应用标准分配算法。它需要充分了解目标架构的能力和限制，与指令调度和循环优化阶段的仔细交互，以及用于最小化溢出宽向量或专用矩阵数据所带来高成本的复杂策略。在此做出的选择，对于在现代异构硬件上将优化后的IR转换为高性能机器代码非常重要。