コンパイラのための新しい中間表現 MLIR の論文を読む

LLVM IRについて調べていると、最近よくMLIRという言葉を目にするようになった。MLIRは"Multi-Level Intermediate Representation"の略称であって決して"Machine Learning"ではないのだが、LLVM IRを置き換えるために開発されている新しい中間表現形式である。

LLVM IRももちろん柔軟性のある中間表現形式であるが、より柔軟かつ幅広いDSLの中間表現を受け入れるために開発されているものだ。Googleに在籍していたLLVM CreaterのChris Lattnerによって開発されたものである。ちなみにChris Lattnerは本論文が発行されたときにはすでにSiFiveに移っている。

https://arxiv.org/abs/2002.11054

この論文を読みながら、MLIRについて理解を進めていくことにしよう。

MLIRの動機

LLVM IRを代表するような中間言語を使用するシステムは大部分においてうまく動作をしているが、より抽象度の高いプログラミング言語を用いる場合はLLVM IRだけでは上手く行かないケースが散見される。例えば、Swift/Rust/Juliaなどは直接LLVM IRには変換されず、複数段階のIR表現形式を使用して最終的にLLVM IRに変換されている。

f:id:msyksphinz:20200517124959p:plain

論文中では、既存のIRについて以下のような表現が使用されている。

LLVM中間表現(IR)は大まかには「ベクトルを持つC」
JVMは「ガベージコレクタを持つオブジェクト指向型システム」の抽象化

C++コードなどの非常に抽象化されたソースコードの解析はLLVM IRでは非常に難しく、これを解決するのがMLIRと言うことになっている。

スタティック・シングル・アサインメント(SSA)ベースのIRデータ構造の標準化
宣言的システムによりdialects(方言)を数多くサポートする
共通のインフラストラクチャの提供

MLIRの基本

以下がMLIRの基本的な考え方となっている。

すべてをカスタマイズできるようにする。これにより将来の問題にも適用できるようにする。
SSAとリージョンを導入する。SSAはともかく、リージョンの概念を導入して入れ子による制御フローやデータフローをサポートする。LLVM IRではリージョンによる入れ子の表現はあまり見たことが無く（というかほとんどLLVM IRはアセンブリに近しい気がする）、リージョンの概念を導入することで問題を抽象化する。
プログレッシブな変換。既にLLVM IRについても、LLVM IRから始まりSelectionDAG、MachineInstr、MCInstへと段階的に変換が行われているが、これをより柔軟化する。
高レベルのセマンティクスを維持する。これは一度表現形式を低レベルに変換してしまうと効率的な解析が難しくなることに起因する。可能な限り高レベルの表現形式を維持する。
IR検証。IRを検証可能な表現形式で維持する。
宣言的な書き替えパタン。つまり変換パタンなどの定義を非常に簡単に定義できなければならない。
ソース位置の追跡とトレーサビリティ。つまりシステム内で変換結果を見ても元のプログラミング言語との対応が取れやすくしなければならない。