LLVM/C++日記　計算グラフを作ってみる

前回からの続き．LLVM/C++の勉強日記です．以下は，わいわいswiftcオンラインで発表した内容のあらすじです．Discord上での発表は非常にスムーズで，質疑応答で，fixstarsの人が色々コメントや質問をしてくれたりして，非常の楽しいものでした．特に，計算グラフの実装のコードをちょっと褒めてもらえたのでうれしかったです．

計算グラフを作ってみた

コードはこちら．計算グラフとは，代数的な計算を抽象的にメモリ上で表現し，後から微分したり，色々できたりする感じの構造のことです．例えば，以下のコードは，my_funcにx-yという数値の計算結果を代入するのではなく，x-yという式を代入し，それをコンパイルして，後から，実行できるようにしています．

int main() {
   auto visitor = IRVisitor();
   Variable x, y;
   Func my_func = x - y;
   Func my_func2 = x + y;
   visitor.visit(&my_func);
   visitor.visit(&my_func2);
   std::cout << my_func(1, 2) << std::endl;
   std::cout << my_func2(20    , 2) << std::endl;
   return 0;
}

まさに，コードを実行することで，関数が定義されるわけで，"Define by run"と言われたりします．計算グラフは，色々なシーンで利用されますが，昨今は，ニューラルネットワークの学習時の計算に，設計したモデルをパラメータで微分した関数が必要なことから，計算グラフがPyTorch, TensorFlowなどで利用されています．私の理解が間違っていなければ，この"Define by run"の考えをニューラルネットワークのライブラリに最初に導入したのがPFNのChainerです．この考え方は，他のニューラルネットワークのライブラリに取り入れられ，広まっていきました．ニューラルネットワークの研究の発展に，Chainerは明らかに貢献したと行っても過言ではないと私は思います．

実現したいこと

まぁ，計算グラフを作って実行するだけなら，C++だけで実装することもできるのですが，Halideの真似事をしたかったので，LLVMで，計算グラフのバイナリを動的に生成することにしました．

実装

上記のコードは，ちゃんとclang++でコンパイルできます．C++の演算子をオーバーロードして実装しました．実を言うと，この実装では，綺麗に抽象化できておらず，もう一段階実装を考えないといけません．これは，Halideの実装を見ればわかるのですが，C++上で宣言する変数のプレースホルダと，実際にパースして構築する抽象構文木のノードの実装を分けないと，メモリ管理がうまくできないのです．この辺は，次回にどこかで解説したいなと考えていますが，すでに，同じリポジトリで実装済みではあります．先に見たい人は，こちらをどうぞ．

LLVMが生成した関数の実行

swiftcでも言及しましたが，LLVMが生成した関数は，C++の関数へのポインタとして取得できるのですが，C++の構文上，決め打ちした関数の型しか受け取れません．この辺が次の課題になってきます．これもとりあえず，実現方法すでに見つけております．これも，次回どこかで発表したいなぁと思っています．

残課題

LLVMのJITコンパイラにはいくつか種類があり，現在は，MCJITとORC JITというのがあり，ORC JITが新しいもので，これからは，LLVMの開発の方向性としては，ORC JITを推奨していくそうです．実際に，LLVMのチュートリアルも，途中からORC JITを使うように変更されているようです．この二つの違いがいまいちよくわからないので，この辺をもうちょっと深堀していこうかなと思っています．

発表資料