多级中间表示的原理

传统的单级中间表示（IR），例如通用编译器中使用的（如LLVM IR，GCC的GIMPLE），在直接应用于机器学习编译方面面临重大挑战。主要问题是语义差异：这些低级IR缺乏有效表示和操作高级机器学习构造和计算图所需的抽象。当高级结构已在大量低级指令和内存操作中丢失时，像操作符合并、复杂数据布局转换或机器学习特有的代数简化等优化，很难（甚至无法）可靠地表达和实现。

为弥补这一差异，现代机器学习编译器采用多级中间表示的原理。编译过程不使用单一IR，而是包含一系列不同的IR，每个IR在不同的抽象层次上运行。可以将其想象为从模型的高级、特定于框架的表示，逐步降级到目标硬件上可执行的机器代码的结构化过程。

抽象层级理念

多级IR系统明确定义了对被编译程序的各种视角：

1. 高级抽象： 这一层与TensorFlow或PyTorch等机器学习框架中使用的理念非常相似。它通常将计算表示为对多维张量操作的粗粒度操作图（例如，Convolution、MatMul、ReLU）。此层级的优化侧重于图结构本身，例如合并兼容的操作符、消除冗余计算，或对数据布局做出策略性选择（例如，NCHW对比NHWC）。该表示保留了操作的高级语义信息。

2. 中级抽象： 可以有一个或多个中间层，弥合高级图与低级硬件细节之间的差异。这些层可能将张量计算表示为循环嵌套（仿射循环），引入分块和并行化的思想，或使用更独立于硬件的线性代数原语表示操作。这里的优化通常涉及复杂的循环转换（借助多面体建模等技术，这在第四章会详细讲述）、内存层次优化，以及向特定目标并行化策略迈出的初步步骤。

3. 低级抽象： 这一层类似于传统的编译器IR（如LLVM IR）或特定于硬件的表示（如用于GPU的SPIR-V或特定于目标的汇编前身）。它处理标量操作、向量指令、内存地址、寄存器以及适合直接代码生成的控制流。优化包括指令选择、寄存器分配、指令调度和特定于目标的代码生成细节。

逐步降低：主要机制

这些抽象层级之间的转换通过一个称为逐步降低的过程实现。编译不会一次性完成。相反，表示会逐步地从高级别转换为低级别。

每个降低步骤将一个抽象层级的构造转换为下一个更低层级中等价但更详细的构造。例如：

• 一个高级的Convolution操作可能被降低为一组实现卷积算法的嵌套循环。
• 这些嵌套循环（可能以仿射方言表示）随后可能被降低为LLVM类IR中的标量和向量操作组合。
• 最后，LLVM类IR被降低为特定于目标的机器代码。

值得注意的是，优化通常在降低发生之前应用于特定的抽象层级内部。图合并发生在图层级；循环分块发生在循环/张量层级；寄存器分配发生在低层级。这种关注点分离使编译器设计更具模块化和易于管理。优化可以针对其最直接表达和有效的表示进行设计和实现。

这是通过多级IR系统中不同抽象层级逐步降低的视图，这通常在MLIR等框架中通过不同的方言实现。优化（椭圆）通常在降低到下一个层级之前，在特定层级内部进行。

多级方法的优点

使用多个IR层级并结合逐步降低，为构建先进的机器学习编译器提供了多个重要优势：

1. 关注点分离： 高级优化（如算法选择或合并）与低级细节（如寄存器分配或特定指令选择）相互分离。这简化了优化遍的设计和实现。
2. 定向优化： 优化可以在最合适的抽象层级应用，在该层级中，必要的信息易于获取，并且转换最易于表达。
3. 模块化和可重用性： 编译器的不同组件（例如，从各种框架导入的前端、不同的优化策略、面向多样化硬件的后端）可以通过特定层级中定义明确的IR接口进行交互。
4. 表示能力： 该系统能很好地表示不同阶段的相关信息，从图拓扑和高级操作，直到特定硬件的约束。
5. 可扩展性： 添加对新硬件或新高级操作的支持，通常涉及定义新的降低路径或在特定现有层级添加优化，而不是重新设计整个编译器。

像MLIR（多级中间表示）这样的框架，我们接下来将详细研究它，提供了定义和管理这些多层级（在MLIR中称为方言）及其间降低过程的基础设施。理解多级抽象和逐步降低的这一原理，是了解现代机器学习编译器如何应对为多样化和高要求硬件目标优化模型复杂性的必要条件。