egglog入門

はじめに

egglogはDatalog[1]っぽいプログラミング言語で、e-graphによって実行されます。
汎用的なプログラミング言語ではありませんが^[2]、数式とかプログラムの最適化が得意でherbie[3]でも使われたりしています。

以下は有向グラフの2点間の最短距離を求めるegglogのサンプルプログラムです。
ちなみにこれだけなら別にめちゃめちゃ最適なアルゴリズムで実行されるというわけでもなくそんなに有用ではないですが、雰囲気はわかると思います。

webデモ

path.egg

; ;から行末まではコメント
; egglogの文法はLispっぽいが見た目がそうなだけ

; 二点間の距離を定義する
(function edge (i64 i64) i64)
; 二点間の最短距離
(function path (i64 i64) i64 :merge (min old new))

(rule ((= (edge x y) d))
    ((set (path x y) d)))

(rule ((= (path x y) d1) (= (edge y z) d2))
    ((set (path x z) (+ d1 d2))))

; 辺のデータを入れていく
(set (edge 1 2) 1)
(set (edge 2 3) 2)
(set (edge 3 4) 3)
; EqSat実行
(run 100) ; 最大100回までイテレーションする。正直100という数字は適当だが文法上必要なので入れている。別に1000とかでも良い
; 1-4間の最短距離を出力
(extract (path 1 4)) ; 6

; 1-4間にもっと短い経路があることを入力
(set (edge 1 4) 5)

; EqSat実行
(run 100)

; 1-4間の最短距離を出力
(extract (path 1 4)) ; 5

Equality Saturationとe-graph

egglog自体e-graphをうまく使うためのインターフェースといった感じがあり、Equality Saturationとe-graphの理解は外せないのでここで説明します。

Equality Saturationが解く問題

手元にあるS式^[4]があり、それをあらかじめ決めたルールに従って書き換えていってより良い形^[5]にするという問題を考えます。このとき書き換える前と後はなにかの基準で同じ意味を持つようにルールを作ります。

例えば、数式を表すS式(/ (* a 2) 2)を以下のルールに従って書き換えていきます。
ご存知の通り$(a \times 2) / 2 = a$^[6]なので、最終的にはaに書き換えられればOKです。

ルール

x, y, zは任意のS式を表します。

1. (* x 2) を (<< x 1)に書き換える。<<は左シフトの意味
2. (* x 1) を xに書き換える。
3. (/ (* x y) z) を (* x (/ y z))に書き換える。
4. xが0でない場合、(/ x x) を 1に書き換える^[7]。

上のルールを適切な順番で適用していくと、(/ (* a 2) 2)をaに書き換えられますが、順番が重要です。

• ルール1を最初に選んだ場合 (/ (* a 2) 2) ➡️ (/ (<< a 1) 2) ...これ以上適用できるルールがない
• ルール3を最初に選んだ場合 (/ (* a 2) 2) ➡️ (* a (/ 2 2)) ルール4を適用 ➡️ (* a 1) ルール2を適用 ➡️ a ...これはOK🎉

つまり適当にルールを適用していく(そして元の式を消す^[8])だけでは最適な結果になるとは限らないのです。

これは最適化問題の頻出系なのでとりあえずビームサーチすればどうにかなりそうに見えますが、この問題に対してはもっと良いアルゴリズム(Equality Saturation)とそれを可能にするデータ構造(e-graph)があります。

Equality Saturation

Equality Saturationはこの問題を次のように解きます。

1. 手元のS式に対して適用できるルールを全部適用して適用前、適用後の情報を全部覚えておく
2. 1を繰り返す。そのうちどのルールをどう適用しても新しい情報を得られなくなる(Saturationする😎)^[9]ので、そしたら次に行く
3. 手元にあるS式の中から(なにかの基準で)一番良いものを選んでそれを出力する

https://blog.sigplan.org/2021/04/06/equality-saturation-with-egg/ から引用した図

実質すべての可能性を探索しているので必ず最適な解が得られます💪
ただ、過程で出てきたすべての情報を覚えなければいけないのでデータ構造を工夫しなければなりません。それがe-graphです。

e-graph

e-graphは上で言うS式とその合同^[10]関係をグラフにして保存します。グラフの表現力はすごいので、指数関数的に増えていくS式を効率的に表すことができます。

e-graphはnodeとe-classを保存します。nodeの子はe-classで、e-classは合同なnodeの集合です。また、nodeは常にただ一つのe-classに属します。
上の例でいうとnodeはS式を表し、e-classは合同なS式の集合を表します。例えば(* a 1)とaはルール2によって合同だとみなせるので同じe-classに属すはずです。

e-graphに対してルールを適用していくと、効率的に空間を使いながら、今までの情報を失うことなく新しい情報をe-graphに足していくことができます。

eggのサイトにある図. (/ (* a 2) 2)を書き換えていく様子を表している

図の実線で囲われている四角はnodeを表し、点線で囲われている四角はe-classを表しています。
nodeから伸びている線はすべてe-classにつながっていることに注意してください。

上の例でいうと(/ (* a 2) 2)は図の左端のように表現できます。例えば2はそれ自体がnodeで2と合同なnodeは今のところ存在していないので2一つで一個のe-classです。(/ (* a 2) 2)には2回2が出てきますがどちらも同じnodeを指しているので同じe-classに線が来ています。

(* a 2)を(<< a 1)に書き換えたのが左から二番目です。(* a 2)と(<< a 1)はルール1によって合同なので同じe-classに属しています。

そういう具合にe-graphにルールを適用していくと最終的にこれ以上ルールを適用できなくなり図の右端のようになります。(/ (* a 2) 2)とaが同じe-classに属しているので、(/ (* a 2) 2)はaになりました🎉。

図の右端の右上の*のnodeは自身が属するe-classを指しているので、(* a 1)、(* (* a 1) 1)、(* (* (* a 1) 1) 1)...のすべてがe-graphに保存されていることに注意してください。すごい

e-graphを具体的にどう実装するかは本記事では触れませんがegg: Fast and extensible equality saturationによく書いてあるので興味があれば読んでみてください。

egg

eggはe-graphのライブラリです。egglog自体はeggに依存していませんが実質同じようなものがegglogに入っていると考えていいです。

eggの特徴は、ユーザーが明示的にEquality Saturationを実行するように設計されていることです。
古典的なe-graphの実装ではユーザーがe-graphになにか操作するたびに自動的にEquality Saturationを実行するのですが、eggではユーザーのタイミングでまとめてEquality Saturationを実行するのでより効率的になっています。詳しくはegg: Fast and extensible equality saturation。

なのでegglogでも上の方にあるサンプルコードのように、Equality Saturationを明示的に(run 100)みたいなやつで実行する必要があります。

egglog

前置きが終わったのでegglogの話に入ります。

環境構築

手軽に試したい場合はWeb上で動くデモがあります。

ローカルで実行する場合は特にバイナリ等はまだ配布されていないので、cargoをでビルドします。
cargoはRustの開発環境についてきます。

cargo install --git https://github.com/egraphs-good/egglog

自作ですがVSCodeの拡張もあるので入れておくと便利です。公式の拡張も一応ありますが全然更新されてないのでおすすめしません。
egglog-language

実行方法

他の言語処理系と同じように、egglogにソースコードを渡すと実行されます。拡張子は.eggにするのが慣習です。

egglog ${file}

egglog -hで他に何が出来るかわかりますが、特に重要なのは以下です。

• egglog --desugar ${file}: 糖衣構文を展開して出力する
• egglog --to-dot ${file}: 実行結果のe-graphのグラフを出力する

VSCodeの拡張を入れている場合は糖衣構文の展開結果をマウスオーバーで見ることができます。

(datatype BDD (ITE i64 BDD BDD))は(sort BDD) (function ITE (i64 BDD BDD) BDD)に展開されることがわかる

egglog --to-dot path.eggの結果。わかりやすい

リファレンス

rustdocにあるドキュメントが実質リファレンスとして使えます。

egglogのサンプルコードが読みたい場合はtestsを読むと良いです。

以下では重要かと思われる機能を紹介します。
上の方でegglog自体e-graphをうまく使うためのインターフェースといった感じがあると書きましたが、以下で説明する機能はすべてegglog処理系がグローバルに持っているe-graphになにか操作するものだと思うとわかりやすいです。

文法

egglogのパーサーはlalrpopで実装されているので比較的短いBNFっぽいDSLを見れすべての文法がわかります。

組み込みのデータ型と関数

src/sortに組み込み型と関数が定義されています。例えばi64って何ができたっけ？となった場合にはsrc/sort/i64.rsに目を通すと良いです。

以下に現在サポートされている組み込み型を列挙します。これで全てです。

• Unit いわゆるユニット型。リテラルは()
• bool 真偽値。リテラルはtrueとfalse
• i64 符号付き64bit整数 今のところ符号なしの整数型はない
• f64 64bitの浮動小数点数
• Rational 有理数 (rational 1 2)とかで作る(この例は1/2を作る)
• String 文字列
• set 集合 本記事では扱わない
• map 連想配列 本記事では扱わない
• vec 配列 本記事では扱わない。再帰的なデータ構造を使えばそれで足りる場面も多い。

egglogにはジェネリクスのような機能がまだないのでset、map、 vecの使い方は少し特殊になっています。テストコードを見ると使い方がわかります。

Sort

Sortはe-classの識別子のような概念です。たぶん数学的に意味のある用語なのでしょう。同じSortを持つnodeは同じe-classに属します。組み込み型もSortですが、後述するunionができないという制限があります。
Sortは以下のように宣言します。

(sort NameofSort)

function

Sortだけあってもしょうがないので、functionでSortを持つnodeを定義できます。

他のプログラミング言語の関数とはだいぶ違います。単に引数のSort列をキーに、返り値のSortを値に持つ連想配列を宣言していると考えるほうが近いです。そういう意味では後述するruleが関数の本体に近いです。

例

(sort Math)
(function Const (i64) Math) ; (Const 123)とかでMathのnodeを作れる

例えば、ある2つのConstがあったとき、引数が同じe-classなら結果も常に同じSort(Math)になります。つまり、(Const 0)と(Const 0)は同じSortを指します。

union

unionを使うと2つのSortをマージすることができます。
組み込み型には使えません。例えば0と1をunionして同じSortになってしまうとめちゃくちゃになってしまいます。

例

(sort Math)
(function Const (i64) Math) ; 定数
(function Var (String) Math) ; 変数

(union (Var "x") (Const 123))

上のコードのe-graph

set

setはnodeに値をセットします。setを使うと組み込み型を返すfunctionに値をセットすることができます。これはunionではできません。

(function F (i64) i64)

(set (F 0) 1)
; (union (F 0) 1) はエラーになる

unionとsetの使い分けはごちゃごちゃしている部分がありますが、間違えたら適切なエラーメッセージが出てくるのでそれに従って書き換えればだいたいなんとかなります。

rule

正直これまでに説明した機能だけだと、だから何なんだ？って感じですが、ruleを使うと面白くなってきます。

ruleは上で説明したe-graphのルールに相当します。

ruleはe-graphにパターンマッチしてなにかアクションを起こす機能です。
そしてunionとsetはアクションです。

後述するrunコマンドを実行するとe-graphに繰り返しruleを適用していきます。

文法の該当部分を抜粋するとこんな感じです。

LParen "rule" <body:List<Fact>> <head:List<Action>> <ruleset:(":ruleset" <Ident>)?> <name:(":name" <String>)?> RParen => ...
...
pub Fact: Fact = {
    LParen "=" <mut es:Expr+> <e:Expr> RParen => ...
    <CallExpr> => ...
}
...
NonLetAction: Action = {
    LParen "set" LParen <f: Ident> <args:Expr*> RParen <v:Expr> RParen => ...
    LParen "delete" LParen <f: Ident> <args:Expr*> RParen RParen => ...
    LParen "union" <e1:Expr> <e2:Expr> RParen => ...
    LParen "panic" <msg:String> RParen => ...
    LParen "extract" <expr:Expr> RParen => ...
    LParen "extract" <expr:Expr> <variants:Expr> RParen => ...
    <e:CallExpr> => ...
}

pub Action: Action = {
    LParen "let" <name:Ident> <expr:Expr> RParen => ...
    <NonLetAction> => <>,
}

例

(sort Math)
(function Const (i64) Math) ; 定数
(function Var (String) Math) ; 変数
(function Add (Math Math) Math) ; 足し算

(rule
    ((= expr (Add (Const x) (Const y)))) ; 文法上はFactのリスト。`=`でマッチした対象を変数に束縛している
    ((union expr (Const (+ x y)))) ; Actionのリスト ; `+`はi64の足し算
)

; 式を定義する
; e-graphにこの式を追加するというような意味合いもある。知らないものはEquality Saturationできないので。
(let e (Add (Var "a") (Add (Const 1) (Add (Const 2) (Const 3))))) ; a + (1 + (2 + 3))

; Equality Saturationを実行
(run 100)

(extract e) ; (Add (Var "a") (Const 6)). a + 6

上のコードのe-graph

上の例ではruleを使ってAddの引数がどちらもConstの場合にはそれを事前に計算してしまうようにしています。ruleにないVarとかは無視して出来るところだけやってくれるのでコンパイラで言うConst Propagationみたいなことを実質やっています。

実際これはすごいことです！例えばこれを拡張していけばきれいにsympy.simplifyやTI-Nspire™ CX IIについてるCASのようなものを作ることができます(破綻しないようにルールを追加していくのは非常に難しいですが😰)。しかも計算が間に合えば常に最適な結果を出力してくれます。

:merge

:mergeによってfunctionに違う値が設定されたときの挙動を変化できます。:mergeはfunctionの宣言のときに指定します。

(function F (i64) i64 :merge (min old new)) ; :mergeの中では`old`と`new`が使えるので複数の値がある場合にどうするかを指定できる。この場合は小さい方を採用する

(set (F 0) 0)
(set (F 0) 1) ; (min 0 1) => 0　なので意味ない

let

もう上で使ってしまいましたが、letは変数に値を束縛します。再代入はできません。
この文法じゃあグローバルな名前空間がぐちゃぐちゃになってしまうじゃないかと思うかもしれませんが回避策もあります。

それが、pushとpopです。この2つのコマンドでスコープを作ることができます。スコープを作るときにはpushを、スコープから抜けるときにはpopを使います。

(push) ; スコープを作る
(let x 0)
(pop) ; スコープから抜ける

(let x 1) ; 上でxを定義したスコープはすでにないのでOK

run

runはEquality Saturationを実行します。runを実行するとruleに従ってe-graphを書き換えていきます。runは最大で引数にいれた回数だけ実行されます。saturateしたらその時点で終わります。

他にも色々オプションがありますが、とりあえず適当に(run 100)とかやっておけばいいです。

extract

extractはe-graphから引数のe-classの中で一番低コストなnodeを出力します。何をもって低コストかをコントロールする方法もあるのですが本記事では触れません。

糖衣構文

いくつか便利な糖衣構文を紹介します。

datatype

datatypeはsortとfunctionまとめた糖衣構文です。

(datatype Math
    (Const i64)
    (Var String)
    (Add Math Math)
)

は

(sort Math)
(function Const (i64) Math)
(function Var (String) Math)
(function Add (Math Math) Math)

に展開されます。

rewrite

rewriteはruleの糖衣構文です。
簡単なruleはrewriteで書くことができます。

(rule
    ((= expr (Add (Const x) (Const y))))
    ((union expr (Const (+ x y))))
)

と

(rewrite
    (Add (Const x) (Const y))
    (Const (+ x y))
)

は同じ意味です。

実際には機械的に

(rule ((= v2___ (Const x))
       (= v4___ (Const y))
       (= rewrite_var__ (Add v1___ v3___))
       (= v2___ v1___)
       (= v4___ v3___))
      ((let v6___ (+ x y))
       (let v5___ (Const v6___))
       (union rewrite_var__ v5___))
         :name "(rewrite (Add (Const x) (Const y)) (Const (+ x y)))")

と展開されます

rewriteに限らずegglogが糖衣構文を展開した際はv{int}___のような名前を自動で付けます。これをユーザーが他の場所使うと普通に衝突するので使わないでください。

relation

functionは連想配列のようなものだと書きましたが、relationは値がユニット型の連想配列、つまりセットのようなものを宣言する糖衣構文です。

(relation Prime (i64)) ; 素数かどうかを表したいとする。単純に素数かそうでないかを表すだけなので返り値の意味はない。ユニット型でいい。

は

(function Prime (i64) Unit :default ()) 

に展開されます

(Prime 5) ; :default ()が設定されているので、これでPrimeのnodeを作ることができる
(set (Prime 2) ()) ; :default ()が設定されていないとこうしなければならない

ちなみにユーザーが定義したsortはe-graphにないものが必要になった場合新しいsortが自動で作られますが、ユニット型はそうはいかないので:defaultが必要になっています。

(sort S)
(function F () S)

(F) ; 自動で新しいSが作られる

declare

declareは新しいSortを変数に割り当てます。

(sort S)
(declare x S)

は

(sort S)
(function v1___ () S)
(let x (v1___))

に展開されます。

ただ、現状可視化したときには関数名しか出てこずxがどこにあるかわからないのであまりおすすめしません。

(sort S)
(function x () S)
(x)

現状上のように書かないとこの図はでてこない

おわりに

本記事ではegglogにある機能の一部を紹介しました。
もっと色々知りたい方はソースコードやテストコードを読んでみてください。

宣伝

egglogでコンパイラの最適化パスを実装してみたい方は一旦CFGをRVSDGに変換するとやりやすいです。
個人的にbriltで実験しているので興味のある方はぜひ🙇

脚注

1. Datalogについてはよくわからないので本記事では触れません🙇

2. チューリング完全ではあると思うが例えばwebサーバーを書くといったことには適さないという意味

3. 浮動小数点数の演算をより誤差が少なくなるように書き換えるプログラム

4. S式みたいな雰囲気のもの。(x . y)みたいなやつは考えない

5. 例えば葉の数が少ないとか

6. $a \times 2$は絶対にオーバーフローしないとします

7. $x=0$の場合、$0/0=1$となってしまいとても危ないので明示的に禁止している。ゼロ除算#ゼロ除算に基づく誤謬のようなことがそのまま起きる。え、じゃあルールを作るのってめっちゃ難しくない！？と思ったあなた、その通りです。ルールを自動で導出しようという研究もあるので興味のある方はRewrite Rule Inference Using Equality Saturationがおすすめです

8. ASIAN KUNG-FU GENERATIONの名曲が思い出されます

9. 実用的にはSaturateしなくてタイムアウトで終了することもある

10. congruence。"等しい(equal)"とは少し意味合いが違うがとりあえず特に気にしなくて良い