""" ============================================================================= CS349E · 前置讲 P4 · 01_compiler_pipeline.py —— 走一遍编译器的流水线：文本 → token → AST → IR → 优化 → 代码 ============================================================================= 【这个脚本要回答的问题】 编译器把一段程序从"人写的样子"变成"机器跑的样子"。它是怎么一步步做到的？每一步的产物长什么样？ 我们用 mini_compiler 处理一段 SwiGLU 激活函数的代码（LLM 的 FFN 里每层都有），看清每一个阶段： 演示 1：前端——词法分析（字符 → token）和语法分析（token → 抽象语法树 AST）。 演示 2：降低到 IR——三地址码、SSA（每个名字只赋值一次）。 演示 3：中端优化——常量折叠、代数化简、公共子表达式消除、死代码删除；每一步都检查结果没变。 演示 4：后端——生成 numpy 代码和 C 代码；再交给系统的 C 编译器，看 -O0 和 -O2 生成的机器指令有什么不同。 【怎么读输出】演示 1–3 是确定性的；演示 4 的汇编取决于你的 C 编译器和 CPU（编写时是 Apple clang、arm64）。 【运行】python3 01_compiler_pipeline.py （约 1 秒） 【依赖】numpy；从 ../lecture01 导入 console；同目录的 mini_compiler、cbuild（没有 C 编译器时跳过 C 部分）。 """ import os import re import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import print_header, print_observations, print_table # noqa: E402 from cbuild import assembly, find_cc # noqa: E402 from mini_compiler import (PASSES, ast_to_str, compile_c_kernel, compile_numpy, eval_ast, ir_to_str, # noqa: E402 lower, op_counts, parse, run_ir, to_c_source, to_numpy_source, tokenize) SOURCE = """# SwiGLU 的逐元素部分：silu(g) * u，故意写得不太好 c = 0.5 * 2.0; # 一个编译时就能算出来的常数 s = 1 / (1 + exp(-g)); # sigmoid(g) silu = g * s * c; # 乘以 c（= 1） gate = g / (1 + exp(-g)); # 和 silu 数学上相同，但写法不同 y = silu * u + 0; # 加 0 dead = sqrt(abs(u)) * 3; # 算了但没有人用 """ OUTPUTS = ["y", "gate"] # --------------------------------------------------------------------------- # 演示 1：前端 # --------------------------------------------------------------------------- def demo1_frontend(): """ 词法分析：把字符流切成 token（数字、名字、运算符），丢掉空白和注释。 语法分析：按文法把 token 组织成树。树的结构体现了运算的优先级和结合性——不需要括号了。 """ print_header("演示 1：前端——词法分析与语法分析") print("源代码：") print(" " + SOURCE.strip().replace("\n", "\n ")) toks = tokenize(SOURCE) print(f"\n词法分析：{len(SOURCE)} 个字符 → {len(toks)} 个 token（含结尾的 eof）。第二条语句的 token：") line2 = SOURCE.splitlines()[2] start = SOURCE.index(line2) seq = [f"{t.kind}:{t.text}" for t in toks if start <= t.pos < start + len(line2)] print(" " + " ".join(seq)) prog = parse(SOURCE) print(f"\n语法分析：{len(prog)} 条语句。第二条语句 s = 1 / (1 + exp(-g)) 的 AST：") print(" " + ast_to_str(prog[1]).rstrip().replace("\n", "\n ")) print_observations([ "词法分析只关心'单词'：注释和空白在这一步就消失了，后面的阶段再也看不到它们。", "AST 把优先级变成了树的结构：除法是根，1 + exp(−g) 是它的右子树；括号的作用已经体现在树的形状里，括号本身不再需要。", "AST 仍然很接近源代码（一条语句一棵树）。它适合做语法检查和简单的解释执行，但不适合做优化——" "优化需要把每个运算拆成一步、看清'哪个值被谁用了'，这就是下一步的 IR。", ]) return prog # --------------------------------------------------------------------------- # 演示 2：降低到 IR # --------------------------------------------------------------------------- def demo2_lowering(prog): """ 把 AST 展开成"三地址码"：每条指令最多一个运算、两个参数、一个结果。 结果的名字遵守 SSA：%s1、%silu1……每个名字只被赋值一次；临时值叫 %tN。 """ print_header("演示 2：降低到 IR——三地址码与 SSA") instrs, outs, inputs = lower(prog, OUTPUTS) print(f"输入变量：{inputs}；输出：{OUTPUTS}") print(ir_to_str(instrs, outs)) print_observations([ f"6 条语句展开成 {len(instrs)} 条指令；每条指令只做一件事，所以'这个值从哪里来、被谁用到'一目了然——这就是数据流图（02 脚本）。", "赋值语句本身没有生成指令：y = … 只是把名字 y 绑定到最后一条指令的结果上。", "SSA 的好处：如果源程序给同一个变量赋值两次，IR 会生成两个不同的名字（%y1、%y2），" "所以'两个名字相同的操作数一定是同一个值'——后面的公共子表达式消除就靠这一点。", "LLVM IR、MLIR、PyTorch 的 FX graph / Inductor IR、XLA 的 HLO 都是 SSA 形式的数据流 IR。", ]) return instrs, outs, inputs # --------------------------------------------------------------------------- # 演示 3：中端优化 # --------------------------------------------------------------------------- def demo3_passes(instrs, outs, inputs, prog, n=10_000, seed=0): """ 依次运行四个 pass，每一步之后：数指令、在随机输入上运行 IR，和 AST 解释器的结果比较。 【为什么每一步都要验证】优化必须**保持语义**。编译器测试里最常见的方法正是"差分测试"： 同一个输入，优化前后的结果必须一致。 """ print_header("演示 3：中端优化——四个 pass，每一步都检查结果没变") rng = np.random.default_rng(seed) env = {"g": rng.standard_normal(n) * 3, "u": rng.standard_normal(n)} ref = eval_ast(prog, env) rows = [["（降低后）", len(instrs), "—", "—"]] cur, cur_outs = instrs, outs for name, p in PASSES: nxt, nxt_outs = p(cur, cur_outs) res = run_ir(nxt, nxt_outs, env) err = max(np.abs(res[k] - ref[k]).max() for k in OUTPUTS) removed = sorted(set(i.dst for i in cur) - set(i.dst for i in nxt)) rows.append([name, len(nxt), f"{err:.1e}", ", ".join(removed) if removed else "（无）"]) cur, cur_outs = nxt, nxt_outs print_table(["pass", "剩余指令", "和 AST 解释器的最大差", "删掉的指令"], rows) print("\n优化后的 IR：") print(ir_to_str(cur, cur_outs)) c0, c1 = op_counts(instrs), op_counts(cur) print("\n各种运算的条数（优化前 → 优化后）：" + "，".join(f"{op} {c0.get(op, 0)}→{c1.get(op, 0)}" for op in sorted(c0))) print_observations([ f"{len(instrs)} 条指令优化到 {len(cur)} 条，结果和原程序完全一致（最大差 0）。", "常量折叠把 c = 0.5 × 2.0 在编译时算成 1；代数化简删掉了 ×1 和 +0；死代码删除去掉了没人用的 sqrt(abs(u))。", "公共子表达式消除发现 silu 和 gate 都要算 −g、exp(−g)、1 + exp(−g)，只算一次——exp 是最贵的运算，" "这一下就省了一半的 exp。但它没能发现 g × (1 / t) 和 g / t 数学上相等：CSE 只认'写法相同'，不懂代数。", "注意 x × 0 → 0 这类化简在浮点里是不安全的（x 是 Inf 或 NaN 时结果应该是 NaN），真实编译器只在 -ffast-math 下做；" "同理，(a + b) + c → a + (b + c) 也会改变舍入（P3 §4.2）。'优化必须保持语义'，而浮点的语义比代数更挑剔。", ]) return cur, cur_outs # --------------------------------------------------------------------------- # 演示 4：后端 # --------------------------------------------------------------------------- def count_asm(asm): """ 粗略统计汇编：指令行数（去掉以 . 开头的伪指令和标签）、SIMD 指令数、函数调用。 【怎样认出 SIMD】arm64 的向量寄存器写作 v0.4s（4 个 float32）、q0；x86 的写作 xmm / ymm / zmm。 """ lines = [l.strip() for l in asm.splitlines() if l.startswith("\t") and not l.strip().startswith(".")] simd = [l for l in lines if re.search(r"\bv\d+\.(4s|2s|2d)\b|\bq\d+\b|[xyz]mm\d+", l)] calls = [l.split()[-1] for l in lines if re.match(r"(bl|call|callq)\b", l)] return len(lines), len(simd), calls def demo4_backend(opt, outs, inputs, prog, n=100_000, seed=1): """ 生成两种代码：numpy（每条指令一次数组运算）和 C（一个循环，每个元素做完所有运算）。 验证两者结果；再看 C 编译器在 -O0 / -O2 下生成的汇编。 """ print_header("演示 4：后端——生成 numpy 代码和 C 代码，再交给 C 编译器") print("生成的 numpy 代码（每行一次数组运算）：") print(" " + to_numpy_source(opt, outs, inputs).replace("\n", "\n ")) print("\n生成的 C 代码（一个循环，中间结果都在局部变量里）：") csrc = to_c_source(opt, outs, inputs) print(" " + csrc.replace("\n", "\n ")) rng = np.random.default_rng(seed) env = {"g": rng.standard_normal(n).astype(np.float32) * 3, "u": rng.standard_normal(n).astype(np.float32)} ref = eval_ast(prog, {k: v.astype(np.float64) for k, v in env.items()}) fnp = compile_numpy(opt, outs, inputs) res_np = fnp(**env) rows = [["生成的 numpy 代码", f"{max(np.abs(res_np[k] - ref[k]).max() for k in OUTPUTS):.1e}"]] if find_cc() is None: print("\n没有找到 C 编译器，跳过 C 部分。") print_table(["实现", "和 FP64 参照的最大差"], rows) return res_c = compile_c_kernel(opt, outs, inputs)(**env) rows.append(["生成的 C 代码（-O2）", f"{max(np.abs(res_c[k] - ref[k]).max() for k in OUTPUTS):.1e}"]) print() print_table(["实现", "和 FP64 参照的最大差"], rows) rows = [] for flags in (("-O0",), ("-O2",), ("-O2", "-fveclib=Accelerate")): asm = assembly(csrc, flags) if asm is None: rows.append([" ".join(flags), "（这个编译器不支持）", "—", "—"]) continue n_ins, n_simd, calls = count_asm(asm) uniq = sorted(set(calls)) rows.append([" ".join(flags), n_ins, n_simd, f"{len(calls)} 次（{', '.join(uniq)}）" if calls else "无"]) print("\nC 编译器（后端的后端）生成的汇编：") print_table(["编译选项", "指令行数", "SIMD 指令", "函数调用"], rows) print_observations([ "两种生成的代码结果一致（差别只是 float32 的舍入）。同一个 IR 可以有很多个后端：numpy、C、CUDA、Triton……" "这正是'中间表示'这个名字的意思：前端和后端可以分别替换。", "-O0 几乎逐行翻译 C 代码，不用 SIMD；-O2 做了寄存器分配、指令选择、循环展开和向量化，SIMD 指令出现了。", "但 -O2 下 exp 仍然是对 C 库 expf 的**标量**调用（一个向量里的 4 个数各调用一次）：通用的 C 库没有向量版本。" "如果编译器支持 -fveclib=Accelerate（macOS），它会改调向量版的 vexpf。ML 编译器通常自己生成 exp / tanh 的" "多项式逼近来避免这个问题（02 脚本会看到它对速度的影响，第 1 讲 §5.5 的 exp 瓶颈也是同一件事）。", "我们的小编译器只做到'生成 C'，剩下的（寄存器分配、指令调度、向量化）交给了 clang / gcc——" "Triton、TVM、XLA 也是这样分层的：它们生成 LLVM IR 或 PTX，再交给 LLVM / NVIDIA 的后端（第 13 讲）。", ]) if __name__ == "__main__": program = demo1_frontend() ir, outputs, inp = demo2_lowering(program) optimized, opt_outs = demo3_passes(ir, outputs, inp, program) demo4_backend(optimized, opt_outs, inp, program)