ML 编译器的结构

ML 编译器并非单一的整体程序。相反，它是一个模块化流程，旨在逐步将高层数学描述转换为低层机器码。虽然像 GCC 或 Clang 这样的标准编译器将 C++ 转换为汇编代码，但 ML 编译器将计算图（例如 ResNet 或 Transformer 模型）转换为优化过的核二进制文件。

为处理将线性代数映射到各种硬件后端（从 CPU、GPU 到 TPU、FPGA 等）的复杂之处，ML 编译器采用了分层架构。这种设计将模型表示问题与硬件执行问题分离开来。

我们通常将编译器堆栈分为三个主要阶段：前端、中端（优化器）和后端。

数据流通过机器学习编译器堆栈的核心组成部分

前端：获取与转换

前端是用户使用的深度学习框架与编译器内部之间的接口。它的主要职责是获取模型定义并将其转换为编译器可以处理的格式。此过程常被称为“降低”。

当您运行 PyTorch 模型时，该框架会维护一个动态操作图。前端采用以下两种方法之一捕获此图：

1. 追踪： 使用模拟数据运行模型并记录执行的操作。
2. 脚本/解析： 分析 Python 抽象语法树 (AST) 以构建控制流的静态表示。

捕获后，前端将特定于框架的运算符（如 torch.nn.Conv2d）映射到编译器自己的高层中间表示 (IR)。此 IR 通常是定向无环图 (DAG)，其中节点表示数学张量操作，边表示数据依赖。

中端：优化引擎

中端是编译器应用转换以提升性能而不改变数学结果的阶段。此阶段平台无关或弱平台相关。它在高层 IR 上进行操作。

此处的优化分为两类：

图级别优化 这些修改会改变计算图的结构。一个典型例子是运算符融合。如果图中包含矩阵乘法紧接着一个 ReLU 激活：

$Y = \text{ReLU}(X \times W + b)$

编译器将这些融合到一个单独的核中，避免了在为 ReLU 再次读取之前将 $(X \times W + b)$ 的中间结果写回主内存的需要。我们将在第 3 章中细致地介绍这一点。

张量级别优化 图优化完成后，编译器将图进一步“降低”为张量或循环 IR。在此阶段，“卷积”等高层操作被分解为显式的嵌套循环。优化器接着应用以下技术：

• 循环分块： 将大循环分解为更小的块以适应缓存。
• 向量化： 将标量操作转换为 SIMD（单指令多数据）指令。

后端：代码生成

后端负责将优化后的 IR 映射到目标硬件的特定指令集架构 (ISA)。这是唯一需要为特定设备（例如 NVIDIA H100 GPU 与 ARM Cortex CPU）进行大量定制的阶段。

后端执行两个核心任务：

1. 资源分配： 它确定如何将特定的硬件寄存器和共享内存区域分配给 IR 中的变量。
2. 代码生成：： 它将 IR 转换为最终的机器码。

许多现代 ML 编译器，如 TVM 或 XLA，不直接生成二进制机器码。相反，它们为更低层级的编译器框架（例如用于 CPU 的 LLVM 或用于 NVIDIA GPU 的 NVVM）生成代码。这使得 ML 编译器能够专注于张量相关的优化，同时依靠 LLVM 等成熟工具来处理指令调度和寄存器分配。

多层 IR 的作用

标准软件编译器通常使用单一的主要 IR（如 LLVM IR）。然而，机器学习中的抽象层级差异太大，无法通过单一格式弥合。对图融合有利的表示形式通常不利于循环调度。

因此，现代 ML 编译器的结构由多层中间表示决定。代码流经一系列方言，从高层抽象（张量、图）开始，到低层抽象（指针、内存偏移）结束。

随着编译器降低代码（在 X 轴上向右移动），领域特定语境减少，而硬件特定细节增多。

在下一节中，我们将对比此编译流程的触发方式，具体比较提前编译 (AOT) 与即时编译 (JIT) 执行。