egg教程(一):e-graphs and equality saturation的概念

2023-10-17 11:21:36 浏览数 (3)

e-graph是一种支持equality saturation优化技术的数据结构。

e-graph(Gregory Nelson’s PhD Thesis, 1980)和equality saturation(Tate et. al., 2009)都是在此之前发明的。egg目的是让他们更加简单,易拓展和高效。

本教程将从高层次探讨这些概念。更多细节可参见论文。

为什么要这样做?

在构建几乎所有与编程语言相关的工具时,您通常会使用您语言中的term(我们将交替使用term、expression和program等词)。您可能会尝试:

1. 将term转换为 "更好的 "等价term(optimize);

2. 根据规范从头构建一个term(synthesis);

3. 证明两个term是等价的(verification)。

Term rewriting 是一种常用的处理技术,可以实现上述任何目标。你可以从一个term t 和一组rewriting开始,每个重写的形式都是 l → r。对于每个重写,我们都会尝试将模式 l 与term t 进行匹配,在某个subterm上生成一个替换 σ,然后将该替换应用到右侧的模式 r 上,并替换匹配的subterm。

例如,考虑term 42 × (7 7) 和重写 x x → 2 × x:

1. 左侧将匹配子项 7 7,生成代词 σ = {x:7}。

2. 应用该替换后,r[σ] = 2 × 7(r[σ] 表示将替换词 σ 应用于 r)。

3. 用 r[σ] = 2 × 7 代替匹配子项 7 7,得到结果:42 × (2 × 7)。

Term rewriting(以及其他涉及term操作的编程语言技术)的一个问题是选择问题。重写是破坏性的,也就是说,一旦进行了重写,初始术语就会消失。如果你有多种重写方法可供选择,那么选择其中一种就像选择岔路口一样:返回并选择另一种方法是昂贵的,或者是不可能的,因此你必须坚持你所选择的道路。此外,在正确的时间选择正确的改写一般来说是非常困难的;有时,在局部看来不错的选择可能在整体上是错误的。

例如,C 语言优化器希望将 a * 2 重写为更便宜的 a << 1,但 (a * 2) / 2 又如何呢?贪婪的局部方法会 "转错弯",结果是 (a << 1) / 2,它比输入表达式更好,但却掩盖了我们本可以取消乘除的事实。

有一些技术可以缓解这种情况(回溯和数值编号),但从根本上说,绕过它的唯一方法就是同时进行所有选择。但重写是破坏性的,因此对一个词条进行 n 次重写会产生 n 个term,然后对这 n 个term中的每个term进行 n 次重写又会产生 n2 个term,以此类推。如果你想探索这棵树的 m "层 "深处,你就必须存储 nm 个term,而这是非常昂贵的。

e-graph可以解决这个表示问题。当重写一堆term时,结果往往在结构上是相似的。e-graph可以非常紧凑地存储大量相似表达式的集合。使用e-graph,你可以同时应用许多重写,而不会造成空间的成倍增长。Equality saturation是一种为优化程序而进行此类重写的技术。

什么是e-graph?

e-graph 是一种数据结构,用于维护表达式的等价关系(实际上是全等关系,见下一节)。e-graph是一组等价类(e-classes),每个等价类都包含等价的 e-node。e-node是带有子节点的运算符,但子节点不是其他运算符或值,而是 e-class。在 egg 中,这些由 EGraph、EClass 和 Language(e-nodes)类型表示。

即使是小的e-graph也能表示大量的表达式,其数量与e-nodes的数量成指数关系。这种紧凑性使得e-graph成为一种引人注目的数据结构。我们可以将表示(或包含)一个term的含义定义如下:

1. 如果一个e-graph的任何一个e-class都代表一个term,那么这个e-graph就代表term。

2. 如果一个 e-class 的任何一个 e-node 都代表一个term,那么这个 e -class就代表term。

3. 如果 e-class ni 代表term ti,则 e -node f(n1, n2, ...) 代表term f(t1,t2,...)。

下面是一些e-graph。我们把 e-class 描绘成围绕着等价 e-node 的虚线框。请注意,边是从 e-node 到 e-class(而不是到其他 e-node),因为 e-node 的子节点是 e-class。先不用担心它们是如何构成的,我们很快就会重写。

可以通过称为 e-matching(等价匹配)的过程查询e-graph中的模式,该过程会在e-graph中搜索代表与给定模式匹配的term的 e-class。egg 的 e-matching(受 De Moura 和 Bjørner 的启发)由搜索器和模式 API 完成。修改e-graph的核心操作有两个:添加(add)和合并(union),前者用于向e-graph中添加e-node,后者用于合并两个e-class。

这些操作可以结合起来对e-graph进行重写。重要的是,这种重写过程只是加法,这意味着电子图永远不会忘记term的旧版本。此外,重写一次可以添加整类term。要在e-graph上执行 l → r 重写,首先要在电子图中搜索 l,得到 (c, σ) 对,其中 c 是匹配的 e-class,σ 是替换。然后,将术语 r[σ] 添加到 e 图中,并与匹配的 e-class c 结合。

下面我们就以上图中的四幅e-graph为例,来说明这一切。

1. 初始 e-graph 表示 term (a × 2) / 2 。由于每个 e-class 只有一个 e-node ,因此该 e-graph 基本上是一棵共享的抽象语法树(2 不重复)。

2. 应用重写 x × 2 → x << 1 已记录了 a × 2 = a << 1 的事实,而没有忘记 a × 2。请注意新添加的 a << 1 是如何指向现有的 "a " e-node的,而"<<" e-node已被联合到与模式 x × 2 匹配的"×"e-node相同的e-class中。

3. 应用重写 (x × y) / z → x × (y / z) 实现了除法与乘法的关联。这个重写对于发现我们正在寻找的 2 的抵消是至关重要的,尽管我们之前应用了 "错误的 "重写,但它仍然有效。

4. 应用 x / x → 1 和 x × 1 → x 重写并不会增加任何新的 e-node,因为所有的 e-node 都已经存在于 e-graph 中。结果只是联合了 e-class,也就是说,尽管现在 e-graph 代表了更多的项,但应用这些重写后,e-graph 实际上变小了。事实上,观察右上方"×" e-node 的左侧子节点就是它自己;这个循环意味着这个 e-class 代表了无限(!)的项集 a、a ×1、a × 1 ×1,以此类推。

不变式和重建

e-graph有两个修改e-graph的核心操作:add(向e-graph中添加e-node)和 union(合并两个e-class)。这些操作维护两个关键(相关)不变式:

1. 一致性

e-graphj不仅维护表达式的等价关系,还维护同构关系。同构关系的基本含义是,如果 x 等价于 y,那么 f(x) 必须等价于 f(y)。因此,当用户调用 union 时,除了给定的两个 e-class 外,其他许多 e-class 可能需要合并以保持一致性。

例如,假设术语 a x 和 a y 分别在 e-class 1 和 2 中表示。在以后的某个时刻,x 和 y 变得等价(也许用户在包含它们的 e-class 中调用了 union)。E-class 1 和 2 必须合并,因为现在两个 " "运算符有了等价参数,使它们等价。

2. e-node的唯一性

在同一 e 类或不同 e 类中,不存在两个不同的 e 节点,它们具有相同的运算符和等同的子节点。这在一定程度上是通过 add 执行的散列整理以及 union 和 rebuild 执行的重复数据删除来实现的。

egg 采用延迟的方式来维护这些不变性。具体来说,调用 union(或应用调用 union 的重写)的效果可能不会立即反映出来。要恢复e-graph不变式并使这些效果可见,用户必须调用重建方法。

egg在这里的选择可以实现更高的性能。维持同构关系会使核心e-graph数据结构复杂化,并且在每次联合时都需要对e-graph进行昂贵的遍历。为了提高性能,egg 选择放宽这些不变式,只有在调用重建时才恢复不变式。关于 egg 的论文更详细地说明了为什么这种设计选择会让速度更快。有关 rebuild 的作用以及何时必须调用 rebuild 的更多信息,请参阅 rebuild 文档。还需注意的是,运行程序会在重写迭代之间调用 rebuild。

Equality Saturation

Equality saturation (Tate et. al., 2009) 是一种程序优化技术,egg通过类似下面伪代码的runner api中实现了equality saturation的变体。

代码语言:javascript复制
fn equality_saturation(expr: Expression, rewrites: Vec<Rewrite>) -> Expression {
    let mut egraph = make_initial_egraph(expr);

    while !egraph.is_saturated_or_timeout() {
        let mut matches = vec![];

        // read-only phase, invariants are preserved
        for rw in rewrites {
            for (subst, eclass) in egraph.search(rw.lhs) {
                matches.push((rw, subst, eclass));
            }
        }

        // write-only phase, temporarily break invariants
        for (rw, subst, eclass) in matches {
            eclass2 = egraph.add(rw.rhs.subst(subst));
            egraph.union(eclass, eclass2);
        }

        // restore the invariants once per iteration
        egraph.rebuild();
    }

    return egraph.extract_best();
}

上文已经介绍了大部分内容,但我们需要定义两个新term:

1. 当e-graph检测到改写不再增加新信息时,就会出现饱和。考虑一下交换重写 x y → y x。在应用了一次之后,由于e-graph没有忘记最初的 x y 项,所以第二次重写不会再增加新信息。如果所有的重写都处于这种状态,我们就说e-graph已经饱和,这意味着e-graph包含了从给定的重写中推导出的所有可能的等价关系。

2. Extraction是一种从e-graph中挑选单个代表term的过程,根据某些成本函数,该过程是最优的。egg的Extractors提供了这个功能。

总而言之,equality saturation会探索程序的所有可能变体,这些变体可以从一组重写中衍生出来,然后提取出最佳变体。这就解决了上文所述的选择问题,因为equality saturation基本上每次迭代都会应用每个重写,利用e-graph避免指数爆炸。

这听起来有点好得不像真的,虽然它肯定有其注意事项,但 egg 解决了这种方法的许多缺点,使其在实际问题中的应用变得可行。我们已经讨论了重建如何使 egg 的e-graph变得快速,稍后的教程将讨论分析如何使这种方法变得灵活,并能处理更多语法重写问题。

0 人点赞