从手工调策略到自动合成——北大团队用LLM解决e图爆炸难题

编译器优化的终极目标，是在有限的计算资源里找到最优的程序实现。传统的逐条重写规则就像贪心算法，走一步看一步，很容易错过全局最优解。等式饱和技术的出现改变了这一局面——它把所有可能的程序变体压缩进一张叫做e图的数据结构里，在搜索结束后才挑出成本最低的那个。理论上，这是一种接近完美的优化范式。

但现实远比理论残酷。当重写规则数量增长到成百上千条时，e图的规模会以指数级膨胀，吃光内存、拖垮性能。要让等式饱和在真实的编译场景中跑起来，光有规则远远不够，还得有一套精心设计的策略来控制搜索过程。目前，这些策略几乎完全依赖人类专家手工编写，费时费力还难以迁移。

北京大学的尹晨昀、肖有为、罗宇泽、邹雨阳和梁云团队在最新论文中提出了EggMind框架，用大语言模型来自动合成等式饱和策略。这套方法在向量化编译基准上将最终代价降低了45.1%，峰值内存减少了69.1%，运行速度提升2.21倍，并且在XLA张量编译和逻辑综合场景中同样有效。

等式饱和为什么好用又难用

等式饱和的核心机制并不复杂。给定一组重写规则，系统反复将这些规则应用到e图上，直到没有新的等价关系可以发现为止。最后用一个提取器从e图中选出成本最低的程序。

传统的破坏式重写有一个根本缺陷：每次应用一条规则，之前的形式就被覆盖了。这意味着如果你先选了A方案再应用规则R1，后续的规则R2可能已经找不到它需要的中间形式了。等式饱和解决这个问题的方式很优雅——它用e图同时保留所有中间形式，让不同规则的搜索路径可以在任意节点汇合。

这种机制带来了两个根本性优势。第一，它不会过早地丢弃潜在有用的中间形式；第二，它可以发现需要多条规则协同作用才能产生的优化。比如在向量化编译中，一条乘加融合优化可能需要先做向量提升、再做常量折叠、最后做指令选择，三个步骤按特定顺序执行才能生效。破坏式重写很难捕捉这种跨步骤的组合优化，因为中间步骤的结果往往会被后续重写覆盖掉。

但代价也很直观：每次应用重写规则，e图都可能膨胀。规则之间还会互相触发，产生连锁反应。一条简单的加法交换律可能产生新的中间形式，而这些新形式又能触发另一条规则，形成雪崩式的增长。论文中的数据很有说服力：在矩阵乘法162x162的基准上，完整等式饱和的峰值内存高达数GB，运行时间超过600秒超时。对于更大的输入规模，内存直接爆炸，完全无法完成搜索。

图1：完整等式饱和会导致e图爆炸，而精心设计的策略能在有限资源下逼近最优解。

图改编自原论文 Figure 1

图2：直接让LLM生成等式饱和策略代码会遇到缺少可复用抽象、缺少有效反馈、缺少可控性三个核心障碍。

图改编自原论文 Figure 2

手动策略的瓶颈

业界对此的应对方式是设计分阶段的执行策略。以向量化编译领域知名的Isaria系统为例，人类专家将重写规则分成若干组，规定它们的执行顺序和重复次数，并在关键位置插入e图收缩操作。这套策略在特定基准上效果不错，但它完全手工打造，换个问题就得重来。

论文中的Table 1清晰地展示了当前策略引导系统的设计空间碎片化现状。Isaria提供可复用的调度方案，但依赖人工设计；基于MCTS的方法可以自动化搜索，但产出的策略不可复用，每个新用例都要从头搜索；引导式方法能通过外部信号来定向搜索，但同样无法生成可复用的策略工件。没有一个现有系统同时满足自动化、LLM引导和可复用这三个条件。

更棘手的是，近年来出现了一些自动规则合成工具，比如Enumo和Ruler，它们能从语义规范中推导出大量重写规则。规则数量上去了，搜索空间更大了，手工设计策略变得更加不可行。这些工具的出现等于在等式饱和的供给侧注入了大量新规则，但需求侧的策略设计能力没有跟上。

用大语言模型来自动化策略生成，看上去是一条可行的路。让LLM直接编写egglog后端代码确实能跑通，但效果很不稳定。问题出在三个地方：

第一，后端代码里只有底层指令，看不到高层策略语义。LLM生成的代码缺乏可复用性，换个测试用例就得从头来过。

第二，反馈信号要么太粗要么太细。最终代价、运行时间这些指标无法告诉模型哪里出了问题，而完整的e图状态数据量太大，塞不进上下文窗口。

第三，没有约束机制。LLM动不动就生成一条把所有规则一股脑跑完的策略，或者把e图收缩得太早，搜索还没展开就结束了。

EqSatL：让策略成为可读、可复用的工件

EggMind的核心设计思想是把策略从不透明的后端脚本提升为一个有结构、可检查、可复用的领域特定语言，叫做EqSatL。

EqSatL把策略设计分解成三个控制层面。

规则集划分解决的问题是：哪些重写规则应该放在一起执行。EqSatL用语义标签来组织规则，而不是简单的规则列表。系统会请LLM把领域内的重写规则按照功能分组——比如向量化中的暴露规则、编译规则、优化规则各成一组。标签只需定义一次，后续的策略合成就在这个标签空间里操作。

调度构建解决的问题是：这些规则组按什么顺序执行、执行几轮。EqSatL用一棵流树来表达调度方案，支持顺序执行、重复循环和阶段性简化。每个阶段结束后必须执行一次e图收缩操作，这就像给搜索设了一个安全阀——允许跨阶段的交互探索，但不让e图失控增长。

简化控制解决的问题是：收缩e图时应该保留什么、丢弃什么。EqSatL引入了一种基于偏好的简化机制，让LLM为每个阶段建议一组优先保留和可以裁剪的模式，而不是全局一刀切。

图3：EqSatL将策略表达为规则集划分、流调度、简化控制三个层面的组合。

图改编自原论文 Figure 4

有了EqSatL，策略不再是散落在脚本里的底层指令，而是一个结构清晰的工件。同一个策略可以通过后端降级生成多个针对不同测试用例的egglog脚本，实现跨用例复用。

Agentic工作流：搜索策略的策略

EqSatL定义了策略的语言，那怎么找到好的策略呢？EggMind设计了一个离线合成的智能体工作流。

整个过程是迭代式的。每一轮迭代中，工作流维护两样东西：一个策略仓库和一个重写动机缓存。策略仓库分为当前最佳、有潜力待改进、刚生成待评估三类。重写动机缓存记录的是从成功运行中提炼出来的结构化证据。这种设计确保了每次迭代不是孤立的探索，而是在已有证据基础上的增量改进。

每轮迭代的流程是：策略师从当前上下文中选择下一步动作——可能是提议一个全新的策略结构，也可能是对某个有潜力的策略进行微调，还可能是向分区器或简化器请求额外的指导。选定动作后，工具池里的四个专用智能体各司其职。生成器提出新策略，评估器支持三种预算模式——全预算评估用于最终比较，缩减预算用于参数调整，最小预算用于快速筛选。分区器和简化器提供可操作性指导，帮搜索避开发散的陷阱。评估完成后，表现最好的策略更新最佳槽位，有潜力但需要继续改进的策略保留在候选池中，新生成但尚未评估的策略标记为待处理。

图4：每轮迭代中，智能体工作流提出策略、评估执行、更新缓存，逐步积累可复用的证据。

图改编自原论文 Figure 5

这套设计的巧妙之处在于，它把原本开放式的代码生成问题转化成了一个结构化的搜索问题。搜索的对象不是任意代码，而是EqSatL策略工件；搜索的证据不是原始日志，而是经过提炼的重写动机；搜索的约束不是硬性规则，而是可调的可操作性指导。

从证明树中提取经验

EggMind的第二个核心创新是证明导出的重写动机缓存。

当一条策略成功运行并找到了更优的程序实现时，egglog后端会生成一棵证明树，记录从输入表达式到最优输出之间的重写路径。这棵证明树蕴含了丰富的信息，但直接拿来做反馈太大太杂。

EggMind的做法是从证明树中提取局部的重写动机——也就是一小段有意义的规则链。这些动机被提升到EqSatL的标签层面，去重后存入缓存。每条动机都附带一个可量化的效用分数，表示它在多大程度上降低了最终代价。

在后续的策略合成中，生成器会参考缓存中的高分动机。这相当于让LLM从过去的成功经验中学习，而不是每次从零开始猜测。

缓存的维护也有讲究。随着证据积累，缓存会根据效用分数进行策展：高分动机被优先保留，低分动机被淘汰。这确保了缓存始终包含最有价值的参考信息。

图5：EggMind在向量化基准的进化用例和留出用例上，代价和运行时间均优于完整等式饱和和Isaria基线。

图改编自原论文 Figure 7

可操作性指导：给搜索加安全围栏

第三个创新点是可操作性指导，包括两个组成部分。

分区建议来自分区器智能体。它构建一个规则依赖图，用组合优化算法来评估不同分区方案的质量。目标函数奖励前向激活、惩罚反向流动和同组纠缠，从而生成的调度方案倾向于有序的前向探索，而不是混乱的全面展开。

简化提示来自简化器智能体。它为每个阶段提出一组偏好/裁剪模式对。在e图收缩时，已经匹配了偏好模式的节点会被优先保留，匹配裁剪模式的节点会被降级处理。这种机制不是硬性删除，而是通过附加惩罚分数来温和地引导简化方向。

两者共同构成了搜索的安全围栏。分区建议防止规则之间产生不可控的交互，简化提示防止e图在某个阶段过度膨胀。论文中专门设计了一个高增长场景来测试简化提示的效果：当暴露规则和向量优化规则被混在同一个粗粒度阶段中运行时，启发式简化因为看不到后续需要哪些中间形式，容易做出短视的裁剪决策。而有了简化提示的引导，系统能够在正确的位置保留关键的中间形式，代价增量从不使用优化时的+55降低到了使用提示时的+15甚至+0，同时运行时间保持在合理的秒级范围。

实验结果：不只是向量化

论文在三个场景上验证了EggMind的效果。

向量化编译是主要评估基准，沿用了Isaria和Diospyros的标准测试集。EggMind合成的策略在15个测试用例中，最终代价的几何均值比完整等式饱和降低了45.1%，峰值内存降低了69.1%，运行时间加速2.21倍。合成过程本身耗时约13分钟，使用53次模型调用，花费约3.16美元——这是一次性成本，可以摊销到后续所有在线运行中。

消融实验揭示了各组件的贡献。去掉重写动机缓存后，代价恶化了78.3%；去掉重复构造后，代价恶化了118.7%，内存从639MB飙升到2.5GB。规则分区器的移除使代价恶化23.1%，但运行时间反而降低了63.3%，说明分区的主要价值在于提升优化质量而非节约搜索成本。

图6：去掉任何一个核心组件都会导致性能下降，动机缓存和重复构造的影响最为显著。

图改编自原论文 Figure 8

XLA张量编译是第二个评估领域。研究团队选取了XLA HLO的代数简化流程，通过扩展重写空间来构造更具挑战性的基准。在17个评估用例中，EggMind的策略在13个上优于无指导基线，其余4个持平，在线运行时间的几何均值加速了11.89倍。

逻辑综合是第三个场景，使用EqMap工具进行FPGA LUT重映射。研究团队用Enumo合成了布尔逻辑和物理单元之间的交互重写规则。一个具体的例子是：两条独立的AOI21单元可以通过一条合成规则直接合并为一个AOI22单元，将原本需要多步布尔归一化、因式分解和提升操作才能完成的变换压缩为单步变换。在carry-chain基准上，随着位宽从2增长到10，搜索空间急剧膨胀。EggMind合成的策略成功将搜索引导到了低代价的单元映射，而无指导搜索和自由代码演化均无法稳定达到同等效果。在位宽为5的配置下，EggMind的策略甚至逼近了理论最优解，代价偏差仅为0。

离线合成的真正价值

EggMind的设计哲学值得单独说一说。

传统的LLM辅助编译优化走的是在线路线——在每次优化任务中实时调用LLM来辅助搜索。这种方式灵活，但成本高、不稳定，且无法积累经验。

EggMind选择了离线路线。策略合成发生在部署之前，一次合成就能产出多个可复用的策略。到了在线阶段，策略直接在EqSatL层面执行，不再需要调用LLM。合成成本可以被大量在线运行摊销，而策略本身是可检查、可比较、可迭代改进的。

论文中的对比实验也印证了这一判断。研究团队让多个LLM智能体在没有EggMind框架的情况下自由搜索策略，分别测试了Kimi Code、Codex和Claude Code。结果发现，基于egglog后端的自由搜索全部无法找到有效策略，收益为零；而基于EqSatL的自由搜索虽然能找到一些有效的策略，但收益远低于EggMind的系统化合成。

这说明问题的关键不在于LLM有多强，而在于给LLM提供什么样的抽象层和反馈机制。EqSatL提供的可复用抽象、证明导出的重写动机提供的可操作反馈、以及可操作性指导提供的安全约束，三者缺一不可。

局限与展望

论文也坦诚地讨论了方法的局限性。目前的评估主要集中在三个领域，更广泛的适用性还需要进一步验证。EqSatL的设计假设重写规则可以按语义分组，对于规则之间高度纠缠的场景，这种分组可能不够自然。此外，分区器的组合优化求解在规则规模极大时可能成为瓶颈。

在逻辑综合的实验中，论文也观察到一个有趣的现象：当重写空间被Enumo大幅扩展后，原始的EqMap流程在大位宽下变得极其缓慢，而EggMind合成的策略能够在合理时间内找到接近最优的映射方案。这暗示了EggMind的一个潜在应用场景——作为规则合成工具链的下游消费者，自动消化和驾驭日益增长的规则空间。

从更大的视角来看，EggMind代表了LLM在系统领域应用的一种新范式：不是让LLM直接生成最终系统代码，而是用LLM来合成中间层的控制策略。这种思路可以推广到其他需要复杂搜索控制的编译和优化场景中——只要能把搜索过程表达为一种结构化的领域特定语言，LLM就有可能在这个语言空间里进行高效的策略合成。

结语

等式饱和是一种理论上接近完美的程序优化技术，但它对搜索策略的高度依赖一直是落地应用的主要障碍。EggMind通过引入EqSatL领域特定语言和LLM引导的智能体工作流，把策略设计从手工编码变成了可自动合成、可跨用例复用的结构化工件。证明导出的重写动机缓存让系统能够从成功经验中持续学习，可操作性指导则为搜索过程装上了安全围栏。

实验结果表明，这套方法在向量化编译、张量图优化和逻辑综合三个场景中均能有效改善资源-质量权衡。更重要的是，它把LLM的角色从在线辅助工具转变为离线策略合成器，使得合成成本可以被摊销，策略可以被检查和复用。

对于编译器开发者来说，EggMind提供了一条实用路径：先用Enumo等工具自动扩展重写规则，再用EggMind自动合成搜索策略，最终得到一个既全面又可控的e图优化器。这比手动为每个新问题设计策略高效得多。

参考资料

本文配图改编自原论文《LLM-Guided Strategy Synthesis for Scalable Equality Saturation》（北京大学，尹晨昀等），仅用于论文解读与学术交流。