现代深度学习编译器通过管理两个不同抽象阶段,与传统语言编译器有别。当标准C++编译器可能直接将源代码转换为LLVM IR等指令级中间表示(IR)时,AI编译器则需首先考量数学结构,然后再关注汇编代码结构。这种需要促成了一种多级IR架构,主要分为图级(高级)IR和内核级(低级)IR。图级中间表示高级IR是编译器在解析前端框架代码(如PyTorch或TensorFlow)后,立即获得的模型视图。在此阶段,该表示保持声明性。编译器了解需要计算 什么,但尚未确定 如何 计算。在图级IR中,数据的基本单元是 张量。编译器记录数据类型(float32, int8)、形状(维度)和内存布局(NCHW与NHWC)等属性。执行的基本单元是 操作符。此图中的单个节点,例如 Conv2D 或 Softmax,表示大量的算术运算。思考Transformer模型中的一个典型层。在高级IR中,一个全连接层后跟一个激活函数表现为一组最少的节点。底层循环的复杂程度通过操作符抽象被隐藏。digraph G { rankdir=TB; node [shape=box, style=filled, fontname="Helvetica", fontsize=10]; edge [fontname="Helvetica", fontsize=9, color="#868e96"]; input [label="输入张量\n[批次, 128]", fillcolor="#a5d8ff", color="#228be6"]; weights [label="权重\n[128, 512]", fillcolor="#ffe066", color="#f59f00"]; bias [label="偏置\n[512]", fillcolor="#ffe066", color="#f59f00"]; matmul [label="操作: 矩阵乘法", fillcolor="#ffc9c9", color="#fa5252"]; add [label="操作: 相加", fillcolor="#ffc9c9", color="#fa5252"]; relu [label="操作: ReLU", fillcolor="#ffc9c9", color="#fa5252"]; output [label="输出张量\n[批次, 512]", fillcolor="#b2f2bb", color="#40c057"]; input -> matmul; weights -> matmul; matmul -> add [label="中间结果\n[批次, 512]"]; bias -> add; add -> relu; relu -> output; }一个密集层的高级图表示。节点表示逻辑数学运算而非CPU指令。此级别的主要用途是全局优化。因为编译器可以看到神经网络的整个拓扑结构,它可以执行代数简化和结构调整。例如,如果编译器遇到以下序列:$$Y = \text{ReLU}(\text{相加}(\text{二维卷积}(X, W), B))$$它不必担心寄存器分配或缓存行。它侧重于 操作符合并,思考这三个操作能否合并为一个内核调用,以减少与全局内存(VRAM)的内存往返。降低过程从高级IR到低级IR的转换常被称为“降低”。这是一个在语义上具有破坏性的过程。一旦 Conv2D 节点被降低,编译器就失去其正在执行卷积的语义知识。相反,该操作被其实现细节取代:嵌套循环、内存加载、乘累加指令和存储。这种区分有其价值,因为某些优化只能在特定级别进行。一旦代码分解成循环,就无法有效进行操作符合并,因为结构变得过于模糊。反之,你无法在图级执行循环平铺或向量化,因为循环尚未在该级别正式存在。内核级中间表示低级IR类似于C或汇编语言的受限子集。它是命令式的而非声明式的。在此阶段,抽象从张量转向 缓冲区(带有指针的已分配内存块),并从操作符转向 循环嵌套。在Apache TVM(特别是张量中间表示,TIR)或MLIR的Affine/Linalg方言等框架中,低级IR将图中所有隐式内容显式化。编译器现在必须管理:迭代域: 循环的边界($i$ 从 $0$ 到 $N$)。内存访问: 计算平面内存缓冲区中的特定索引($A[i \times 步长 + j]$)。作用域和存储: 判断数据是存在全局内存、共享内存(GPU)还是寄存器中。思考图示例中的相同矩阵乘法。在低级IR中,它被扩展为适合硬件特定调整的迭代结构。digraph LowLevel { rankdir=TB; node [shape=note, style=filled, fontname="Courier", fontsize=10]; edge [color="#adb5bd"]; alloc [label="分配缓冲区 C[M*N]", fillcolor="#eebefa", color="#be4bdb"]; loop_i [label="循环 i = 0 到 M", fillcolor="#d0bfff", color="#7950f2"]; loop_j [label="循环 j = 0 到 N", fillcolor="#d0bfff", color="#7950f2"]; loop_k [label="循环 k = 0 到 K", fillcolor="#d0bfff", color="#7950f2"]; compute [label="寄存器_C += A[i,k] * B[k,j]", fillcolor="#ffc9c9", color="#fa5252"]; store [label="存储 C[i,j] <- 寄存器_C", fillcolor="#99e9f2", color="#15aabf"]; alloc -> loop_i; loop_i -> loop_j; loop_j -> loop_k; loop_k -> compute; loop_j -> store [label="归约后"]; }降低后相同操作的表示。侧重于迭代顺序、内存分配和标量计算。此级别是性能工程主要工作发生的场所。编译器尝试重塑这些循环嵌套以最大化数据局部性。循环平铺(将大循环分解成更小的块以适应缓存)和 向量化(使用AVX-512或NEON等SIMD指令)等技术在此处应用。IR特点比较分析要设计一个有效的编译器通道,必须确定哪个IR级别能为该任务提供正确的原语。在图级别进行内存规划是不精确的,因为缓冲区生命周期尚未完全确定。在循环级别执行死代码消除,与更早地修剪图相比效率不高。以下分解阐明了这两个抽象层之间的功能差异。{ "data": [ { "type": "table", "header": { "values": [ "<b>特征</b>", "<b>高级IR (图)</b>", "<b>低级IR (内核)</b>" ], "align": "left", "fill": { "color": "#e9ecef" }, "font": { "family": "Arial", "size": 12, "color": "#495057" } }, "cells": { "values": [ [ "数据单元", "操作", "控制流", "内存模型", "优化侧重", "示例系统" ], [ "张量 (形状 + 类型)", "逻辑 (Conv2D, 矩阵乘法)", "数据依赖边", "抽象 / 隐式", "合并, 布局, 简化", "TVM Relay, XLA HLO, TF Graph" ], [ "缓冲区 (指针 + 偏移)", "标量 (加载, 存储, FMA)", "循环, If/Else, 跳转", "显式 (分配/释放)", "平铺, 向量化, 循环展开", "TVM TIR, MLIR Affine, LLVM IR" ] ], "align": "left", "fill": { "color": [ "#f8f9fa", "#ffffff", "#f8f9fa" ] }, "font": { "family": "Arial", "size": 11, "color": "#212529" }, "height": 30 } } ], "layout": { "title": { "text": "IR抽象层比较", "font": { "size": 16 } }, "margin": { "l": 20, "r": 20, "t": 40, "b": 20 }, "height": 350 } }功能比较,区分图级表示与内核级表示的范围和能力。灰色地带:多级方言在MLIR(多级中间表示)等现代基础设施中,高级和低级之间的界限变得不那么严格。编译器可能采用涉及多种方言的渐进式降低策略,而非两种单一状态。例如,MLIR中的 Linalg 方言可以说处于中间位置。它有效表示矩阵乘法等操作(保留语义意图),但对适合循环分析的缓冲区进行操作。这使得编译器能够利用图级信息执行循环合并,这是一种通常专用于循环级IR的技术。这种混合方法允许编译器在感知更广的操作符语境的同时做出平铺决策,解决了“阶段排序”问题,即一个级别的优化可能会无意中使下一个级别的代码变差。理解这种二分法对于后续实践环节很有助益。在查看TVM的Relay IR时,你将看到声明性图。在关于调度的后续章节中,你将看到该图如何降低为命令式TIR,以便针对特定硬件约束进行调整。