当运行一个 CUDA Kernel 时发生了什么？

背景

在高性能计算和 AI 领域，CUDA 是 NVIDIA GPU 编程的核心框架。然而，对于大多数开发者而言，CUDA 内核（Kernel）的执行过程往往是一个“黑盒”。我们通常只关心代码逻辑是否正确，而忽略了从主机代码编译到 GPU 硬件执行之间复杂的底层机制。

本文源自 Hacker News 上的一篇深度技术文章，旨在通过一个具体的向量加法示例，深入剖析 CUDA 程序从源代码到最终在 RTX 4090 上执行完毕的全过程。这不仅涉及编译器前端和后端的工作，还涵盖了驱动程序、内存管理、指令集架构（ISA）转换以及硬件层面的寄存器分配和内存访问。

文章指出，仅仅执行一行简单的内核启动代码，背后可能涉及数千万条 CPU 指令、多个设备文件操作、数百次 ioctl 调用以及一次内存映射的门铃寄存器写入。这种“透明度”的提升，得益于 AI 辅助工具带来的“可读性过渡”（legibility transition），使得开发者能够以前所未有的深度理解计算机系统的运作原理。

核心内容

1. 编译过程：从源代码到机器码

CUDA 程序的编译并非单一过程，而是由 nvcc 驱动的一系列编译器协作完成的。nvcc 是一个驱动程序，它调用多个其他编译器并将输出组合在一起。

• 主机代码处理：主机代码（Host Code）被发送到标准的主机编译器（如 GCC 或 Clang）。
• 设备代码处理：设备代码（Device Code，即 __global__ 函数）经过以下步骤：
  1. cicc：基于 LLVM 的编译器将 CUDA C++ 代码转换为 PTX（Parallel Thread Execution）。
  2. ptxas：将 PTX 转换为特定架构的 SASS（Streaming ASSembler）。

PTX：虚拟指令集架构

PTX 是一种虚拟 ISA，具有以下特点：

• 无限寄存器：它拥有无限多的类型化寄存器，不关心硬件实际拥有的寄存器数量。
• 设备无关性：PTX 是设备无关的，因此生成的代码较为冗长。例如，计算一个内存地址可能需要三条 PTX 指令，因为 CUDA 指针默认是“通用”的（可能指向全局、共享或本地内存），需要显式转换。
• 示例：在 PTX 中，计算索引 i 并加载数据涉及 mad.lo.s32、setp.ge.s32、cvta.to.global.u64、mul.wide.s32 和 add.s64 等多条指令。

SASS：硬件特定指令集

ptxas 将 PTX 编译为针对特定架构（如 sm_89，对应 Ada Lovelace 架构）的 SASS。SASS 代码更紧凑，体现了硬件特性：

• 寄存器折叠：PTX 中的十几个虚拟寄存器在 SASS 中可能只映射到七个物理寄存器（如 R1-R9）。
• 指令融合：PTX 中的 mul.wide 和 add 序列在 SASS 中融合为一条 IMAD.WIDE 指令。
• 特殊寄存器：S2R 指令用于将硬件维护的特殊寄存器（如 SR_CTAID.X 对应 blockIdx.x，SR_TID.X 对应 threadIdx.x）复制到普通寄存器中以便进行算术运算。
• 常量内存：内核参数（指针 a, b, c 和大小 n）以及启动几何信息存储在常量内存的 Bank 0 中。这是为了利用常量缓存的广播特性，使得网格中的所有线程可以在一次操作中获取相同的参数。

二进制文件结构

• Cubin：包含 SASS 的 ELF 文件，包含符号表和 .text.vadd 节。
• Fatbin：将 Cubin 和 PTX 捆绑在一起的格式。如果目标 GPU 架构不支持预编译的 SASS，驱动程序可以在加载时 JIT 编译 PTX。
• 主机可执行文件：最终的可执行文件包含主机代码、嵌入的 Fatbin（其中包含压缩的 PTX 和 SASS）。

2. 内核启动：主机如何触发 GPU

当调用 vadd<<<4096, 256>>>(...) 时，发生了一系列复杂的交互：

1. 主机存根（Host Stub）：编译器生成的主机代码负责打包内核参数。这些参数被放置在特定的内存位置，与 SASS 中常量内存的偏移量（如 0x160, 0x168 等）相对应。
2. QMD（Queue Management Descriptor）：驱动程序创建一个 QMD 结构，其中包含内核的参数、启动几何信息以及指向代码的指针。
3. 提交到 GPU：驱动程序通过 ioctl 调用将 QMD 提交给 GPU 的硬件队列。
4. 门铃寄存器（Doorbell Register）：最后，驱动程序写入一个内存映射的门铃寄存器，通知 GPU 硬件有新的工作负载需要处理。
5. GPU 执行：GPU 硬件读取 QMD，分配线程束（Warps），加载指令，执行计算，并将结果写回全局内存。

3. 执行细节：从 Warps 到结果

• 线程映射：4096 * 256 = 1,048,576 个线程，每个线程处理一个浮点数。
• 内存访问：每个线程计算其全局索引 i，检查边界，计算内存地址，加载 a[i] 和 b[i]，执行加法，并将结果存储到 c[i]。
• 结果验证：最后，数据从设备复制回主机，并打印验证结果。

关键要点

• 编译流水线复杂：CUDA 编译涉及 nvcc 协调 cicc（生成 PTX）和 ptxas（生成 SASS）等多个组件。
• PTX 是虚拟 ISA：PTX 具有无限寄存器且设备无关，代码较为冗长，旨在提供向前兼容性。
• SASS 是硬件指令：SASS 是特定于架构的机器码，经过寄存器分配、指令融合等优化，效率更高。
• 常量内存优化：内核参数存储在常量内存 Bank 0，利用广播特性减少内存带宽占用。
• Fatbin 机制：最终二进制文件包含 SASS 和压缩的 PTX，PTX 作为后备方案，允许在旧硬件上 JIT 编译。
• 启动开销巨大：启动一个内核涉及大量 CPU 指令、设备文件操作、ioctl 调用和内存映射寄存器写入。
• QMD 的作用：QMD 结构是驱动程序与 GPU 硬件之间的关键接口，包含所有必要的执行上下文信息。
• 可读性提升：借助现代工具和 AI 辅助，开发者可以深入理解底层硬件行为，打破“黑盒”状态。

意义与影响

理解 CUDA 内核执行的底层细节对于高性能计算开发者至关重要。首先，它揭示了性能优化的潜在空间。例如，了解常量内存的广播特性可以帮助开发者更有效地组织参数；理解寄存器分配和指令融合有助于编写更高效的 PTX 或 SASS 代码。

其次，这种透明度对于调试和性能剖析具有重要意义。当遇到性能瓶颈或硬件错误时，深入理解从主机到设备的整个链路可以帮助开发者快速定位问题，而不是盲目地调整参数。

最后，文章提到的“可读性过渡”反映了 AI 时代技术文档和工具链的演变。随着 AI 辅助工具的发展，理解复杂系统（如 GPU 架构）的门槛正在降低，这将促进更多开发者参与到高性能计算和系统级优化的研究中，推动整个行业的技术进步。