Skip to content
Published on

编译器/解释器设计完全指南:从解析器到 LLVM、AI 编译器(TVM/XLA)

分享
Authors

引言

软件工程师需要理解"编译器",并不仅仅是为了造一门语言。如今在 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 的特点:

  • 类型系统i32i64floatdouble、指针等类型都十分明确
  • 基于 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 个通用寄存器(raxrbxrcxrdxrsirdirsprbpr8~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 的内部工作流程:

  1. 把 Python 字节码转换为 Numba IR
  2. 通过类型推断(type inference)确定静态类型
  3. 生成 LLVM IR
  4. LLVM 后端编译为机器码
  5. 把编译好的代码缓存起来

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 上最重要的优化,把多次内核调用合并成一次。

# 融合前: 3GPU 内核调用
y1 = conv2d(x)        # 内核 1 (内存读写)
y2 = batch_norm(y1)   # 内核 2 (内存读写)
y3 = relu(y2)         # 内核 3 (内存读写)

# 融合后: 1GPU 内核调用
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 论文。