TT Lab

한국어 English 日本語 中文

开始
博客

博客

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

对话收听 · 10分钟两位主持人为你讲解本文

引言

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


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 的内部工作流程:

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

# 融合前: 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 论文。

登录后即可点赞

评论

还没有评论。

登录后即可发表评论