- Authors

- Name
- Youngju Kim
- @fjvbn20031
- 引言
- 1. 编译器基础:完整流水线
- 2. 中间表示 (IR, Intermediate Representation)
- 3. 优化 Pass (Optimization Passes)
- 4. 代码生成:x86-64 汇编
- 5. JIT 编译:运行时优化
- 6. AI/ML 编译器生态
- 7. MLIR:下一代编译器基础设施
- 8. 小测验:检验理解程度
- 结语
引言
软件工程师需要理解"编译器",并不仅仅是为了造一门语言。如今在 AI/ML 领域,torch.compile()、TVM、XLA、MLIR 这类工具已经成为核心基础设施,它们的内部工作原理无一例外都建立在编译器理论之上。
本文将从头到尾梳理编译器设计的完整流水线。
1. 编译器基础:完整流水线
编译器是把源代码转换为可执行代码的程序,这个过程分为多个阶段。
源代码
↓ 词法分析 (Lexing)
Token 流
↓ 语法分析 (Parsing)
AST (抽象语法树)
↓ 语义分析 (Semantic Analysis)
Annotated AST
↓ IR 生成
中间表示 (IR)
↓ 优化 (Optimization)
优化后的 IR
↓ 代码生成 (Code Generation)
机器码 / 汇编
1.1 词法分析 (Lexical Analysis / Lexing)
词法分析器(Lexer 或 Tokenizer)把源代码字符串转换为 Token(词法单元) 的流。
Token 是数字字面量、标识符、关键字、运算符、标点符号这类具有意义的最小单位。
import re
from enum import Enum, auto
from dataclasses import dataclass
from typing import List, Optional
class TokenType(Enum):
NUMBER = auto()
IDENT = auto()
PLUS = auto()
MINUS = auto()
STAR = auto()
SLASH = auto()
LPAREN = auto()
RPAREN = auto()
ASSIGN = auto()
SEMICOL = auto()
LET = auto()
EOF = auto()
@dataclass
class Token:
type: TokenType
value: str
line: int
KEYWORDS = {'let': TokenType.LET}
TOKEN_PATTERNS = [
(r'\d+(\.\d+)?', TokenType.NUMBER),
(r'[a-zA-Z_]\w*', TokenType.IDENT),
(r'\+', TokenType.PLUS),
(r'-', TokenType.MINUS),
(r'\*', TokenType.STAR),
(r'/', TokenType.SLASH),
(r'\(', TokenType.LPAREN),
(r'\)', TokenType.RPAREN),
(r'=', TokenType.ASSIGN),
(r';', TokenType.SEMICOL),
]
def tokenize(source: str) -> List[Token]:
tokens = []
line = 1
pos = 0
while pos < len(source):
if source[pos] in ' \t\r':
pos += 1
continue
if source[pos] == '\n':
line += 1
pos += 1
continue
matched = False
for pattern, ttype in TOKEN_PATTERNS:
m = re.match(pattern, source[pos:])
if m:
value = m.group(0)
# 处理关键字
if ttype == TokenType.IDENT and value in KEYWORDS:
ttype = KEYWORDS[value]
tokens.append(Token(ttype, value, line))
pos += len(value)
matched = True
break
if not matched:
raise SyntaxError(f"Unexpected character '{source[pos]}' at line {line}")
tokens.append(Token(TokenType.EOF, '', line))
return tokens
# 使用示例
src = "let x = 10 + 20 * 3;"
for tok in tokenize(src):
print(tok)
1.2 语法分析 (Parsing) 与 AST
解析器接收 Token 流,构建出 抽象语法树(AST, Abstract Syntax Tree)。AST 表示程序的层级结构。
from dataclasses import dataclass, field
from typing import Union
# AST 节点定义
@dataclass
class NumberLit:
value: float
@dataclass
class Identifier:
name: str
@dataclass
class BinOp:
op: str
left: 'Expr'
right: 'Expr'
@dataclass
class LetStmt:
name: str
value: 'Expr'
Expr = Union[NumberLit, Identifier, BinOp]
Stmt = Union[LetStmt]
# Recursive Descent Parser
class Parser:
def __init__(self, tokens: List[Token]):
self.tokens = tokens
self.pos = 0
def peek(self) -> Token:
return self.tokens[self.pos]
def consume(self, expected: Optional[TokenType] = None) -> Token:
tok = self.tokens[self.pos]
if expected and tok.type != expected:
raise SyntaxError(f"Expected {expected}, got {tok.type} at line {tok.line}")
self.pos += 1
return tok
def parse_stmt(self) -> Stmt:
if self.peek().type == TokenType.LET:
return self.parse_let()
raise SyntaxError("Expected statement")
def parse_let(self) -> LetStmt:
self.consume(TokenType.LET)
name = self.consume(TokenType.IDENT).value
self.consume(TokenType.ASSIGN)
expr = self.parse_expr()
self.consume(TokenType.SEMICOL)
return LetStmt(name=name, value=expr)
def parse_expr(self) -> Expr:
return self.parse_additive()
def parse_additive(self) -> Expr:
left = self.parse_multiplicative()
while self.peek().type in (TokenType.PLUS, TokenType.MINUS):
op = self.consume().value
right = self.parse_multiplicative()
left = BinOp(op=op, left=left, right=right)
return left
def parse_multiplicative(self) -> Expr:
left = self.parse_primary()
while self.peek().type in (TokenType.STAR, TokenType.SLASH):
op = self.consume().value
right = self.parse_primary()
left = BinOp(op=op, left=left, right=right)
return left
def parse_primary(self) -> Expr:
tok = self.peek()
if tok.type == TokenType.NUMBER:
self.consume()
return NumberLit(value=float(tok.value))
if tok.type == TokenType.IDENT:
self.consume()
return Identifier(name=tok.value)
if tok.type == TokenType.LPAREN:
self.consume()
expr = self.parse_expr()
self.consume(TokenType.RPAREN)
return expr
raise SyntaxError(f"Unexpected token {tok}")
2. 中间表示 (IR, Intermediate Representation)
把 AST 直接转换成机器码效率低下。借助 中间表示(IR),则可以实现与具体语言无关的优化。
2.1 三地址代码 (Three-Address Code)
这是最简单的 IR 形式,每个运算最多只有三个操作数。
# 源代码: x = (a + b) * (c - d)
t1 = a + b
t2 = c - d
x = t1 * t2
2.2 SSA(静态单赋值)形式
SSA 是每个变量都 只被赋值恰好一次 的 IR 形式,这让优化分析大大简化。
# 普通代码
x = 1
x = x + 2 # x 被赋值了两次
# SSA 转换后
x_1 = 1
x_2 = x_1 + 2
在分支汇合的位置,使用 Phi 函数。
if cond:
x_1 = 10
else:
x_2 = 20
# 汇合点
x_3 = phi(x_1, x_2)
2.3 LLVM IR
LLVM IR 是实际工业级编译器所使用的强大 IR。
; 函数定义: int add(int a, int b)
define i32 @add(i32 %a, i32 %b) {
entry:
%result = add i32 %a, %b
ret i32 %result
}
; 更复杂的示例: 条件分支 + SSA
define i32 @max(i32 %a, i32 %b) {
entry:
%cmp = icmp sgt i32 %a, %b ; a > b ?
br i1 %cmp, label %then, label %else
then:
br label %merge
else:
br label %merge
merge:
; phi 函数: 根据来自哪个基本块选择对应的值
%result = phi i32 [ %a, %then ], [ %b, %else ]
ret i32 %result
}
LLVM IR 的特点:
- 类型系统:
i32、i64、float、double、指针等类型都十分明确 - 基于 SSA:每个值只被定义一次
- 虚拟寄存器数量不限:
%0、%1、%2... 数量没有限制 - 语言无关:C、C++、Rust、Swift 都会转换成同一套 LLVM IR
3. 优化 Pass (Optimization Passes)
下面来看编译器优化的核心技术。
3.1 Constant Folding(常量折叠)
在编译期把能够算出结果的常量表达式提前计算好。
# 优化前
x = 2 * 3 + 4
# 优化后 (编译器已提前算好)
x = 10
def constant_fold(node):
"""简单的常量折叠实现"""
if isinstance(node, BinOp):
left = constant_fold(node.left)
right = constant_fold(node.right)
if isinstance(left, NumberLit) and isinstance(right, NumberLit):
ops = {'+': lambda a,b: a+b, '-': lambda a,b: a-b,
'*': lambda a,b: a*b, '/': lambda a,b: a/b}
result = ops[node.op](left.value, right.value)
return NumberLit(value=result)
return BinOp(op=node.op, left=left, right=right)
return node
3.2 Dead Code Elimination(死代码消除)
移除不会被执行、或结果不会被使用的代码。
def is_reachable(block, cfg):
"""在控制流图中用 BFS 搜索可达的基本块"""
visited = set()
queue = [cfg.entry]
while queue:
current = queue.pop(0)
if current in visited:
continue
visited.add(current)
queue.extend(cfg.successors[current])
return block in visited
3.3 Loop Unrolling(循环展开)
减少循环开销,提升指令级并行度(ILP)。
// 原始代码 (循环 4 次)
for (int i = 0; i < 4; i++) {
a[i] = b[i] + c[i];
}
// 展开后 (消除循环)
a[0] = b[0] + c[0];
a[1] = b[1] + c[1];
a[2] = b[2] + c[2];
a[3] = b[3] + c[3];
3.4 Inlining(函数内联)
消除函数调用开销,并创造进一步优化的机会。
// 原始代码
inline int square(int x) { return x * x; }
int y = square(5);
// 内联后
int y = 5 * 5; // → 通过常量折叠得到 25
4. 代码生成:x86-64 汇编
这是把 IR 转换为实际机器码的过程。
4.1 寄存器分配
x86-64 拥有 16 个通用寄存器(rax、rbx、rcx、rdx、rsi、rdi、rsp、rbp、r8~r15),需要把虚拟寄存器映射到这些物理寄存器上。
图着色(Graph Coloring) 算法是最著名的寄存器分配方法。同时存活(live)的变量不能共享同一个寄存器。
4.2 x86-64 汇编示例
; C 函数: int add(int a, int b) { return a + b; }
; Linux x86-64 System V ABI: 参数依次是 rdi, rsi
add:
push rbp
mov rbp, rsp
mov DWORD PTR [rbp-4], edi ; 保存 a
mov DWORD PTR [rbp-8], esi ; 保存 b
mov edx, DWORD PTR [rbp-4]
mov eax, DWORD PTR [rbp-8]
add eax, edx ; a + b
pop rbp
ret ; 结果保存在 eax 中
; 优化后的版本 (-O2)
add:
lea eax, [rdi + rsi] ; 只需 1 条指令!
ret
5. JIT 编译:运行时优化
JIT(Just-In-Time,即时编译) 是在程序运行过程中编译代码的技术,是克服 Python 执行速度慢这一问题的核心手段。
5.1 Numba JIT
Numba 通过 LLVM 把 Python/NumPy 代码编译成机器码。
import numpy as np
from numba import njit, prange
import time
# 纯 Python 版本
def matmul_python(A, B):
N = A.shape[0]
C = np.zeros((N, N))
for i in range(N):
for j in range(N):
for k in range(N):
C[i, j] += A[i, k] * B[k, j]
return C
# Numba JIT 版本 (含并行化)
@njit(parallel=True)
def matmul_numba(A, B):
N = A.shape[0]
C = np.zeros((N, N))
for i in prange(N): # 并行循环
for j in range(N):
for k in range(N):
C[i, j] += A[i, k] * B[k, j]
return C
N = 256
A = np.random.rand(N, N).astype(np.float32)
B = np.random.rand(N, N).astype(np.float32)
# 首次调用时触发编译 (预热)
_ = matmul_numba(A, B)
t0 = time.perf_counter()
C_python = matmul_python(A, B)
print(f"Python: {time.perf_counter()-t0:.3f}s")
t0 = time.perf_counter()
C_numba = matmul_numba(A, B)
print(f"Numba JIT: {time.perf_counter()-t0:.4f}s")
# 结果: Python ~30s vs Numba ~0.01s (快数百倍)
Numba 的内部工作流程:
- 把 Python 字节码转换为 Numba IR
- 通过类型推断(type inference)确定静态类型
- 生成 LLVM IR
- LLVM 后端编译为机器码
- 把编译好的代码缓存起来
6. AI/ML 编译器生态
6.1 torch.compile() 的内部工作原理
torch.compile() 是 PyTorch 2.0 的核心功能,由三个组件构成。
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(128, 256)
self.linear2 = nn.Linear(256, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = self.relu(self.linear1(x))
return self.linear2(x)
model = SimpleNet().cuda()
x = torch.randn(32, 128).cuda()
# 应用 torch.compile
compiled_model = torch.compile(model, mode='max-autotune')
# 首次调用: 触发编译 (耗时数秒)
out = compiled_model(x)
# 之后的调用: 复用已编译的内核 (速度快)
for _ in range(100):
out = compiled_model(x)
编译栈的三个阶段:
第 1 阶段 — TorchDynamo(图捕获)
拦截 Python 字节码,将其转换为 torch.fx.Graph。
import torch._dynamo as dynamo
def my_func(x, y):
return torch.sin(x) + torch.cos(y)
# 查看 FX 图
explanation = dynamo.explain(my_func)(torch.randn(4), torch.randn(4))
print(explanation.graphs[0].print_tabular())
# opcode | name | target | args
# --------|--------|---------------|----------
# call_fn | sin | torch.sin | (x,)
# call_fn | cos | torch.cos | (y,)
# call_fn | add | operator.add | (sin, cos)
# output | output | output | (add,)
第 2 阶段 — AOTAutograd(自动微分预编译) 预先生成前向(forward)与反向(backward)传播图。
第 3 阶段 — TorchInductor(内核生成) 把 FX 图转换成 Triton(GPU)或 C++(CPU)内核代码。
6.2 TVM:深度学习编译器
Apache TVM 能把深度学习模型针对各种硬件(GPU、CPU、FPGA、NPU)做优化。
import tvm
from tvm import relay
import tvm.relay.testing
import numpy as np
# 把 PyTorch 模型转换为 Relay IR
import torch
import torchvision
model = torchvision.models.resnet18(pretrained=False).eval()
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
# 经由 TorchScript 转换为 TVM Relay
scripted_model = torch.jit.trace(model, input_data)
shape_list = [("input0", input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
# 设置目标平台并优化
target = tvm.target.Target("cuda")
# AutoTuning: 搜索最优内核参数
from tvm import auto_scheduler
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=200, # 测试次数
runner=auto_scheduler.LocalRunner(repeat=10),
measure_callbacks=[auto_scheduler.RecordToFile("tuning_log.json")],
)
tuner.tune(tune_option)
# 用调优结果进行编译
with auto_scheduler.ApplyHistoryBest("tuning_log.json"):
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
print("TVM 编译完成!")
6.3 Operator Fusion(算子融合)
这是 GPU 上最重要的优化,把多次内核调用合并成一次。
# 融合前: 3 次 GPU 内核调用
y1 = conv2d(x) # 内核 1 (内存读写)
y2 = batch_norm(y1) # 内核 2 (内存读写)
y3 = relu(y2) # 内核 3 (内存读写)
# 融合后: 1 次 GPU 内核调用
y3 = conv_bn_relu(x) # 单一的融合内核
# 中间结果保留在 L1/L2 缓存中
# 节省内存带宽 = 性能提升
XLA(Accelerated Linear Algebra)是 TensorFlow/JAX 的编译器,使用 HLO(High Level Operations)作为 IR。
import jax
import jax.numpy as jnp
# JAX + XLA: 用 jit 装饰器启用 XLA 编译
@jax.jit
def compute(x, W, b):
# XLA 会自动把 matmul + add + relu 融合起来
return jax.nn.relu(x @ W + b)
x = jnp.ones((32, 128))
W = jnp.ones((128, 256))
b = jnp.zeros(256)
# 首次调用: 触发 XLA 编译
result = compute(x, W, b)
print(result.shape) # (32, 256)
7. MLIR:下一代编译器基础设施
MLIR(Multi-Level Intermediate Representation)是 Google 设计的编译器框架,把多个层级的 IR 统一到同一套基础设施中。
TensorFlow Graph
↓ (TF Dialect)
MLIR HLO Dialect
↓ (Linalg Dialect)
Loop Nest
↓ (Affine Dialect)
Affine Loops
↓ (LLVM Dialect)
LLVM IR
↓
机器码
每个 Dialect 都代表某个特定层级的抽象,编译过程会逐步向更低层级转换(lowering)。
8. 小测验:检验理解程度
Q1. LL(1) 解析器与 LR(1) 解析器在解析能力上有什么区别?
答案:LR(1) 解析器能够处理范围更广的文法。
解析:LL(1) 表示 Left-to-right、Leftmost derivation、1 个 lookahead,是一种自顶向下(top-down)的解析器,无法处理左递归(left-recursive)文法,并且只能看一个 Token 就决定该套用哪条规则。LR(1) 表示 Left-to-right、Rightmost derivation、1 个 lookahead,是一种自底向上(bottom-up)的解析器。LALR(1) 是 LR(1) 在内存占用上更高效的变体,被 GCC、Bison 等实际编译器广泛采用。LR 解析器能处理的文法类别比 LL 解析器宽得多,绝大多数编程语言的文法都可以用 LALR(1) 表示。
Q2. SSA 形式为什么有利于优化?
答案:因为变量的数据流分析会变得极其简单。
解析:在 SSA 中,每个变量只被定义一次,因此变量的定义(def)与使用(use)关系十分清晰。这带来了:(1) 常量传播 — 若定义处的值是常量,就能在所有使用处立即替换;(2) 死代码消除 — 没有 use 的 def 可以直接删除;(3) 寄存器分配 — 活跃性分析(liveness analysis)随之简化;(4) 循环不变代码外提 — 依赖关系分析更容易。LLVM、GCC 的 middle-end 优化 Pass 全部都建立在 SSA 之上。
Q3. LLVM IR 有哪些特点?为什么能实现与语言无关的优化?
答案:因为 LLVM IR 是一种不依赖特定语言或架构的通用底层 IR。
解析:LLVM IR 的核心特点是:(1) 强类型系统,保证类型安全;(2) 基于 SSA,便于分析与优化;(3) 虚拟寄存器数量不限(物理寄存器分配放到代码生成阶段);(4) 显式的内存模型与对齐信息。语言无关性来自于把前端(C→IR、Rust→IR、Swift→IR)与后端(IR→x86、IR→ARM、IR→RISC-V)彻底分离 — 中间的优化 Pass 只处理 IR,因此可以套用到任意语言与架构的组合上。
Q4. TVM 中的 operator fusion 是如何提升 GPU 性能的?
答案:因为它能减少 GPU 内存带宽瓶颈,并把缓存效率发挥到极致。
解析:GPU 运算真正的瓶颈,很多时候并不是运算量(FLOPS)本身,而是内存带宽(HBM ↔ SRAM)。没有融合时,每个算子都要把中间结果写入 HBM(全局内存),下一个内核再重新读取一次。融合之后,中间结果留在 L1/L2 缓存或寄存器中,HBM 访问次数大幅减少。elementwise 运算(ReLU、BatchNorm 等)尤其是内存受限(memory-bound)型,因此融合带来的效果格外显著。TVM 在 Relay→TIR 的转换过程中,会自动检测出可以融合的算子组。
Q5. torch.compile() 内部使用的编译栈是什么?
答案:TorchDynamo → AOTAutograd → TorchInductor 这三个阶段构成的编译栈。
解析:(1) TorchDynamo 工作在 Python 字节码层面,在执行 Python 代码的同时,把其中的 PyTorch 运算部分捕获为 FX Graph。为了应对 Python 的动态特性(控制流、依赖数据的分支),它使用基于 guard 的特化机制。(2) AOTAutograd 接收前向 FX Graph,提前(ahead-of-time)生成出反向传播图。(3) TorchInductor 把 FX Graph 转换为实际的 GPU 内核代码(Triton)或 CPU 代码(C++),operator fusion、tiling、vectorization 等优化都在这一步完成。可以通过 mode 选项在 'default'、'reduce-overhead'、'max-autotune' 之间选择。
结语
编译器设计不是纸上谈兵的理论。今天 AI/ML 基础设施的核心 —— torch.compile()、TVM、XLA、MLIR,全都建立在数十年编译器研究的积累之上。
理解 Lexer → Parser → AST → IR → Optimization → Code Gen 这条流水线之后,你就能自然地明白:为什么 GPU 编程里内核融合如此重要,为什么 LLVM 会被这么多语言用作后端,为什么 SSA 形式是优化的基础。
下一步,推荐阅读 LLVM 的 Kaleidoscope 教程、TVM 官方文档,以及 PyTorch 2.0 论文。