中间表示（IR）的作用

机器学习框架为设计复杂模型提供了便捷接口，但它们操作的抽象层次与硬件执行的物理现实相距甚远。当你在PyTorch中编写一个简单的矩阵乘法时，框架看到的是对两个张量对象的一次函数调用。然而，硬件只识别寄存器、内存地址和基本指令集。中间表示（IR）弥合了这一差距。

IR在机器学习编译器中的主要作用是将模型定义与其执行环境分离。这种分离实现了被称为 $M \times N$ 解决方案的模块化设计。如果没有统一的IR，跨 $N$ 种不同硬件后端（NVIDIA GPU、ARM CPU、TPU）支持 $M$ 种不同框架（PyTorch、TensorFlow、JAX）将需要构建 $M \times N$ 个独立的编译器。通过就一种通用的中间格式达成一致，编译器工程师只需构建 $M$ 个前端将框架转换为IR，以及 $N$ 个后端将IR转换为机器码。这将工程工作量减少到 $M + N$。

保留高级语义

在GCC或LLVM等传统软件编译器中，IR通常表示低级标量操作。例如，一个遍历数组的循环被表示为一系列指针运算、比较和跳转。尽管这对于通用代码是高效的，但它会破坏机器学习操作的高级意图。

以2D卷积为例。在高级ML IR中，这被表示为一个单一的原子节点：conv2d(input, weight)。这种表示保留了操作的语义含义。如果编译器立即将其下沉为嵌套循环和指针运算，它将失去执行特定领域优化的能力。识别 conv2d 节点并用高度优化的cuDNN内核替换它，比分析由七个通用 for 循环组成的嵌套来确定它们共同表示一个卷积要容易得多。

下图展示了IR如何作为中心枢纽，在下沉到硬件特定代码之前保持高级语义。

计算从框架到硬件的流程。高级IR捕捉意图，而低级IR处理实现细节。

多级方言方法

现代机器学习编译器，例如基于MLIR（多级中间表示）基础设施的编译器，不依赖单一的IR格式。相反，它们使用一种表示层次结构，通常称为“方言”。这种层次结构解决了优化方面的不足，通过让编译器在最适合特定优化任务的抽象级别上工作。

1. 图级别： 在最高级别，IR类似于依赖图。节点对应于张量代数操作。这里的优化侧重于图的结构。例如，代数简化（如 $A \times 0 = 0$）或算子融合（将乘法和加法合并到一个内核中）在此处进行。
2. 张量/循环级别： 一旦结构性优化完成，编译器会“下沉”表示。抽象的 conv2d 节点被扩展为如何计算它的明确定义，通常涉及嵌套循环和标量计算。循环分块、向量化和内存分配在此处发生。
3. 硬件级别： 最后，IR被下沉到接近汇编的格式，将操作映射到目标设备上可用的特定内部函数，例如GPU上的Tensor Cores或CPU上的AVX-512指令。

静态分析与形状推断

IR的另一个作用是促进静态分析。Python是一种动态语言；变量类型和张量形状可以在运行时改变。然而，为了生成高效的机器码，编译器需要确定性。

当模型被导入IR时，编译器执行形状推断和类型检查。它通过图传播元数据，以确定每个中间张量的内存需求。

$$\text{输出}_{\text{尺寸}} = \frac{\text{输入}_{\text{尺寸}} - \text{核}_{\text{尺寸}} + 2 \times \text{填充}}{\text{步幅}} + 1$$

通过使用从算子定义中获得的公式提前计算这些维度，编译器可以静态分配内存缓冲区。这消除了运行时内存管理的开销，运行时内存管理在即时执行模式中是延迟的一个主要来源。

不变性与副作用

在大多数机器学习编译器IR中，图被构建为纯数据流图。这意味着操作没有副作用。一个算子接受输入并产生输出，而不修改全局状态。这种不变性对并行化很有帮助。如果IR中的两个节点不相互依赖数据，编译器可以安全地安排它们在不同的流或核心上同时执行，而无需担心竞态条件。

这与Python的命令式风格形成对比，在Python中，变量可以被覆盖，或者列表可以在原位修改。IR有效地将逻辑“冻结”为一个静态快照，为积极的重写和优化提供稳固的支撑。