💻 computer science

LLM-Guided Strategy Synthesis for Scalable Equality Saturation

本文提出了 EggMind 框架，通过引入领域特定语言 EqSatL 和结合证明推导重写模式缓存等技术的 LLM 引导智能体工作流，实现了可复用且高效的等式饱和策略自动合成，显著降低了优化成本并缓解了 e-graph 爆炸问题。

原作者： Chenyun Yin, Youwei Xiao, Yuze Luo, Yuyang Zou, Yun Liang

发布于 2026-04-21

📖 1 分钟阅读☕ 轻松阅读

原作者： Chenyun Yin, Youwei Xiao, Yuze Luo, Yuyang Zou, Yun Liang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文介绍了一个名为 EggMind 的新系统，它利用人工智能（大语言模型，LLM）来自动设计一种叫“等式饱和（Equality Saturation）”的编译器优化策略。

为了让你更容易理解，我们可以把整个过程想象成在一个巨大的迷宫里寻找最短路径。

1. 背景：什么是“等式饱和”？

想象你手里有一个复杂的数学公式（比如 a + b * c）。在计算机眼里，这个公式可以写成无数种等价的形式（比如 b * c + a，或者展开成更复杂的式子）。

传统方法：像是一个急躁的向导，一旦看到一种写法，就立刻决定“就按这个改”，然后继续往下走。如果第一步选错了，后面就回不来了，可能错过更好的方案。
等式饱和（EqSat）：像是一个超级耐心的探险家。它不急着做决定，而是把公式能变成的所有样子（成千上万种）都画在同一张巨大的地图上（这叫 e-graph）。等把所有路都画完了，它再根据“谁最省资源”的标准，从地图里挑出最好的一条路。

问题出在哪？
这张地图太大了！如果规则太多，地图会瞬间爆炸，把计算机的内存撑爆（就像你试图在一秒钟内画出整个地球的所有街道，电脑直接死机）。

2. 核心痛点：如何控制这张“爆炸”的地图？

为了让等式饱和能工作，我们需要一个**“策略”（Strategy）。这就好比给探险家一张分阶段的寻宝图**：

先只探索简单的路（第一阶段）。
把走不通的死胡同清理掉（简化）。
再探索复杂的路（第二阶段）。
最后再清理一次，选出最佳路线。

现在的麻烦是：设计这张“寻宝图”非常难。以前全靠人类专家凭经验手写，既慢又容易出错。而且，现在的 AI 能自动生成很多新的“改写规则”，规则越多，地图越容易爆炸，人类更不知道该怎么安排这些规则了。

3. EggMind 的解决方案：AI 当“总导演”

EggMind 就是一个由 AI 驱动的自动化导演系统，它不再让 AI 直接去写复杂的底层代码（那就像让 AI 直接去画每一块砖，容易画崩），而是让 AI 编写**“分镜头剧本”**。

它做了三件关键的事：

A. 发明了一种“剧本语言” (EqSatL)

以前 AI 直接写代码，就像让一个不懂乐理的人直接指挥交响乐团，容易乱套。
EggMind 发明了一种简单的**“剧本语言” (EqSatL)**。

比喻：这就好比给 AI 一个乐高积木说明书。AI 不需要知道怎么造每一块砖，只需要说：“先搭底座（Expose 阶段），再搭塔身（Compile 阶段），最后加个尖顶（Opt 阶段），中间记得把多余的砖头扔掉（Simplify）”。
这样，AI 就能专注于宏观策略，而不是陷入微观的代码细节。

B. 建立“成功案例库” (Rewrite Motif Caching)

AI 在试错时，经常不知道“为什么这个方案好”。
EggMind 会像一个聪明的图书管理员，把那些成功的“寻宝路径”提炼成**“模式”**（Motif）。

比喻：如果 AI 发现“先走 A 路再走 B 路”通常能省钱，它不会把整条路都记下来，而是记下一个标签：[先 A 后 B] = 省钱。下次 AI 再设计剧本时，就会直接调用这个“省钱秘籍”，而不是从头瞎猜。

C. 安装“防崩溃保险” (Tractability Guidance)

AI 有时候太激进，可能会设计出一个让地图瞬间爆炸的剧本（比如“把所有路一次性全打开”）。
EggMind 给 AI 配了一个**“安全顾问”**。

比喻：当 AI 想写一个疯狂的剧本时，安全顾问会提醒：“嘿，这个组合会让内存爆炸！我们换个温和点的顺序吧。”它通过计算规则之间的依赖关系，确保 AI 设计的剧本是安全、可控的。

4. 结果如何？

实验证明，EggMind 非常成功：

更省钱：它找到的优化方案，比完全放任不管（Full EqSat）省了 45% 的“代价”（比如运行时间或能耗）。
更省内存：它让计算机在优化过程中少用了 69% 的内存，避免了电脑死机。
更通用：这套方法不仅适用于向量计算，还能用到芯片设计（XLA）和逻辑合成等其他领域。

总结

EggMind 就像是给编译器优化请了一位“超级 AI 导演”。
它不再让 AI 去死磕具体的代码细节，而是教 AI 如何写剧本（分阶段、定规则、做简化）。通过这种“宏观指挥 + 经验传承 + 安全监控”的组合，它成功解决了“地图爆炸”的难题，让编译器能自动找到既快又省的最佳方案，而且不需要人类专家手把手教。

这就好比以前我们要找宝藏，只能靠人拿着地图一个个试，容易迷路累死；现在有了 EggMind，AI 能自动画出最合理的寻宝路线图，既快又稳，还能把经验传给下一个人。

1. 研究背景与问题 (Problem)

等式饱和 (Equality Saturation, EqSat) 是一种强大的程序优化范式，它利用 e-graph（等式图）紧凑地表示大量等价程序，并推迟到最终提取阶段才根据成本模型选择最优解。与传统的破坏性重写不同，EqSat 能发现固定顺序重写所遗漏的优化。

然而，EqSat 在实际应用中面临 e-graph 爆炸 (e-graph explosion) 的主要瓶颈：

资源不可行性：随着重写规则的增加，搜索空间呈指数级增长，导致内存和运行时间迅速耗尽。
策略设计的局限性：为了控制爆炸，需要设计特定的 策略 (Strategy) 来组织搜索（如规则集划分、调度顺序、简化时机）。目前这些策略主要依赖人工设计，难以自动化，且无法适应大规模规则集。
现有自动化工具的不足：虽然已有工具（如 Enumo, Ruler）能自动生成重写规则，但这反而扩大了搜索空间，加剧了爆炸问题。
直接利用 LLM 的困难：
- 缺乏抽象：直接让 LLM 生成底层后端代码（如 egglog 脚本）缺乏可复用的策略抽象，难以表达高层意图。
- 反馈缺失：LLM 难以从最终成本或巨大的 e-graph 状态中获得有效的反馈，无法理解策略为何成功或失败。
- 不可控性：缺乏引导时，LLM 容易生成导致灾难性内存增长或不稳定搜索的策略。

2. 核心方法：EggMind (Methodology)

为了解决上述问题，作者提出了 EggMind，一个基于 LLM 引导的端到端框架，用于合成可复用的 EqSat 策略。其核心由三个支柱组成：

2.1 领域特定语言：EqSatL

EggMind 引入了一种名为 EqSatL 的领域特定语言 (DSL)，将 EqSat 的控制逻辑显式化为可检查、可复用的工件（Artifact），将其与底层实现解耦。EqSatL 包含三个控制层面：

规则集划分 (Ruleset Partitioning)：将重写规则按语义标签（Tags）分组，决定哪些规则可以交互。
调度构建 (Schedule Construction)：定义规则交互的时间顺序，包括阶段（Phase）、序列（Sequence）和重复（Repeat）结构。
简化控制 (Simplification Control)：显式控制何时以及如何收缩 e-graph（例如通过启发式或 LLM 提示的简化），防止状态无限膨胀。

2.2 代理工作流 (Agentic Workflow)

EggMind 采用离线合成工作流，通过迭代优化策略：

角色分工：包含 Strategist（策略制定者）、Generator（生成器）、Evaluator（评估器）、Partitioner（划分器）和 Simplifier（简化器）。
状态管理：维护一个策略仓库（当前最佳、有潜力的、待评估的）和一个重写模式缓存。
迭代过程：Strategist 根据上下文提出策略更新 -> 调用工具执行评估 -> 更新工作流记忆。

2.3 关键技术组件

基于证明的重写模式缓存 (Proof-Derived Rewrite Motif Caching)：
- 不直接缓存庞大的 e-graph 状态或原始证明，而是从成功的 EqSat 运行中提取 重写模式 (Rewrite Motifs)。
- 模式是语义标签的短链（如 vec-lifting -> vec-mac-opt），抽象了具体的规则细节。
- 这些模式作为“证据”被缓存，用于指导 LLM 在后续迭代中生成更有效的策略结构。
可解性引导 (Tractability Guidance)：
- 基于依赖的规则集风险模型：构建规则依赖图，评估不同阶段划分的风险，避免导致 e-graph 不稳定的交互（如防止后向激活或同阶段纠缠）。
- LLM 引导的简化提示：针对每个阶段，LLM 提供结构偏好（哪些模式应优先保留，哪些应被剪枝），作为简化过程中的软性约束，平衡收缩与保留有用中间态。

3. 主要贡献 (Key Contributions)

EggMind 框架：首个开源的、基于 LLM 引导的框架，用于合成可复用的 EqSat 策略，实现了从“一次性脚本”到“跨案例复用策略”的转变。
EqSatL 语言：提出了一种 DSL，将策略控制显式化为可验证的工件，支持规则集组织、调度构建和简化控制。
新型代理工作流：结合了证明衍生的模式缓存和可解性引导，解决了 LLM 在 EqSat 搜索中缺乏有效反馈和易导致不稳定的问题，实现了稳定且高质量的策略搜索。

4. 实验结果 (Results)

作者在三个主要领域进行了评估：

4.1 向量化编译器基准 (Vectorization Benchmarks)

对比对象：全量 EqSat (Full EqSat) 和专家调优的 Isaria 策略。
性能提升：
- 质量：相比 Full EqSat，最终成本降低了 45.1%；相比 Isaria，降低了 20.6%。
- 资源：峰值内存使用量降低了 69.1%（最大降幅达 95%）。
- 速度：相比 Isaria，运行时间几何平均加速 2.21 倍；相比 Full EqSat，在大型案例上加速达 6.85 倍。
消融实验：证明了模式缓存 (Motif Caching) 和重复结构 (Repeat) 对性能至关重要。移除模式缓存会导致成本恶化 78.3%。

4.2 基于 XLA 的张量编译器 (XLA-based Tensor Compiler)

在代数简化任务中，EggMind 在 17 个案例中的 13 个实现了比全量 EqSat 更低的最终成本。
运行时间几何平均加速 11.89 倍，证明了该方法在不同优化领域的可迁移性。

4.3 逻辑综合案例研究 (Logic Synthesis Case Study)

在结合 Enumo 生成的增强重写空间（包含更复杂的逻辑与物理单元映射规则）下，EggMind 展现了关键优势。
结果：相比原始规则集（ $i=4$ ），成本降低 33.76%；相比增强规则集的全量搜索（ $i=5$ ，通常导致内存溢出或超时），EggMind 在保持稳定运行时间的同时，成本降低了 8.75%，峰值内存降低 51.94%。
意义：证明了策略合成对于利用“增强重写空间”的必要性——没有好的策略，更大的搜索空间只会导致爆炸。

5. 意义与影响 (Significance)

自动化编译器构建的新范式：EggMind 将 EqSat 策略设计从手工调优转变为自动化的离线合成过程，使得基于 e-graph 的编译器构建更加可行。
解决“规则爆炸”悖论：证明了虽然自动规则合成扩大了搜索空间，但通过 LLM 引导的策略合成，可以有效地管理这种复杂性，将“不可行”的搜索空间转化为“可解”的优化过程。
可复用性与效率：通过 EqSatL 和模式缓存，策略知识可以在同一领域的不同案例间复用，显著降低了后续优化的边际成本。
LLM 在系统优化中的新应用：展示了 LLM 不仅仅是生成代码，通过适当的抽象（DSL）和反馈机制（模式缓存、可解性引导），LLM 可以成为复杂系统控制策略的高效搜索者。

总结：EggMind 通过引入显式的策略抽象语言 (EqSatL) 和智能的 LLM 代理工作流，成功解决了等式饱和中 e-graph 爆炸的难题，实现了在资源受限下的高质量程序优化，为下一代自动化编译器优化奠定了重要基础。