ML 编译器并非单一的整体程序。相反，它是一个模块化流程，旨在逐步将高层数学描述转换为低层机器码。虽然像 GCC 或 Clang 这样的标准编译器将 C++ 转换为汇编代码，但 ML 编译器将计算图（例如 ResNet 或 Transformer 模型）转换为优化过的核二进制文件。为处理将线性代数映射到各种硬件后端（从 CPU、GPU 到 TPU、FPGA 等）的复杂之处，ML 编译器采用了分层架构。这种设计将模型表示问题与硬件执行问题分离开来。我们通常将编译器堆栈分为三个主要阶段：前端、中端（优化器）和后端。digraph G { rankdir=TB; node [fontname="Helvetica", shape=box, style=filled, color="#dee2e6"]; edge [fontname="Helvetica", color="#868e96"]; subgraph cluster_0 { label = "ML 框架"; style=dashed; color="#adb5bd"; Framework [label="PyTorch / TensorFlow\n模型定义", fillcolor="#e9ecef"]; } subgraph cluster_1 { label = "ML 编译器堆栈"; color="#ced4da"; style=solid; Frontend [label="前端\n(导入器与解析器)", fillcolor="#74c0fc"]; HighIR [label="高层 IR\n(计算图)", fillcolor="#bac8ff"]; Optimizer [label="优化器\n(优化阶段与分块)", fillcolor="#74c0fc"]; LowIR [label="低层 IR\n(循环与索引)", fillcolor="#bac8ff"]; Backend [label="后端\n(代码生成)", fillcolor="#74c0fc"]; } subgraph cluster_2 { label = "硬件"; style=dashed; color="#adb5bd"; Binary [label="可执行二进制文件\n(ELF / CUBIN)", fillcolor="#e9ecef"]; } Framework -> Frontend; Frontend -> HighIR; HighIR -> Optimizer; Optimizer -> LowIR; LowIR -> Backend; Backend -> Binary; }数据流通过机器学习编译器堆栈的核心组成部分前端：获取与转换前端是用户使用的深度学习框架与编译器内部之间的接口。它的主要职责是获取模型定义并将其转换为编译器可以处理的格式。此过程常被称为“降低”。当您运行 PyTorch 模型时，该框架会维护一个动态操作图。前端采用以下两种方法之一捕获此图：追踪： 使用模拟数据运行模型并记录执行的操作。脚本/解析： 分析 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）进行大量定制的阶段。后端执行两个核心任务：资源分配： 它确定如何将特定的硬件寄存器和共享内存区域分配给 IR 中的变量。代码生成：： 它将 IR 转换为最终的机器码。许多现代 ML 编译器，如 TVM 或 XLA，不直接生成二进制机器码。相反，它们为更低层级的编译器框架（例如用于 CPU 的 LLVM 或用于 NVIDIA GPU 的 NVVM）生成代码。这使得 ML 编译器能够专注于张量相关的优化，同时依靠 LLVM 等成熟工具来处理指令调度和寄存器分配。多层 IR 的作用标准软件编译器通常使用单一的主要 IR（如 LLVM IR）。然而，机器学习中的抽象层级差异太大，无法通过单一格式弥合。对图融合有利的表示形式通常不利于循环调度。因此，现代 ML 编译器的结构由多层中间表示决定。代码流经一系列方言，从高层抽象（张量、图）开始，到低层抽象（指针、内存偏移）结束。{"layout": {"title": "ML 编译中的抽象层级", "xaxis": {"title": "抽象层级", "showticklabels": false}, "yaxis": {"title": "信息密度", "showticklabels": false}, "height": 400, "width": 600, "showlegend": true}, "data": [{"x": [1, 2, 3], "y": [90, 50, 10], "mode": "lines+markers", "name": "领域特定信息（例如，卷积）", "line": {"color": "#74c0fc", "width": 4}}, {"x": [1, 2, 3], "y": [10, 50, 90], "mode": "lines+markers", "name": "硬件特定信息（例如，寄存器）", "line": {"color": "#fa5252", "width": 4}}]}随着编译器降低代码（在 X 轴上向右移动），领域特定语境减少，而硬件特定细节增多。在下一节中，我们将对比此编译流程的触发方式，具体比较提前编译 (AOT) 与即时编译 (JIT) 执行。
