AOT 与 JIT 编译

部署环境决定了您的机器学习模型何时以及如何从高级图转化为可执行机器码。这种转化时机是区分即时 (JIT) 编译和提前 (AOT) 编译的主要因素。理解这种区别对于针对特定目标优化模型非常重要，例如边缘设备上的低延迟推理或云集群上的高吞吐量训练。

编译时间线

AOT 和 JIT 的核心区别在于编译栈相对于应用运行时何时执行。

在提前 (AOT) 工作流中，编译过程离线进行。您通常在训练完成后，获取模型定义，并通过编译器栈处理，生成一个独立的二进制文件或库文件。此产物包含运行模型所需的优化机器码，但不包含编译器本身。在运行时，应用只需加载并执行此二进制文件。

在即时 (JIT) 工作流中，编译在程序运行时动态发生。框架（如 PyTorch 或 JAX）会监控执行。模型或函数首次被调用时，JIT 编译器会捕获图，对其优化，并即时生成机器码。该代码随后被缓存并执行。后续调用会跳过编译步骤，直接使用缓存的内核。

下图展现了这些工作流的结构区别。

AOT 和 JIT 工作流的比较。在 AOT 中，编译器在部署前运行一次。在 JIT 中，编译器是执行循环的一部分。

即时 (JIT) 编译

JIT 编译是模型研究和训练期间的主流做法，因为它保留了灵活性。当使用 torch.compile 或 XLA（加速线性代数）等工具时，框架会尝试融合运算符并优化执行，而不强迫用户完全放弃 Python。

JIT 的主要特点是预热开销。数据首次流经模型时，延迟会显著升高，因为系统正在忙于编译图。

考虑 JIT 编译函数的性能表现。系统需要分析输入张量形状，追踪操作，应用优化，并生成 GPU 内核。如果后续调用中输入形状发生变化（动态形状），JIT 编译器可能需要触发重新编译，以生成针对新维度优化的内核。

JIT 的优点

1. 易用性： 它与 Python 控制流自然结合。在编译发生之前，您通常可以使用标准 Python 调试工具。
2. 动态优化： 编译器在运行时了解确切的形状和数据类型，使其能够为特定实例生成高度专业化的代码。

JIT 的缺点

1. 运行时开销： 编译器基础设施必须存在于生产环境。这会增加内存占用和依赖项大小。
2. 不可预测的延迟： 如果新的执行路径或张量形状触发重新编译，用户请求可能会经历突然的延迟（抖动）。

提前 (AOT) 编译

AOT 编译将机器学习模型视为 C++ 程序。目标是生成一个专用可执行文件，只依赖于最小运行时，而不是完整的训练框架。这种做法是部署到移动设备、嵌入式系统或像 FPGA 这样的专用加速器的标准方式。

在 AOT 环境中，您通常必须提供静态形状或为动态维度定义严格的边界。编译器进行大量分析，预先确定内存需求。这使其能够静态分配内存，避免执行期间动态内存管理的开销。

AOT 的优点

1. 可预测的性能： 由于所有编译都在事前发生，不会有运行时意外。推理时间从首次运行起保持一致。
2. 可移植性： 生成的产物通常是一个共享库（.so 或 .dll），无需在目标机器上安装 Python 或 PyTorch，即可从 C、C++、Java 或 Rust 调用。
3. 硬件效率： AOT 编译器可以应用需要长时间计算的激进优化，因为这段编译时间不会影响最终用户。

AOT 的缺点

1. 刚性： 模型图必须完全静态。动态控制流（基于张量数据的循环和条件语句）难以捕获，且通常需要重写模型部分。
2. 复杂的工具链： 设置交叉编译环境（例如在 x86 服务器上为 ARM 移动芯片编译）会增加构建流程的复杂性。

性能比较

为清晰显示这些策略的作用，我们可以查看一系列推理请求的执行时间。下图比较了标准解释器（如标准 PyTorch eager 模式）、JIT 编译器和 AOT 编译模型。

10 次运行的延迟比较。注意 JIT 由于编译开销导致的初始峰值，随后性能超越解释器。AOT 从第一步起提供持续的低延迟。

选择合适的策略

AOT 和 JIT 之间的选择很少是关于哪个能产生更快的内核，现代编译器通常对两者使用相同的底层优化过程。具体的选择应根据部署限制来定。

何时使用 JIT：

• 您处于研究和实验阶段。
• 您拥有充足的内存和 CPU 资源（例如云服务器）。
• 您的模型大量依赖动态 Python 功能或变化的输入形状。

何时使用 AOT：

• 您要部署到资源有限的边缘设备（移动、物联网）。
• 启动时间非常重要（例如自动制动系统不能等待编译）。
• 您需要将模型集成到 Python 运行时不可用的更大 C++ 应用中。

在后续章节中，我们将主要查看这些过程生成的中间表示 (IR)。无论是即时触发还是提前触发，编译器最终都会将图转化为这种 IR 来实现其功能。