JIT系统中的中间表示

即时（JIT）编译从根本上改变了优化和代码生成 何时 发生，将这些阶段从离线的预编译（AOT）过程转移到实际的运行时执行。这种时间上的变化对JIT编译器中使用的中间表示（IR）提出了独特且严格的要求。与AOT场景不同，AOT场景下编译时间不那么要紧，JIT系统中的IR设计不仅要重视表达能力，还要重视其构建、操作和降级的速度，因为这些直接影响用户感受到的应用程序延迟。

JIT中间表示的核心要求

机器学习 (machine learning)JIT编译器的效率在很大程度上取决于其IR的能力。有几个特性非常重要:

1. 多层抽象与灵活性： JIT编译通常首先捕获接近源框架级别的操作（例如，PyTorch中的Python操作或TensorFlow图节点）。IR必须忠实地表示这些高级结构，包括动态控制流和框架特定的语义。随后，它需要支持通过可能多个中间级别进行渐进式降级，从而实现图级优化（如融合）、张量/循环级转换（如平铺），并最终映射到低级硬件指令。这要求IR能够表示不同粒度的计算。
2. 操作效率： 由于编译发生在执行期间，构建、遍历和转换IR所花费的时间是重要的开销。IR数据结构必须轻量，并允许快速模式匹配、重写和分析。复杂或操作缓慢的IR可能会抵消JIT专用化带来的性能优势。
3. 动态信息的表示： JIT编译的一个主要目的就是使用运行时信息。IR必须对表示动态属性提供一流的支持，最常见的是张量形状和数据类型，这些在编译开始时可能无法完全获知。这通常涉及在IR结构本身中使用符号维度、约束或类型占位符。这种动态信息是运行时专用化的主要促成因素。
4. 可扩展性： 机器学习随着新算子、硬件目标和优化技术的出现而迅速发展。JIT系统的IR需要具有可扩展性，以便直接添加新操作、数据类型，甚至全新的抽象级别（通常通过方言机制，如MLIR中所见，在第2章讨论），而无需对编译器的核心基础设施进行根本性更改。
5. 高效降级途径： 尽管灵活性非常重要，但IR还必须提供明确且高效的途径，将高级表示降级为可执行代码。这涉及一系列转换过程，它们逐步减少抽象，解决动态属性（如果可能），并将操作映射到特定硬件的结构或像LLVM IR这样的低级IR。

IR中动态性的处理

JIT IR处理的一个核心难题是表示最初未知或可变的信息。张量形状是典型的例子。AOT编译器可能要求所有张量维度都是静态常量。然而，JIT编译器经常遇到某些维度取决于运行时输入的张量。

IR可以使用以下机制来处理：

• 符号维度： 用符号表示未知维度（例如，tensor<Nx1024xf32>）。
• 形状函数/约束： 将操作与定义基于输入形状的输出形状的函数或约束关联，即使某些输入维度是符号。例如，矩阵乘法 C = matmul(A, B)，其中 A 的形状为 $(M, K)$，B 的形状为 $(K, N)$，其IR表示将编码 C 的形状为 $(M, N)$，无论 $M$、$K$ 或 $N$ 是具体值还是符号。
• 类型细化： 随着编译的进行或运行时信息变得可用，JIT编译器可以细化IR中的类型和形状，用从追踪或输入守卫派生的具体值替换符号维度。

这种表示和操作部分指定信息的能力，是实现运行时专用化的基本条件，JIT会生成针对特定执行轨迹中遇到的 实际 张量形状优化的代码。

分层抽象和降级流程

为了平衡高级语义与低级优化需求，JIT IR通常采用分层或多方言的方法。JIT中的典型流程可能如下所示：

JIT编译器中IR降级阶段的视图。编译从捕获的框架操作开始，通过逐步降低抽象的IR实现不同的优化，最终实现目标代码生成。形状专用化通常发生在从高级IR到中级IR的转换期间。

这种分层方法允许在最合适的抽象级别应用优化：例如，在高级IR上进行图融合，在中级IR上进行循环平铺，在低级IR上进行指令调度。JIT编译器协调这些层之间的转换（降级）。

与图获取的关系（追踪 vs. 脚本化）

捕获用户模型的方法影响JIT IR的初始形式：

• 追踪： 当使用示例输入追踪模型时，JIT会观察执行的操作序列。IR通常构建为数据流图，其中节点表示已执行的操作，边表示张量依赖关系。处理追踪过程中遇到的控制流（如条件语句或循环）需要仔细的IR设计，以准确捕获分支逻辑和潜在的形状变化。IR必须表示追踪期间所采取的特定路径，同时可能嵌入 (embedding)处理其他路径所需的信息，如果重新编译发生。
• 脚本化： 当使用宿主语言的受限子集（例如TorchScript或TensorFlow的带有autograph的tf.function装饰器）定义模型时，JIT会直接解析此代码。生成的IR通常更接近抽象语法树（AST），或包含显式的控制流结构（例如MLIR术语中的scf.if或scf.for）。与追踪相比，这为编译器提供了更明确的程序结构进行分析。

在这两种情况下，初始IR都捕获了程序结构，随后利用运行时可用的动态上下文 (context)进行细化和优化。

JIT与AOT IR要求的对比

尽管JIT IR与AOT编译器IR共享一些基本思想（如SSA形式、操作语义），但它们在不同的限制下运行。AOT编译器可以承担昂贵的分析和转换，因为编译时间是在离线进行的。它们通常依赖于关于形状和类型的详细静态信息。

相反，JIT IR必须：

• 生成速度快： 通过追踪或脚本化捕获模型必须快速。
• 优化速度快： 在运行时应用的核心优化计算成本必须低。更复杂的优化可能会被推迟或根据执行次数自适应应用（请参阅第7.6节）。
• 适应性强： 设计为能够很好地结合运行时信息，并在必要时触发重新优化或重新编译（例如，如果张量形状在不同调用之间发生显著变化）。

TensorFlow的XLA使用HLO（高级优化器IR），它是基于图的，适用于积极的融合，而PyTorch的TorchScript使用一种IR，在降级之前最初保留更多Python风格的语义。两者都旨在平衡表示能力与JIT编译的性能要求，体现了此处讨论的原则。这些系统将在第7.7和7.8节中进行更详细的考察。

总而言之，中间表示是任何机器学习 (machine learning)JIT编译系统的根本。它的设计必须权衡忠实地表示高级的、可能动态的程序语义与实现高效的、运行时敏感的优化和代码生成。处理动态信息、支持多层抽象以及促进快速操作的能力，是有效JIT IR的决定性特点。