GPU 内存层次结构映射

高性能深度学习 (deep learning)核函数很少受处理器算术能力的限制。相反，它们的瓶颈在于数据从内存移动到计算单元的速度。将神经网络 (neural network)图编译到硬件时，主要的工程难题是掩盖内存延迟。一块标准的 NVIDIA A100 GPU 能以数百 teraFLOPS 的速率执行浮点运算，但其全局内存带宽大约只在 1.5 - 2 TB/s 之间。如果不对内存层次结构进行明确管理，计算单元（Tensor Cores）就会停滞，等待操作数到来。

硬件感知代码生成需要将“循环分块”或“张量切片”的抽象想法映射到特定的物理内存空间。GPU 内存层次结构呈现为带宽递增而容量递减的金字塔形。编译器必须确定哪些数据存放在全局内存中，哪些子集构成共享内存中的工作集，以及哪些值驻留在寄存器中以进行即时计算。

GPU 内存模型

GPU 内存架构展现了三个不同的地址空间，编译器必须明确地进行管理。与 CPU 相比，其缓存管理对软件来说基本是无感知的（硬件管理），而 GPU 则依赖编译器将数据明确地调度到更快的片上内存层级。

1. 全局内存： 最大、最慢的层级。它位于片外高带宽内存 (HBM) 中。它在核函数启动之间保持数据，并可被所有线程访问。
2. L2 缓存： 一个由硬件管理的统一缓存，由所有流式多处理器 (SM) 共享。它缓冲对全局内存的访问，但程序员不能直接寻址。
3. 共享内存 (SRAM)： 位于 SM 上的物理位置，是一种快速、低延迟的内存暂存器。它在一个线程块内的线程之间共享。这是矩阵乘法和卷积中数据复用的核心层级。
4. 寄存器： 最快的内存，单个线程私有。操作数必须在寄存器中才能被 ALU 或 Tensor Cores 使用。

下图说明了线程层次结构和内存层次结构之间的关系，显示了可见性和访问范围。

该层次结构决定了可见性和速度。数据从全局内存流经 L2 缓存进入共享内存，在那里可供线程块访问，然后加载到寄存器中进行私有线程执行。

全局内存合并访问

层次结构映射的第一步是将数据从全局内存移动到芯片上。此移动的效率由内存合并访问决定。

全局内存通过 32、64 或 128 字节的事务进行访问。当一个 warp（32 个线程的组）运行加载指令时，硬件会检查每个线程请求的内存地址。如果地址是连续且对齐 (alignment)的，硬件会将这些请求合并为最少数量的事务。

例如，如果线程 $i$ 访问地址 $ptr + i \times 4$（加载一个 32 位浮点数），这 32 个线程会访问一个连续的 128 字节块。这会产生一次内存事务，实现 100% 的总线利用率。

然而，如果访问模式是跨步的，比如访问行主序矩阵中的一列，地址就会分散。线程 0 可能会访问地址 $X$，而线程 1 访问 $X + \text{步长}$。尽管每个线程只需要 4 字节，但硬件可能会为每个线程获取一个完整的 32 字节段来满足请求。这种情况会造成“带宽浪费”，其中有效吞吐量 (throughput)仅为理论峰值的一小部分。

编译器通常通过执行向量 (vector)化加载来处理此问题。线程不再加载单个标量，而是由编译器发出指令（例如 PTX 中的 LD.128），让一个线程加载一个 float4 向量。这增加了每条指令传输的数据量，并有助于分摊内存请求的开销。

共享内存分块

一旦数据从全局内存中获取，就必须存储在可以高效复用的地方。这是共享内存的作用。在矩阵乘法 ($C = A \times B$) 等深度学习 (deep learning)操作中，A 和 B 的每个元素都会被多次访问。如果反复从全局内存中获取它们，会立即使带宽饱和。

编译器通过将 A 和 B 的子矩阵（分块）加载到共享内存中来实现分块（或阻塞）。然后，线程使用这些快速访问的分块计算部分乘积，然后再处理下一个分块。

下面的图表突出显示了这些层级之间带宽和延迟的显著差异，说明了分块缓冲的必要性。

共享内存提供了比全局内存高一个数量级的带宽提升。编译器的用意是最大化每字节加载到该层级的计算操作量。

为优化此流水线，现代编译器（例如使用 MLIR NVGPU 语言的编译器）采用异步复制指令（如 NVIDIA Ampere 和 Hopper 上的 cp.async）。这些指令使 GPU 能够将数据从全局内存直接移动到共享内存，而无需寄存器文件作为中介。这降低了寄存器压力，并允许计算单元在后台加载下一个分块时继续处理当前分块（软件流水线）。

寄存器分配与压力

映射的最后阶段是寄存器文件 (RF)。这里存放着矩阵乘法结果 ($C_{tile}$) 的活动累加器。

寄存器是 GPU 上最稀缺的资源。单个流式多处理器拥有一个庞大的寄存器文件（例如 256 KB），但这个空间分配给所有活动线程。如果编译器生成的核函数每个线程所需的寄存器过多，调度器就必须减少活动 warps 的数量以适应。这种减少被称为寄存器压力。

高寄存器压力会造成低占用率。占用率是活动 warps 数量与支持的最大 warps 数量之比。如果占用率过低，就没有足够的活动线程来掩盖内存操作的延迟。如果一个 warp 因等待全局内存加载而停滞，调度器需要立即切换到另一个就绪的 warp。

将 IR 映射到寄存器时，编译器会进行活跃性分析以查明哪些变量可以共享同一个物理寄存器。如果寄存器需求超出硬件限制，编译器会被迫将寄存器“溢出”到局部内存。尽管名称如此，局部内存物理上位于全局内存（片外），因此溢出会导致严重的性能下降。

双缓冲

为维持连续执行，编译器实现了双缓冲（或多级缓冲）。

1. 加载阶段： 编译器发出指令，将分块 $i+1$ 从全局内存加载到共享内存。
2. 计算阶段： 在加载进行中时，线程将分块 $i$（已在共享内存中）的数据加载到寄存器中并进行计算。
3. 同步： 屏障 (__syncthreads()) 确认分块 $i+1$ 已完全加载，然后循环才重复。

通过将下一个分块的内存延迟与当前分块的算术强度重叠，编译器使得 Tensor Cores 始终保持忙碌。这需要对指令调度进行精细控制，通常需使用 LLVM IR 或直接 PTX 生成等低级中间表示，而非高级源代码。