机器学习框架为设计复杂模型提供了便捷接口,但它们操作的抽象层次与硬件执行的物理现实相距甚远。当你在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如何作为中心枢纽,在下沉到硬件特定代码之前保持高级语义。digraph G { rankdir=TB; node [shape=box, style="filled", fillcolor="#f8f9fa", fontname="Helvetica", color="#dee2e6"]; edge [fontname="Helvetica", color="#adb5bd"]; subgraph cluster_frontend { label = "前端(框架)"; style=dashed; color="#ced4da"; PyTorch [label="PyTorch模型"]; TF [label="TensorFlow模型"]; } subgraph cluster_ir { label = "机器学习编译器栈"; style=filled; color="#e9ecef"; fillcolor="#f1f3f5"; HighIR [label="高级IR\n(张量操作图)", fillcolor="#d0bfff"]; Opt [label="优化器\n(融合,布局)", shape=ellipse, fillcolor="#ffffff"]; LowIR [label="低级IR\n(循环与指针)", fillcolor="#ffc9c9"]; } subgraph cluster_backend { label = "后端(硬件)"; style=dashed; color="#ced4da"; GPU [label="NVPTX / CUDA"]; CPU [label="LLVM / x86汇编"]; } PyTorch -> HighIR [label="导入"]; TF -> HighIR [label="导入"]; HighIR -> Opt; Opt -> LowIR [label="下沉"]; LowIR -> GPU [label="代码生成"]; LowIR -> CPU [label="代码生成"]; }计算从框架到硬件的流程。高级IR捕捉意图,而低级IR处理实现细节。多级方言方法现代机器学习编译器,例如基于MLIR(多级中间表示)基础设施的编译器,不依赖单一的IR格式。相反,它们使用一种表示层次结构,通常称为“方言”。这种层次结构解决了优化方面的不足,通过让编译器在最适合特定优化任务的抽象级别上工作。图级别: 在最高级别,IR类似于依赖图。节点对应于张量代数操作。这里的优化侧重于图的结构。例如,代数简化(如 $A \times 0 = 0$)或算子融合(将乘法和加法合并到一个内核中)在此处进行。张量/循环级别: 一旦结构性优化完成,编译器会“下沉”表示。抽象的 conv2d 节点被扩展为如何计算它的明确定义,通常涉及嵌套循环和标量计算。循环分块、向量化和内存分配在此处发生。硬件级别: 最后,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有效地将逻辑“冻结”为一个静态快照,为积极的重写和优化提供稳固的支撑。