夜雨聆风这一篇文章偏算法一点，但也是重要的应用思路（仿佛回到了本科打数模的时候）。梳理了 AI 编译这一领域中四类搜索方法的设计思路、代表工作与各自适用的层次，他们分别是：以 e-graph 和 Equality Saturation 为代表的符号方法，以 cost model 为核心的启发式方法，以 RL、MCTS、进化算法为代表的 AI 搜索方法，以及最近兴起的 LLM 辅助编译优化。
1. AI 编译的核心问题
AI 编译的核心问题和硬件设计空间探索，或者任何一个搜索问题本质都是一样的：在一个组合规模很大的程序变换空间里，在有限时间内逼近某个性能或资源目标的近似最优解，同时保持语义等价。这一问题包含三组约束。
第一组是搜索空间与时间预算。一个 transformer block 的算子级 schedule 搜索空间在 10^15 量级以上。一条 LLVM pass pipeline 即便限定在常用的 100 余个 pass 内，组合数也很大。但编译时间需要可控，在线场景不能花费几天，离线调优也不希望耗费上千 GPU 小时。
第二组是正确性与性能。任何变换都必须保持语义等价，但越激进的变换（融合、重写、量化、近似）越容易出错。
第三组是泛化性与专精度。一个搜索算法越通用，适配多种硬件、多种算子，通常越难逼近手写库（cuBLAS、cuDNN）的性能。
围绕这三组约束，社区在过去十年大致经历了四代方法。第一代是启发式加模板，代表工具有 Halide、AutoTVM、XLA 的 pattern matching。第二代是符号方法加 ILP，代表工具有 egg、TASO、Tensat。第三代是学习型搜索，代表工具有 Ansor、X-RLflow、AlphaTensor。第四代是 LLM 辅助和 agentic 方法，代表工具有 Meta LLM Compiler、Compiler-R1。
另外提一嘴，这四代不是替代关系，而是叠加关系。当前的工业级编译栈（TVM、XLA、Triton、IREE）几乎都是多种方法的混合：在 graph level 用 pattern 匹配的启发式，在 kernel level 用 cost model 加进化搜索，在某些场景下接入 RL 或 LLM。因此应把它们当作互补工具来看。
第二章讲ai编译器的需求，想了一下，大佬们应该不需要，还是删了，分割一下-----------------------
3. 符号方法：E-graph 与 Equality Saturation
严格来说，e-graph 是表示（representation）方法，不是搜索方法。它做的事情是把"一个程序的所有等价形式"这个集合用一种紧凑的数据结构刻画出来。Equality Saturation 是配套的构建过程（用 rewrite 规则把这个集合填满），extraction 才是真正的搜索（从已构建的结构中选最优）。
3.1 phase-ordering 问题
举个例子说明： y = ReLU(BN(Conv(x, W)))
针对这个 pattern 通常有三条 rewrite 规则。规则 A 是 BN folding，已知BatchNorm 在推理时可以被折叠进 Conv 的权重和 bias 里。规则 B 是 Conv-BN 融合，使用硬件提供的 FusedConvBN kernel 把两个算子合并执行。规则 C 是 Conv-ReLU 融合，用一个支持 epilogue 的 FusedConvReLU kernel。
如果先应用规则 A，BN 消失，规则 B 不再可匹配，最终只能在普通 Conv 之上应用规则 C，得到 FusedConvReLU。如果先应用规则 B，得到 FusedConvBN，规则 A 不再可用；除非硬件支持 FusedConvBN 之后再接 ReLU 的 epilogue，否则 ReLU 必须独立执行，部分融合机会丢失。如果硬件提供 FusedConvBNReLU 三合一 kernel，最优解需要直接匹配整个 pattern，任何先单独应用 A 或 B 的路径都会破坏这个机会。
进一步把范围扩大到 Conv-BN-ReLU 之后接 Add（残差连接）这种 ResNet 基本块结构时，规则的相互依赖更复杂。一些规则要求 Add 之前的两个分支具有特定的数据布局，而布局变换又可能与 BN folding 冲突。规则集合到一定规模后，串行应用的最终结果对顺序的敏感度迅速放大，性能差异可以达到几倍。这是 phase-ordering 问题在张量计算图重写中的具体表现。
3.2 E-graph
Equality Saturation（等式饱和）提供了一种思路：不再串行选择 rewrite，而是把所有可达的等价表达式同时记下来，最后再统一抽取最优。承载这个想法的数据结构是 e-graph（equality graph）。
E-graph 由两类节点组成。e-node 是一个算子，加上若干个指向 e-class 的子节点引用。e-class 是一组语义等价的 e-node 的集合。当两个 e-node 的对应子 e-class 全相等时，它们会被合并到同一个 e-class，这一过程称为 congruence closure。
E-graph 用等价类作为基本单位，让一个图同时表示指数级多的等价表达式。因此本质上
应用流程因此变成两阶段。第一阶段是饱和（saturation），反复应用 rewrite 直到 e-graph 不再变化或达到资源限制。第二阶段是抽取（extraction），在 e-graph 上跑一个 cost-aware 的图选取算法（贪婪或 ILP），找出代价最小的等价表达式。
Equality Saturation 这个名字最早出现在 Tate、Stepp、Tatlock、Lerner 的 POPL'09 论文里，配套系统叫 Peggy。技术上是完整的，但当时的 e-graph 实现继承自 1970 年代定理证明器（Nelson、Oppen、Downey-Sethi-Tarjan），在 rewrite 密集的 saturation 工作负载下又慢又难扩展。这十年里它主要作为定理证明器内部组件存在（Z3、Simplify 都用 e-graph 做 congruence closure），独立做编译优化的工作不多。Denali、Equality Saturation 本身、Herbie 是仅有的几个代表，圈子很小。
egg 在 POPL 2021 拿了 Distinguished Paper。它不是发明新概念，而是把 e-graph 工程化到了别人能拿来直接用的程度。它通过 e-class 分析支持把附加信息（如 tensor shape、cost）与 e-class 绑定，让 cost-aware 的 rewrite 与抽取成为可能。后面的很多工作都是基于egg展开的。
Tensat（MLSys'21）把 Equality Saturation 用到了 DNN 计算图优化上。其做法是：复用 TASO 综合的图替换规则（包括 single-pattern 和 multi-pattern rewrite），用 egg 构建 e-graph 来代替 TASO 的回溯搜索，用 ILP 做最终的 graph extraction。实验表明 Tensat 覆盖了比 TASO 更大的搜索空间，并在更短时间内达到更好的性能。这体现了 e-graph 抽象的特点，phase-ordering 问题被结构性地消除：所有 rewrite 同时应用，最优选择推迟到最后一步。
3.3 优点和缺点
优点
结构性消除 phase-ordering
启发式与 AI 搜索面对 phase-ordering 的方式都是"在更大的决策空间里搜更好的顺序"。回溯搜索枚举不同顺序，RL 学习长期价值函数，MCTS 用 UCB 平衡探索与利用。这些都是优化算法层面的努力。
e-graph 不在这个框架内。它通过非破坏性 rewrite 让所有等价形式同时存在，应用顺序维度被直接消去而非更聪明地搜索。这是降维而不是优化。Phase-ordering 严重的场景（规则之间相互干扰、局部 cost 不能预测全局性能），e-graph 的优势会很明显。
可组合可扩展
加一条新的 rewrite 规则到 e-graph，不会影响已有规则的正确性，也不需要重新训练任何模型。等价类只会更大，永远不会更小。
RL agent 增加一个动作通常要重新训练，因为策略网络的输出维度变了，且新动作可能改变状态分布。启发式方法加规则也要重调权重或优先级，特别是当规则之间有冲突时。这种"加东西不破坏旧东西"的性质让 e-graph 适合做长期演化的工程系统：项目早期可以只有几条规则，随着需求增长慢慢加规则集。
变换可信且可追溯
E-graph 上每一条变换都对应一条 rewrite 规则。规则本身是局部等价命题，可以形式化验证；congruence closure 把局部等价扩展到全局。规则正确，结果就保证语义等价，不需要再跑 verifier。同样的机制让推理过程留下完整路径：extraction 给出的最终项由哪些 e-node 组成、每个 e-node 来自哪条规则，都可以反查。
启发式方法的正确性依赖 schedule 原语的设计。TVM、Halide 的 schedule 操作不会改变语义，但这个性质是隐式的，扩展新原语时容易出错。RL 方法靠环境验证，agent 提出动作后编译器实际执行检查，越界动作导致编译失败；学到的策略是黑盒，给出奇怪的 schedule 时不容易判断是 bug 还是优化。LLM 方法的输出可能语法正确但语义错误，进入生产前需要外部 verifier，过程也难追溯。
编译器对正确性敏感、调试密集，e-graph 在这两点上的优势是结构性的。
缺点
只适合离散规则空间
E-graph 的根基是"有限的 rewrite 规则集合"和"项的离散结构"。Schedule 调参面对的是连续参数（tile size 可以是任意 2 的幂、unroll factor 可以是 1 到 N 任意整数）和组合空间（多个参数的笛卡尔积），把这些参数全部塞进 e-graph 既不现实也无收益。
这一层启发式加 learned cost model 更直接：参数空间用结构化方式枚举，cost model 给每个候选打分，进化或贝叶斯优化做选择。RL 也可以处理（动作是参数选择），但工程实践中进化算法常常更稳定。这就是为什么 Ansor、MetaSchedule 这条路至今仍是工业主流，e-graph 在这一层基本缺席。
自身会指数爆炸
E-graph 的压缩性依赖"高共享、规则组合性强"。一旦这个前提被破坏，e-graph 自身会指数级膨胀。Multi-pattern rewrite 是典型触发器：一条规则要求 e-graph 中同时存在两个模式才能应用，新加入的 e-node 又会触发更多 pattern 匹配，迅速失控。
启发式与 RL 在状态空间里走单条路径，不论搜索多深都只占线性内存。e-graph 把整个等价空间持有在内存里，规模一大就走不下去。Tensat 不得不限制 multi-pattern rewrite 只跑 1 到 2 轮，MCTS-guided EqSat 等工作的核心动机也是控制 e-graph 自身的展开速度。这是 e-graph 至今难以处理"规则数量 100+ 的复杂工业场景"的主要原因。
Extraction 是 NP-hard
Saturation 完成后还有一个找最优等价表达式的问题。在带共享的 DAG 上做 cost-aware extraction 等价于 weighted DAG covering，一般情况下 NP-hard。贪婪 bottom-up extraction 速度线性但可能陷入局部最优。ILP extraction 给出最优解但只能处理几千 e-class 量级。
启发式与其他ai方法给出的本来就是"一条解"，没有"是否全局最优"的问题。e-graph 把找最优解变成显式步骤，反而暴露了这个 NP-hard 难题。MCTS 引导的 extraction 是较新的折中，但还在研究阶段。
规则要人工写
E-graph 的能力上限完全由规则集合决定。规则要么是数学等式（结合律、分配律），要么是领域知识（fusion pattern、硬件特定的等价变换），后者完全依赖专家。（这一点太烦了，最近我都成手写rewrite仙人了）
RL 与 LLM 在这一点上有结构性优势：它们可以从数据中学习，不需要显式规则。AlphaTensor 在矩阵乘法分解空间里搜索时，没有人告诉它"应该试什么变换"，奖励信号自动引导它发现规则。LLM Compiler 看了 5460 亿 tokens 的 LLVM 代码，隐式学到了大量优化模式。Ruler 这类规则合成工作正在让 e-graph 也能"自动发现规则"，但成熟度和覆盖面距离 RL 与 LLM 还有差距。
接入成本高
把 e-graph 接入一个已有编译器，需要把 IR 翻译成 e-node 表示、定义 cost function、写规则集合、做 saturation 配置（迭代次数、内存上限、规则优先级），最后还要把 extraction 出来的项回填到原 IR。每一步都不难但加起来是几人月的工程量。
启发式方法的接入是写一个搜索循环，RL 的接入是包一个 Gym 环境，相对都更轻。这就是为什么 e-graph 在新写的工具（Tensat、Cranelift mid-end）里出现得多，而在 LLVM、TVM 这类已有大型代码库的主线里推广较慢。它更适合在新项目立项时就采用，而不是中途切换。
4. Schedule 空间的搜索：启发式方法与 learned cost model
4.1 schedule 抽象
张量程序层最关键的建模来自 Halide（PLDI'13）和 TVM（OSDI'18），它们共同确立了 compute/schedule 分离这套抽象。compute 描述算什么——GEMM 就是几行循环加乘累加；schedule 描述怎么算——tile 多大、循环什么顺序、要不要向量化、要不要流水线、在哪一级缓存中间结果。同一份 compute 配上不同的 schedule,性能可以差几十倍。
这个分离把"搜索一个高性能 kernel"重新表述成了"搜索一个好 schedule"。为什么这一步关键？因为它用结构化的方式定义了整个张量程序层的搜索空间:每个 schedule 维度都是离散参数,参数之间的合法组合由 schedule 原语保证不破坏语义。后续所有自动化工作都建立在这套建模上,本质都是同一件事——在 schedule 原语张出的空间里找好解。
但空间一旦明确,就发现它非常大。一个 GEMM 的 schedule 空间在 10¹⁵ 量级以上。这里值得说清楚的是,"手写 kernel"这条路从未被工业界放弃——cuBLAS、cuDNN、oneDNN 至今仍是性能基线,它们就是高度调优的手工产物。被放弃的是另一种范式:让专家为每个算子、每种 shape 单独填一份固定 schedule。这种穷举式人工编程不可扩展。后续方法要解决的核心问题,是怎么把这个空间自动地搜得动。
4.2 AutoTVM:用 cost model 把搜索从硬件上解放出来
让搜索走不动的瓶颈是硬件测量太慢。一次测量一个候选 kernel 要跑实际硬件、要预热、要多次取中位数,单次几百毫秒到几秒。空间一大,时间预算瞬间耗尽。
AutoTVM(NeurIPS'18)的做法是引入一个学习到的 cost model 来替代大部分硬件测量。开发者先写一份含可调参数的 schedule 模板(tile_size、unroll_factor 这些),系统在硬件上跑少量样本训练一个 GBM 或小神经网络作为 cost model,再用 cost model 在大量候选里筛 top-k,只把 top-k 送到硬件上实测。实测结果反馈回 cost model 继续训练,迭代到预算用完。具体的探索策略支持随机、网格、模拟退火、GA 等几种,模拟退火是默认推荐。
这套"少量真实测量加大量代理评估"的范式把单次评估成本从几百毫秒压到几毫秒,搜索规模随之放大几个数量级。它至今仍是绝大多数 auto-tuner 的骨架。
但 AutoTVM 留下了一个明显的限制:模板还是要人写。每来一类新算子(Conv2D、GEMM、Pooling、LayerNorm)都要专家写一份模板,模板的写法又决定了搜索空间的边界。写得保守就丢失优化机会,而专家也未必能预想到所有有价值的组合。
4.3 Ansor:把模板这一步也自动化掉
Ansor(OSDI'20,在 TVM 里叫 Auto-Scheduler)的目标就是去掉人工模板。它用 hierarchical sketch 从 compute 定义直接生成 schedule 骨架——骨架描述结构性选择(几级 tile、哪些循环并行、在哪一级做 cache read),骨架内部的具体参数用进化搜索来填。Cost model 仍然负责评估,再加一个 task scheduler 在多个子图之间动态分配搜索预算,把算力优先花在对端到端延迟影响最大的子图上。
Ansor 选进化搜索而不沿用 AutoTVM 那套探索策略,根本原因不在算法本身有多优越,而在于它要搜的空间大了一个量级——因为不再受手写模板边界限制。进化算法通过变异和交叉在大空间里跨距走,又对 cost model 的噪声相对鲁棒,配合得很好。
实验数据上,Ansor 在 Intel CPU、ARM CPU、NVIDIA GPU 上分别取得了最高 3.8×、2.6×、1.7× 的加速。论文里有一个值得单独拎出来的观察:Ansor 找到的最优 schedule 落在已有 search-based 方法的搜索空间之外。也就是说,进步不是在同一个空间里搜得更快,而是空间本身变大了,覆盖了手写模板根本触及不到的优化组合。这条洞察在本章后面还会反复出现:搜索空间的表达力是这一类方法的真正天花板,搜索算法的聪明程度反而是次要的。
Ansor 之后这条路又往前走了几步。TVM 推出 MetaSchedule,把 schedule 编程语言化,让 sketch 和搜索策略都用同一套语言来表达,进一步降低写新 schedule 系统的门槛;FamilySeer 注意到不同子图结构常常相似,让它们共享 cost model 来提升搜索效率。整个范式的演化方向是清晰的:人写得越来越少,自动化覆盖得越来越深。
4.4 Cost model 是这套范式的命门
整套方法的有效性都押在 cost model 上。cost model 不准,进化搜索就在错误方向上空耗预算;cost model 不通用,每换一类硬件就要重训。围绕这个核心,工程上常见几种做法:把 schedule 编码成定长特征向量(loop counts、memory access pattern、arithmetic intensity)后用 XGBoost 或 LightGBM 训练;用 TreeLSTM 或 GNN 直接处理 schedule 的树形结构;把 schedule 序列化后用 Transformer 编码(One-Shot Tuner 这类);或者混合分析模型加 ML 残差修正。
跑下来真正绕不开的问题其实是连在一起的。首先是真实测量本身有噪声——CPU 后台进程、GPU 频率调整都会让同一 kernel 多次测量差几个百分点,得多次取中位数才稳定。其次是 cost model 在搜索过程中会漂移:早期采样的低质量候选和后期高质量候选分布不一样,没有持续反馈训练,cost model 会越搜越偏。两者叠加,让跨硬件迁移变得很难——同一段代码在 V100、A100、H100 上的最优 schedule 完全不同,而 cost model 又对噪声和分布敏感,几乎只能 per-device 训练,预训练那点初始化红利很快就消耗掉了。
4.5 性能墙:搜索空间的表达力才是天花板
cost model 的局限直接导致一个现象:在主流硬件上的主流算子(FP32/FP16 GEMM、Conv),auto-tuner 的最优结果常常仍然不如 cuBLAS 和 cuDNN,瓶颈集中在 Tensor Core 的利用率上。原因是 Tensor Core 编排、warp scheduling、bank conflict 这些硬件细节对性能影响很大,但传统 cost model 的特征空间根本捕捉不到。
后续工作的演化路径反过来印证了这一点。Roller(OSDI'22)、AMOS(ISCA'22)、TensorIR(ASPLOS'23)这一批方法的共同点,是直接把 Tensor Core 抽象进搜索空间——不是用更聪明的搜索算法,而是把硬件层级(thread block tile、warp tile、MMA fragment)作为可调维度暴露出来。一旦这个维度纳入搜索,启发式方法在 FP16 GEMM 上能逼近甚至追上 cuBLAS。Triton 和 CUTLASS 走的是另一条路,给开发者更高级的硬件抽象,让人在合适的颗粒度上把硬件细节写出来。两条路得出同一个结论:搜索的上限受限于搜索空间的表达能力,而不是搜索算法的聪明程度——这正是 4.3 节那个观察的延续。
但即便搜索空间表达力扩展了,还有一个工程现实绕不过:编译时间。AutoTVM/Ansor 调一个完整模型常常要几小时到几天,搜得越深、空间越大,时间越长。这是它们在很多实际场景被绕过去的核心原因——TorchInductor、torch.compile 这一路用更轻量的模板加少量 tuning,牺牲峰值性能换可接受的编译时间,在生产环境反而更受欢迎。
尽管打不过手写库、调起来又慢,启发式加 learned cost model 至今仍是研究型生产环境的首选。理由很务实:不依赖预训练大模型,不需要 RL 大量探索;行为相对可解释,cost model 出问题能 debug;对未见硬件适应性较好,只要能跑实测就能训出新 cost model。TVM、XLA、IREE 这些主流框架都仍以这一范式为基础——不是因为它最强,而是因为它在"够用、可控、可维护"这三件事上同时合格。
这一章后面要讲的 AI 搜索方法和 LLM 方法,每一个都在某一维度上比启发式更强,但都还没在这三件事上同时合格。
5. AI 搜索方法:RL、MCTS 与进化算法
启发式方法把"搜索哪个候选"交给某种固定策略(贪婪、进化、模拟退火)。RL 与 MCTS 把这一策略本身变成可学习的对象。沿着上一章的判断——搜索空间的表达力才是天花板——RL/MCTS 这条线的关键贡献不只是把搜索做得更聪明,而是把"什么算可以搜的东西"也重新定义了:从一步一个 schedule 参数,扩展到一整条长决策序列。
5.1 RL 与编译优化的契合点
编译优化在结构上贴合强化学习。状态是当前 IR、e-graph 或 schedule;动作是应用某条 rewrite、选择某个 pass、设定某个 tile size;奖励是编译完成后的运行时,或代理 cost model 估计;一个 episode 是从原始程序到最终 lowered code 的整条路径。
真正让 RL 在这一类问题上有独特价值的是长程信用分配(long-horizon credit assignment)。某个看起来局部不利的 rewrite——比如先做一个看似"反向"的代数化简——可能为后续更激进的优化创造条件。贪婪不会选它,因为短期指标变差;进化算法可以歪打正着撞到,但没有显式机制把"为了第 8 步的收益主动接受第 2 步的损失"这种逻辑学进策略里。RL 的 value function 是为这件事而生的。
5.2 X-RLflow
X-RLflow(VLDB'23)是这条思路在图层 rewrite 上的一个直接示例。它针对 TASO 做替换,把 cost-based 回溯搜索换成深度 RL agent:用 GNN 编码当前 tensor 计算图,输出每条 rewrite 规则的 Q 值,逐步决定下一步用哪条 rewrite。实验表明 X-RLflow 在多类工作负载上超过 TASO,原因正是 agent 可以为长期收益放弃短期收益,这是纯贪婪和 cost-based 搜索在结构上做不到的。
5.3 AlphaTensor:把"用什么算法"也变成搜索对象
DeepMind 的 AlphaTensor(Nature, 2022)把"寻找矩阵乘法的快速算法"建模为一个单人游戏。状态是一个三维 tensor,初始化为矩阵乘法对应的 Tn,m,pTn,m,p;动作是从张量上"减去"一个秩-1 项 u⊗v⊗wu⊗v⊗w;当 tensor 变为 0 时游戏结束,使用的秩-1 项数即为该算法的乘法次数;奖励为负的步数。
这是一个标准的 AlphaZero 式设置——神经网络加 MCTS。结果是 AlphaTensor 在 4×4 二元域矩阵乘上找到了一个 47 次乘法的算法,超过了 1969 年 Strassen 算法在该设置下保持的 49 次乘法的纪录。它还能针对特定硬件(V100、TPU)的实测延迟做优化,找出比标准实现更快的算法。
从 AI 编译角度看,AlphaTensor 的意义不在"打破 Strassen 纪录"这件具体事上,而在于它把搜索的边界推到比 schedule 更基本的地方——连"用什么算法实现这个算子"本身都被纳入搜索空间。上一章把搜索空间表达力作为天花板,这里就是一次显式的天花板抬升:不再只在给定算法的实现里选 schedule,而是连算法本身都可以重选。代价是搜索空间从 10151015 量级直接跳到组合爆炸,只有 MCTS 加神经网络这种规划能力强的组合才能 navigate。
5.4 LLVM 层与 MLIR 层的 RL
经典编译领域,CompilerGym(CGO'22)把 LLVM phase ordering、GCC flag selection、CUDA loop nest generation 等任务封装成 OpenAI Gym 兼容的 RL 环境,允许 RL 研究者直接在编译问题上评估算法。它已经成为这一方向的常用基准。
更近的 MLIR RL 工作(CGO'26)把 RL 应用到 MLIR Linalg 方言层,提出了一种多动作 RL 形式化(动作空间是若干简单子动作的笛卡尔积),以及一个减少 loop interchange 动作空间大小的 level pointer 技巧。这类工作的共同动机是,MLIR 这种多方言 IR 给 RL 提供了比纯 LLVM 更结构化的状态和动作语义——RL 不擅长在状态/动作语义模糊的空间里学,IR 越规整,样本效率越好。
5.5 进化算法与贝叶斯优化:不是所有"AI 搜索"都需要 RL
工程实践里,进化算法加 learned cost model(就是 Ansor 那套)常常比纯 RL 更稳定、调参更少。原因是结构性的:进化算法对奖励噪声更鲁棒,不需要 advantage estimation 这套数值敏感的机制;评估天然并行,CPU/GPU 资源利用率更高;失败模式简单——无非"种群退化"——诊断成本远低于 RL 训练崩溃。
RL 真正划算的场景是长 episode、稀疏奖励、需要规划的任务,比如 phase ordering、graph-level rewriting——这些任务要连续应用很多变换,最终才能评估,贪婪和进化都吃亏。短 episode、密集奖励的 schedule 调参,进化算法或 Bayesian Optimization 通常就够,把 RL 套上来反而是用复杂度换不到多少收益。
5.6 AI 搜索方法的代价
代价有几个绕不开的。训练成本高,AlphaTensor 训练用的算力远超普通 auto-tuner;泛化性有限,一个 RL agent 在 ResNet-50 上学会的策略未必能直接迁移到 ViT;实现复杂,RL 训练涉及环境、奖励设计、回放缓冲、稳定性技巧,复杂度高于一个 evolutionary search 循环;可解释性低,当 agent 给出一个奇怪的 schedule,工程师难以判断是 bug 还是真的好。最后这一点直接踩中 4.5 节那条"够用、可控、可维护"标准里的"可控"和"可维护"两项。
这些代价决定了 RL 和 MCTS 目前主要在研究层和少数高价值场景(手写 GEMM、超大规模模型部署、算法发现)使用,尚未成为主流编译器的默认策略。它在搜索空间表达力上确实更强,但还没能同时在三件事上合格。
6. LLM 辅助编译优化
2023 年以来 LLM 进入编译领域的速度较快（可以说我参与了llm自动设计芯片的最早一个项目了/狗头）。这里要区分两类用法:LLM 作为代码生成器,以及 LLM 作为搜索策略。两者放在 4.5 节那个三件事的尺子下看,跑分模式很不一样。
6.1 LLM 直接生成优化代码
早期标志性工作之一是 Cummins 等人(arXiv 2309.07062)的 Large Language Models for Compiler Optimization。他们训练了一个 7B 参数 transformer,输入未优化的 LLVM 汇编,输出最佳的编译选项序列与优化后代码。
几个数据点值得记下来。训练时让模型同时预测优化前后的指令数与优化后代码本身,这两个辅助任务提升了主任务质量。测试集上,模型相对编译器减少 3.0% 指令数,超过了两个需要数千次编译的 baseline。模型生成的代码在 91% 的情况下可编译,70% 的情况下与编译器输出完全一致。
这一结果表明,LLM 在某种程度上学到了编译器内部的优化逻辑——它不是在"猜"输出,而是吸收了相当一部分 IR 层的因果关系。
6.2 LLM Compiler
2024 年 6 月 Meta 发布的 LLM Compiler 把这条路线推到了更系统的位置。它基于 Code Llama,在 5460 亿 tokens 的 LLVM IR 与汇编代码上做继续预训练,再在一个 compiler emulation 数据集上做 instruction tuning。能力包括:代码大小优化(达到 autotuning 搜索潜力的 77%)、反汇编(从汇编恢复到高级 IR)、优化 pass 推荐(给定 IR,推荐最佳 pass 序列)。模型配以宽松商业许可证发布。
LLM Compiler 代表了一种不同于 RL 或搜索的范式:把编译器优化当作翻译问题——从未优化代码翻译到优化代码、从汇编翻译到 IR,用语言模型的强先验加速这个过程。这条路的搜索空间表达力理论上最广(自然语言加全部代码),但代价是丢失了符号方法的等价性保证,从优化转成了概率猜测。
6.3 LLM 与 RL 的结合:Compiler-R1
更新的方向把 LLM 与 RL 接到一起。Compiler-R1(arXiv 2506.15701, 2025)训练一个 LLM agent 在 CompilerGym 环境中做 LLVM pass 选择,使用两阶段端到端 RL 流水线:先用高质量推理数据集做 SFT,再用基于 outcome 的奖励做 RL。结果显示在 7 个数据集上相对 opt -Oz 平均减少 8.46% 的 IR 指令数。
Compiler-R1 架构上有两个关键点。一是它不让 LLM 一次性预测完整 pass 序列,而是让 agent 与编译环境多轮交互——查看当前 IR、调用工具(如 instrcount)、决定下一动作。这把 6.1 节那种"一次性生成"扩展成了带反馈的闭环,LLM 可以看到自己每一步的效果再决定下一步,大大降低了对单步预测准确率的依赖。二是 RL 训练教会 agent 什么时候用工具、什么时候停,这是纯 SFT 学不到的。
这是 agentic compiler tuning 的早期范式:把 LLM 当作具备工具使用能力的 agent,让它和编译器、profiler 形成闭环。在编译这种"有客观地面真相(运行时/指令数)"的领域,agent loop 的可用性比许多通用 agent 场景要好得多——奖励信号不含糊,失败可以验证,这两点已经足够让 RL 训练落地。
6.4 LLM 路线的正确性问题
LLM 在编译领域最大的瓶颈是正确性。语言模型可能输出语法正确但语义错误的优化代码——改变了浮点行为、违反了 alias 规则、丢失了 volatile 标记。即便正确率 91%,9% 的失败率在生产编译器里就是不可接受。这
目前几条缓解思路:
Verifier-guided generation:用形式化等价检查器(Alive2、translation validation)过滤 LLM 输出。
LLM 作为搜索策略而非生成器:让 LLM 输出变换的描述(例如"在第 3 个循环外做 tile,尺寸 32"),由可信的编译器实施变换。这等于把 LLM 的输出限制在一个由编译器原语张成的离散空间里,绕开了任意代码生成的正确性问题。
混合 LLM 与符号方法:LLM 提议候选 rewrite,符号系统(e-graph)保证等价性。
RL 的 outcome reward:让训练中的失败惩罚替代显式验证,Compiler-R1 是一个例子。
这些方向都还早,但共识已经收敛:LLM 必须与某种正确性保证机制结合,才能进入主流编译流水线。直接用 LLM 一次性生成最终 IR 这条最浪漫的路,在可预见的未来都不太可能成为主流。
7.总结：
总结一个表格
由此可以提炼出几条经验性的对应关系。
Graph-level rewrite(DNN 计算图重写)适合符号方法。规则数量有限、组合性强,e-graph 的优势在这一层最明显。当 e-graph 爆炸时,引入 MCTS 或 RL 引导。
Tensor program / schedule level 仍以启发式 + learned cost model 为主流。Ansor、MetaSchedule 这条路在工业界稳定可用。RL 在这一层的额外收益尚未充分证明能 justify 它的复杂度。
LLVM-level pass ordering 适合 RL 与 LLM。CompilerGym 环境成熟,问题离散化彻底,样本效率不算差。Compiler-R1 这类工作显示 LLM-RL 混合在这一层也有进展。
算法发现(algorithmic level)目前主要由 RL+MCTS 承担(AlphaTensor)。这是搜索可以触及的最深层次,也是最昂贵的层次,目前只在矩阵乘法这类有干净游戏结构的问题上有突破。
跨层优化与新硬件适配方向上,LLM 辅助路线在未来两三年可能带来变化——LLM 是目前唯一能消化非结构化文档(硬件手册、白皮书)的搜索和生成机制。一个新加速器发布,几千页 PDF 加一个早期 SDK,启发式方法和 RL 都要等数据集和 cost model 才能起步,LLM 至少能读着文档先生成一版可跑的 kernel。
