代码生成后端

编译器的中间表示（IR）精确地描述了如何执行计算，其中包含循环分块、向量化宽度和内存访问模式的具体指令。但此表示仍在编译器框架内部。要在物理硬件上运行这些逻辑，编译器必须将这些内部结构转换为目标处理器能理解的格式。这个过程称为代码生成，或简称“codegen”。

代码生成阶段是神经网络的抽象数学运算与 CPU、GPU 或加速器的具体指令集之间的桥梁。现代机器学习编译器不是为每种架构手动编写汇编代码，而是借助成熟的后端框架来处理最后的转换步骤。

从调度到低层 IR

在生成二进制代码之前，编译器会执行一个“降低”阶段。处理张量和算子的高层图，被降低为一种通常称为低层 IR 的基于指针的表示。在此状态下，“张量”这一想法不复存在。相反，编译器看到的是扁平内存缓冲区、明确的循环索引和基本算术运算。

例如，一个高层矩阵乘法运算 $C = A \times B$ 被转换为三重嵌套循环结构。在此阶段，编译器会注入自动调优阶段做出的优化决定。如果自动调优器决定将内层循环展开 4 次，低层 IR 会通过在循环体内重复算术指令四次来明确体现这一点。

这种分离使机器学习编译器能够专注于循环转换，而无需担心寄存器分配或指令调度，这些任务被委派给下游后端。

数据从优化调度到可执行机器码的流程，区分了基于 LLVM 和基于源代码的生成路径。

CPU 的 LLVM 路径

对于 CPU 目标（x86、ARM、RISC-V），大多数机器学习编译器使用 LLVM（低层虚拟机）。LLVM 提供了一种标准化的中间表示和一套强大的工具，可将该表示编译成高效的汇编代码。

机器学习编译器遍历其自身的低层 IR，并构建对应的 LLVM 模块。这包括将基本运算映射到 LLVM 指令：

• 类型： 机器学习编译器中的 32 位浮点数映射到 LLVM 中的 float。
• 控制流： 循环被转换为基本块和分支指令。
• 内联函数： 向量化运算被映射到 LLVM 向量内联函数，LLVM 后端随后将其转换为 AVX-512 (Intel) 或 NEON (ARM) 指令。

通过针对 LLVM，机器学习编译器获得了通用编译器优化领域数十年的工程成果。LLVM 处理寄存器分配、指令调度和死存储消除。机器学习编译器只需在 LLVM IR 中正确表达程序的意图即可。

考虑一个简化示例，说明向量加法在降低到 LLVM IR 时可能的样子。其语法冗长且带有类型，类似于汇编，但具有无限虚拟寄存器：

$$\%res = \text{fadd } <4 \times \text{浮点数}> \%vecA, \%vecB$$

在这里，该操作同时作用于一个包含 4 个浮点数的向量。机器学习编译器生成此通用指令，LLVM 决定哪个具体的 CPU 指令（例如 vaddps）最适合目标处理器来执行。

GPU 的源代码生成

为 GPU 生成代码通常遵循不同的路径。虽然 NVIDIA GPU 有像 NVVM（基于 LLVM）这样的后端，但许多机器学习编译器使用一种称为“源到源编译”的技术。编译器不是直接生成二进制文件，而是生成一段类似 C 语言的源代码字符串。

对于 NVIDIA GPU，编译器生成 CUDA C 代码。对于 AMD GPU 或移动处理器，它可能会生成 OpenCL 或 Vulkan 计算着色器。这种方法很实用，因为 GPU 驱动程序会调用其自身的激进编译器，这些编译器已针对特定代图形卡进行了高度调优。

该过程通常包括：

1. 内核提取： 识别将在 GPU 上运行的循环嵌套。
2. 语法转换： 将内部 IR 表达式转换为字符串。例如，一个内部的 Select(condition, true_val, false_val) 节点可能被转换为字符串 condition ? true_val : false_val 或像 fminf 这样的专用内联函数。
3. 运行时编译： 在程序执行期间，生成的源代码字符串被传递给运行时编译器驱动程序（例如用于 CUDA 的 NVRTC）。这会将字符串编译成 PTX（并行线程执行）代码或二进制 cubin 文件。

这种方法使得机器学习编译器能够轻松检查生成的源代码，这对于调试性能问题很有用。如果自动调优器选择块大小为 128，您将在生成的 C 代码中清楚地看到 __launch_bounds__(128) 或类似的指令。

处理主机-设备分离

代码生成很少是关于生成单个代码块。它通常需要生成两个不同的部分：主机代码和设备代码。

设备代码是在加速器上运行的计算密集型内核（例如，矩阵乘法）。 主机代码在 CPU 上运行，负责：

1. 在设备上分配内存。
2. 将数据从 CPU RAM 移动到设备内存。
3. 配置函数参数（指针、形状）。
4. 启动内核。

当机器学习编译器生成代码时，它将这两个组件打包在一起。主机代码通常被编译成共享库（例如 .so 或 .dll 文件），而设备代码则作为二进制数据块或字符串嵌入在该库中。

函数调用期间时间花费的细分。代码生成器必须优化设备执行，同时最大限度地减少由主机代码管理的参数打包和启动的开销。

JIT 与 AOT 终结化

代码生成后端的最终输出根据部署策略而定。

在**提前（AOT）**场景中，后端将生成的机器码写入磁盘，作为共享对象文件。此库随后可以链接到 C++ 应用程序或移动应用中，使得模型能够在没有机器学习编译器框架的情况下运行。这非常适用于内存和存储有限的边缘设备。

在**即时（JIT）**场景中（这在 Python 实验中很常见），后端在内存中生成代码。它分配一个可执行内存页，直接将机器指令写入其中，并返回一个函数指针。Python 随后通过 ctypes 或类似的外部函数接口调用此函数指针。这在自动调优过程中实现了即时反馈，因为编译器可以在一个紧密的循环中生成、运行和测量调度。

掌握这些后端，您就能精确查看硬件执行的内容。您不必将 matmul 函数视为黑箱，而是可以检查 LLVM IR 或 CUDA 源代码，以确认向量化已启用且内存访问模式得到优化。