关键词：等式饱和、e-graph、编译器、MLIR、持久化、优化
通过将 e-graph 直接嵌入 MLIR，研究人员让等式饱和贯穿整个编译流程，无需反复翻译、不丢失等价信息，并成功复现了 Herbie 浮点精度优化工具。
现代编译器通常由一系列独立的优化遍（pass）组成，每个遍在中间表示（IR）上执行特定的转换，例如常量折叠、死代码消除、循环不变式外提等。这种模块化设计虽然便于维护和扩展，却引入了一个经典难题——阶段顺序问题（phase‑ordering problem）。某个优化 Pass 可能破坏另一个优化 Pass 的应用机会，而编译器的编写者必须手动决定 Pass 的顺序，这往往依赖经验或启发式，无法保证得到最优结果。
等式饱和（equality saturation） 提供了一种全新的思路。它不再“破坏性”地替换表达式，而是将所有等价表达式都保存在一个称为 e‑graph 的数据结构中，并持续应用重写规则向其中添加新表达式。
当 e‑graph 达到饱和（不再能添加新表达式）或达到某个资源限制后，再根据代价模型从中提取出最优表示。这种方法完全避免了阶段顺序问题，因为所有可能的等价形式都被保留，最终选择可以全局考虑。
等式饱和已经在许多优化挑战中取得成功，例如浮点精度提升（Herbie）、高 Level 综合（SEER）、指令选择（Cranelift）等。然而，现有的集成方式存在明显短板：
- 要么将程序翻译到外部等式饱和库如 egg、egglog，优化后再翻译回来，导致翻译开销和等价信息丢失；
- 要么为特定应用定制 e‑graph 实现，难以复用。
本文介绍的工作发表于 CGO 2025，其提出了一种激进而优雅的解决方案：将 e‑graph 直接嵌入编译器的 IR 中，使其成为编译器的一部分，从而在整个编译流程中持久化地保留等价信息。
作者基于 MLIR 框架和其 Python 实现 xDSL，设计了新的 eqsat dialect，并改造了 MLIR 的模式匹配基础设施，使开发者可以用完全相同的方式编写重写规则，既可用于传统破坏性重写，也可用于等式饱和。通过复现 Herbie 的核心功能，作者证明了该方法的灵活性和潜力。
一、等式饱和基础：e-graph 与 e-matching
在深入技术细节之前，我们先回顾 e‑graph 的核心概念。
一个 e‑graph 由若干 e‑class 组成，每个 e‑class 包含多个 e‑node ，这些 e‑node 在语义上是等价的。e‑node 类似于抽象语法树中的节点，但它的子节点指向的是 e‑class 而非具体节点。
例如，表达式 a * 2 可以表示为一个 e‑class（记为 C1）包含一个乘法 e‑node，其左子指向变量 a 所在的 e‑class，右子指向常数 2 所在的 e‑class。当应用重写规则 x*2 → x<<1 时，并不会删除原来的乘法节点，而是向 C1 中添加一个新的左移 e‑node，从而将两种表示都保留在同一个 e‑class 中。
为了在这种共享结构上进行模式匹配，需要 e‑matching 算法。与普通模式匹配不同，e‑matching 必须考虑 e‑class 内的多个候选，并利用变量间的相关性来剪枝。高效的 e‑matching 通常将多个模式编译成一个自动机，并共享公共子模式以减少重复工作。
二、现有方法的不足：外部库与定制实现的局限
图 1 清晰地勾勒了当前三种集成方式的对比：
图 1 | 现有研究仅针对单一抽象层级开展工作，要么是（b）协调外部的等式饱和库，要么是（c）为特定使用场景原生支持 e-graph。相比之下，我们的持久化等式饱和方法（d）能够复用现有的编译器转换逻辑，且能在不同抽象层级间保留推导出的等式信息。整体来说，能看出编译器中等式饱和的这三类实现方式中，传统方案或依赖外部库、或仅适配专属场景，均存在抽象层级单一、等式信息易丢失的问题。而本文的持久化方案将 e-graph 嵌入编译器 IR，实现了跨层级的等式保留，还能直接复用编译器现有转换流程，从根本上突破了传统方案的应用局限。
- 图 1a ：传统编译器流程，优化遍顺序固定。
- 图 1b ：调用外部等式饱和库。程序被翻译到库的 IR，饱和后再翻译回编译器 IR。缺点：翻译开销大，且一旦提取，所有其他等价形式立即丢失，后续遍无法再利用这些信息。
- 图 1c ：深度定制，如 Cranelift 的 acyclic e‑graph。它直接使用 e‑graph 表示 IR，但为了编译速度牺牲了通用性如不支持循环，且无法与其他编译器基础设施（如已有的分析遍）协同工作。
本文提出的方法（图 1d）则是在编译器内部以一等公民的身份实现 e‑graph，使其能够贯穿整个编译流程，并重用现有的分析、转换和模式匹配机制。
三、核心创新：将 e-graph 嵌入编译器 IR
3.1 MLIR 基础：可扩展的 SSA 中间表示
本文的实现基于 MLIR（Multi-Level Intermediate Representation）框架。MLIR 将用户可扩展性作为核心设计，通过 dialect 组织不同抽象层次的操作，并利用 SSA（静态单赋值）形式简化分析与转换。
每个操作可以包含 region，而 region 又由基本块组成，支持图区域（graph region）以允许循环定义，如图 2 所示。这种灵活的结构为在编译器内部嵌入 e-graph 提供了理想的基础。
图 2 | MLIR 的模块化结构。其以 SSA 为基础，融合了方言、操作、属性等设计，为核心的可扩展性提供了支撑，使得嵌入 e-graph 成为可能。
基于 MLIR 的 Python 实现 xDSL，我们引入了一个新的 dialect——eqsat，它通过四个核心操作在 SSA 形式的 IR 中完整表示 e-graph。
3.2 eqsat dialect 设计
作者在 MLIR 中引入了一个新的 dialect——eqsat，它包含四个核心操作，足以在 SSA 形式的 IR 中完整表示一个 e‑graph：
eqsat.eclass： 接受一个或多个值（即 e‑node 的结果），产生一个代表该 e‑class 的 SSA 值。它的每个操作数都对应一个 e‑node。eqsat.const_eclass： 专为常量设计的变体，它直接持有一个常量属性，可以触发 MLIR 的常量折叠机制。eqsat.egraph： 一个图区域（graph region）操作，其体内包含所有与该 e‑graph 相关的eclass操作。图区域允许循环定义，这对于表示由重写引入的循环（如a + 0 → a）至关重要。eqsat.yield： 作为egraph的终结符，将指定的 e‑class 结果暴露给外部。
将普通函数转换为 e‑graph 的过程非常直接：一个编译器遍会遍历函数，将每个操作的结果用 eqsat.eclass 包装，并将所有对原始值的引用替换为对对应 e‑class 结果的引用。Listing 3 展示了这一转换。
Listing 3 | 通过插入 eqsat.eclass 等操作，将普通 IR 转换为 e‑graph 表示。右侧是等价的 e‑graph 可视化，虚线框代表 e‑class。
当应用重写 x*2 → x<<1 时，新的左移操作及其常量被创建，但原有的乘法并未被删除，而是将其结果也添加为同一个 e‑class 的成员（Listing 4）。这正是等式饱和的精髓——保留所有等价形式。
Listing 4 | 应用重写后，eqsat.eclass 现在包含两个成员：乘法和左移，实现了 e‑class 的扩展。
3.3 从传统重写到 e-matching：改造 PDL
MLIR 提供了一个强大的声明式重写框架 PDL（Pattern Definition Language），开发者可以用 PDL 编写模式，然后编译成高效的匹配器（pdl_interp dialect）。
传统上，PDL 匹配器是破坏性的——匹配到的模式会被替换。作者保留了 PDL 的语法，但重新实现了 pdl_interp 的解释器，使其能够感知 eqsat.eclass 的间接层，并支持回溯。
图 3 | 传统 MLIR 重写与 eqsat 中操作关联的核心差异。传统框架中操作的结果与定义源一一对应，而 eqsat 因引入 e-class，一个值对应多个等价操作，打破了原有逆操作关系。
关键挑战在于 pdl_interp.get_defining_op 和 pdl_interp.get_result 这类操作。
* 在普通 IR 中，它们互为逆操作；
* 但在 e‑graph 中，一个值可能由多个 e‑node 定义，get_defining_op 需要返回所有可能的定义操作。
解释器通过维护回溯点来处理：当一条路径失败时，回退到最近的选择点并尝试 e‑class 中的下一个 e‑node。
此外，pdl_interp.create_operation 不再盲目创建新操作，而是先检查是否已存在相同的操作（哈希一致性，hashconsing），以避免重复。pdl_interp.replace 则不再删除原操作，而是将替换值加入对应的 e‑class。
3.4 优化 e-matching：避免指数爆炸
将多个 PDL 模式简单地组合成一个匹配器，直接用于 e‑matching 会导致性能灾难。原因在于默认的 PDL 到 pdl_interp 的降级策略优先进行结构检查，再进行语义检查。
例如，对于模式 f(a, g(a))，匹配器会先获取外层操作的两个操作数，然后对内层操作获取其操作数，最后才检查操作名称是否匹配。在 e‑matching 中，每次获取操作数都可能引入多个候选（因为操作数指向的是 e‑class），导致组合爆炸。
作者提出了两种优化：
- 优化1：尽早进行语义检查：将操作名称等语义检查尽可能提前，一旦发现不匹配立即回溯，避免不必要的结构探索，见图 4。
图 4 | 启用即时语义检查前后的 pdl_interp 匹配代码对比。通过尽早执行语义检查，能够避免回溯点数量的指数级激增。
- 优化2：引入
eqsat_pdl_interp.choose操作：当匹配器从尝试一个模式切换到另一个模式时，原来的回溯点会被保留，造成多个不相关模式的回溯点混杂。choose操作显式地标记分支点，使得回溯时可以直接跳转到当前模式的选择点，隔离不同模式的回溯栈，见图 5。
图 5 | 通过引入 choose 操作分离不同模式的回溯点，避免了跨模式回溯导致的指数级增长问题。
实验表明，这些优化在小规模测试中带来了 2.34 倍的几何平均加速，在后续的 Herbie 案例中达到 2.57 倍。
3.5 利用现有分析实现 e-class 分析
e‑class 分析是 egg 引入的重要扩展，它允许在 e‑class 级别上传播抽象信息（如区间、符号性等）。
MLIR 本身拥有强大的稀疏数据流分析框架，作者巧妙地将其复用：只需为 eqsat.eclass 定义 visit 方法为其操作数的格值的 meet，数据流求解器就能自动在 e‑graph 上传播信息。这意味着任何现有的 MLIR 分析（如常量传播、区间分析）无需修改即可成为 e‑class 分析。
对于常量折叠，作者采用 eqsat.const_eclass 配合 MLIR 的 fold 机制，在插入新操作时立即尝试折叠，这比 egg 的懒加载 modify 更积极。
四、案例研究：浮点精度优化（Herbie 复现）
为了验证 eqsat 的实用性，作者选择复现著名的浮点表达式精度优化工具 Herbie 的核心部分。Herbie 使用等式饱和搜索与原始表达式在实数域等价、但在浮点运算中误差更小的表达式。
4.1 Herbie 工作原理回顾
给定一个浮点表达式，Herbie 会：
1. 生成随机样本点作为输入。
2. 应用一系列数学重写规则（如分配律、结合律、恒等式），在 e‑graph 中构建大量候选表达式。
3. 对每个候选，用高精度计算其在样本点上的真实值，并与浮点计算结果比较，得到误差。
4. 选择误差最小的表达式输出。
Herbie 最初使用 Racket 实现，后来借助 egg 库大幅加速。图 6 展示了本文实现的流程。
图 6 | 结合现有的和全新的编译器转换方法与分析手段，我们复现了 Herbie 工具主循环的单次迭代过程。该图梳理了基于 eqsat 方言复现 Herbie 核心流程的八大步骤，从 FPCore 解析到 MLIR 映射，再到 e-class 插入、PDL 规则转换、等式饱和执行，最终完成精度分析与代码提取。整个流程全程基于编译器原生 IR 与基础设施，无需外部工具转换，验证了 eqsat 方言在实际等式饱和应用中的可行性与兼容性。
4.2 实现流程
- 解析与转换：将 Herbie 的输入语言 FP Core 翻译成 MLIR 的
arith和math方言操作。 - 构建 e‑graph：通过插入
eqsat.eclass将函数体包装成初始 e‑graph。 - 样本生成：运行区间分析确定每个变量的有效范围，并在该范围内均匀采样，避免除以零等错误。
- 重写规则：将 Herbie 的 egglog 规则集（包含前置条件，如要求某个子表达式非负）手工转换为 PDL 模式。这些前置条件通过 e‑class 分析（区间、非零等）来验证。
- 等式饱和：运行改造后的 PDL 匹配器，直到 e‑节点数达到 4000 或无法再添加新节点。
- 计算真实值：对于每个 e‑class，只需计算其中一个 e‑节点的高精度值（使用 MPFR 库，1024 位精度），因为所有节点在实数域等价。
- 误差计算：对每个 e‑节点，用浮点计算其在样本点上的值，与高精度真实值比较，取对数平均误差作为代价。
- 提取：运行贪心选择，将误差作为代价，最终提取出最优表达式并转换回 FP Core 格式，以便与 Herbie 的结果公平比较。
4.3 准确性对比
在 FPBench 的 31 个 Herbie 基准测试上，本文实现与 Herbie 的结果进行了准确性对比，见图 7。
图 7 | 我们的方法在 FPBench 大多数基准测试中能够达到与 Herbie 工具相当的精度。精度上的差异源于重写规则的细微不同，以及因早期超时导致的潜在规则遗漏。对于三个表达式精度结果一致（误差在 1% 以内）的基准测试，我们在视觉上做了淡化处理。总得来说，eqsat 与 Herbie 二者整体精度持平，少数案例的差异源于规则集细微不同、部分函数未实现及 e-nodes 数量限制。这一结果验证了 eqsat 方言在浮点表达式精度优化场景下的实际效果，证明其能复现专业等式饱和工具的核心能力，具备实际应用价值。
总体而言，两者在大多数基准上精度相当。少数差异归因于：
* 本文未实现某些函数，如 hypot；
* 规则集略有不同。本文使用 egglog 规则，Herbie 有自己的一套；
* Herbie 可能使用了一些无法证明其安全性的规则，而本文依赖 e‑class 分析确保规则应用前提成立；
* 非确定性：由于 e‑节点数限制，不同规则应用顺序可能导致不同的最终候选。
有趣的是，在某些基准上本文实现反而优于 Herbie，同样可能是因为规则集差异和探索顺序的偶然性。
4.4 性能分析与讨论
性能方面，本文实现比 Herbie 慢约 400 倍（几何平均，见图 8）。
图 8 | 我们的实现版本比 Herbie 工具慢得多，这主要是因为 Herbie 和 egg 库的实现中包含诸多性能优化手段，而这些优化目前尚未应用到我们基于 xDSL 的实现中。该图对比了 eqsat 与 Herbie 在基准测试中的运行耗时，前者存在显著的速度差距，核心原因包括 xDSL 的 Python 实现效率低于 Herbie 的 Rust+Racket 组合、MPFR 高精度计算未做动态精度优化、还需执行重型区间分析。但这一差距并非算法层面的问题，论文指出将实现迁移到 C++ 的 MLIR 后，性能可得到大幅提升。
主要原因：
* 语言层面：xDSL 基于 Python，而 Herbie 后端是 Rust（egg）和 Racket，两者都远快于 Python。
* 高精度计算开销：本文使用固定 1024 位 MPFR 计算所有样本，而 Herbie 动态调整精度以加速。
* 分析开销：egglog 规则要求复杂的区间分析，Herbie 默认避免在饱和过程中运行分析。
* 优化程度：egg 经过多年打磨，而 eqsat 尚处于原型阶段。
尽管如此，仅考虑等式饱和阶段（不包括高精度计算），组合匹配器相比逐个匹配模式仍有 2.57 倍的加速见图 9，验证了匹配优化的有效性。
图 9 | 将所有的 PDL 模式整合为单个 pdl_interp 执行等式饱和，与逐个匹配每个模式的方式相比，能带来显著的速度提升。该比较验证了本文 pdl_interp 优化的实际效果，将所有重写模式整合为单个匹配例程，相比逐个匹配模式，实现了 2.57 倍的几何平均速度提升。这是因为整合后的匹配例程能复用公共子模式的匹配逻辑，避免重复计算，结合前文的即时语义检查、选择操作优化，让 e-matching 的效率得到了多重提升。
五、相关工作
等式饱和与编译器的结合已有诸多探索。论文将其分为三类：
- 外部库集成：如 egg 和 egglog，被用于 Herbie、SEER、ROVER 等。优点是性能好，缺点是需要来回翻译且无法保留等价信息。
- 定制实现：如 Cranelift 的 acyclic e‑graph，将 e‑graph 作为 IR 的核心，但为速度牺牲了通用性（不支持循环），且难以复用编译器其他部分。
- 混合方法：DialEgg 将 MLIR 与 egglog 连接，降低了翻译开销，但仍需在每次调用后丢弃 e‑graph，无法跨遍持久化。
本文工作属于“深度嵌入”，将 e‑graph 作为编译器的一等公民，既保留了通用性，又实现了持久化，并能重用现有分析、转换和模式匹配。
总结与展望
本文提出了一种将等式饱和持久化嵌入编译器的新方法。通过在 MLIR 中引入 eqsat dialect，并改造 PDL 匹配器，作者展示了如何让 e‑graph 贯穿整个编译流程，避免翻译开销和信息丢失，同时能够利用编译器现有的丰富基础设施。案例研究成功复现了 Herbie 的核心功能，证明了方法的可行性。
未来工作包括：
- 将实现移植到 C++ 版本的 MLIR 中，以大幅提升性能。
- 结合引导式等式饱和（guided equality saturation），利用编译器中的具体目标（如循环深度、指令数）引导探索方向。
- 支持混合代价模型，例如在性能与浮点精度之间权衡。
- 扩展对控制流的支持，目前 PDL 尚不能匹配带有区域的操作（如循环），未来有望实现对控制流本身的等式饱和。
等式饱和曾被质疑其高昂的开销，但随着硬件性能提升和算法优化，它正逐渐从学术玩具变为实用的编译器技术。这项工作为编译器社区提供了一种低门槛、高灵活性的集成方案，有望加速等式饱和在工业编译器中的普及。
关注“鲸栖”小程序，掌握最新AI资讯
本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：https://www.itsolotime.com/archives/22128