こちらはNEC デジタルテクノロジー開発研究所 Advent Calendar 2023の12月15日の記事です.本日は私の誕生日ということもあり,私の人生を振り返りながらFireDucksが中で利用している技術を紹介したいと思います.
FireDucksについてはこちら
コンパイラとの出会い
FireDucksはコンパイラ技術を使ったデータフレームライブラリですが,私がコンパイラに出会ったのは,大学2,3年生のころに受けたコンパイラの授業でした.当時はプログラムを算譜,コンパイルを翻訳と呼ぶような時代だったのですが,先生から提供されたトイ言語用のコンパイラに何か新しい機能を実装するのが最後の課題でした.それまでプログラミングは授業でちょっと習った程度の私には大分ハードルが高い課題で,夜な夜なPCに向かって格闘した記憶があります.このトイ言語にはfor文が無かったのでfor文の実装をテーマに選び,forがあるならi++
もあったほうがいいよねと思ったら++
の実装が非常に大変...というようなことをやりながら,プログラムを読んだり書いたりするのは楽しいことだと学びました.
この経験もあり大学の研究室は自動並列化コンパイラをやっているところに入り,そこでコンパイラでプログラムを高速化するという研究をしてました.この高速化はライフワークとなり,その後就職してからも一貫して高速化に携わってきましたが,十数年コンパイラに関わることはありませんでした.
FireDucksを支える技術その1: MLIR
その後私が数年間関わったのがNECが開発してるスパコンSX-Aurora TSUBASAで,これに搭載されているVector Engine(VE)というアクセラレータ用にプログラムの高速化を行っていました.高速化を行うためにはIntrinsicsが欲しくなり,その為にLLVMを用いてVE用のコンパイラllvm-veを開発しました.これはその後同僚の手によってUpstreamingされて,今もメンテされているようです.
この活動は久しぶりにコンパイラに関わる楽しい活動だったのですが,llvm-veの発表をしにLLVM Developers' Meeting(たしか2019年)に参加したときに知ったのが,FireDucksを支える技術その1のMLIR(Multi-Level Intermediate Representation)でした.MLIRは独自の中間言語を開発するためのフレームワークで,ドメイン特化のIRを設計し,独自のコンパイラを作ったりできるものです.強力なLLVMのインフラを利用できるので,コンパイラを開発したいと思ってる人には非常にありがたいものです.この辺りから,またコンパイラをやりたいなと思い始めました.MLIRについては日本語でも解説記事がありますので,ご興味がある方は読んでみてください.
FireDucksを支える技術その2: TFRT
llvm-veの開発と並行して行っていたのが,tensorflow-veの開発です.これはTensorFlowにVector Engineサポートを追加したものです.当時はTensorFlowにはXLAと呼ばれる実行時コンパイラが組み込まれており,pytorchのGlowやApache TVMなどちょうどDeep Learningコンパイラが流行っている時期でした.DLコンパイラは計算グラフの最適化や多様なアクセラレータのサポートのために重要な技術ですが,当時はまだそこまで整備されておらず,当時のTensorFlowのアクセラレータ周りはGPU依存の部分も結構あり,Vector Engineサポートはそれなりに苦労しました.
この経験から,ドメイン特化の中間言語が設計できるMLIRによって,DLコンパイラのような実行時コンパイラが色々な領域で使われるようになり,APIと実装が分離したアクセラレータ対応がしやすいソフトウェアアーキテクチャ(Disaggregated Computingのソフトウェア版)が広まるといいな,そして色々なソフトウェアのVectorEngine対応が簡単になったらいいなというのがFireDucksの開発の最初のモチベーションでした.
ちょうどそのころGoogleがTFRT(A New TensorFlow Runtime)を発表しました.これがFireDucksを支える技術その2です.TFRTは,TensorFlowのグラフコンパイラが出力する(MLIRで記述された)計算グラフを実行する機能を提供するものです.これは良いものが出たと思ってコードを読み漁ったところ,TFRTのコアはTensorFlowやDeep Learningとは独立しており,MLIRで記述された中間言語中の各命令(Op)を登録されたカーネル関数を使って実行する機能で,実際の実装もコア部分とその周辺であるDL用の中間言語やカーネルとはきれいに分離されており,汎用的に作られてることが分かりました.これを使えば,独自の中間言語に対してカーネルを用意すれば,あとの実行管理(中間言語にしたがってスタックフレームの管理などをしながらカーネルを呼び出してくれる機能)はTFRTがやってくれるので,実行時コンパイラを作るのが大分楽になります.
FireDucksを支える技術その3: Apache Arrow
その後調べものをしたり色々やって,データ分析におけるデータの前処理の高速化が大事だろう,そこで良く使われるpandasが速くなったら嬉しいだろう(pandasがSX-Auroraで動いたら嬉しいだろう)とうことでFireDucksプロジェクトを開始しました.そのときに最もきっかけとなったのがこちらのこちらの論文: Towards Scalable Dataframe Systemsです.これはFireDucksと同じくDataFrameライブラリであるmodinの作者の書いた論文なのですが,(私の理解したところでは)DataFrameに対する様々な処理は,DataFrame Algebra上の操作の組み合わせとして表現可能で,代数上の変換として最適化を行い,各操作に対する実装(カーネル)を用意すれば,DataFrameライブラリが作れると言うことを解説してくれています.しかもその実装例としてmodinがあるということで,modinにMLIRとTFRTを組み込んで,データフレーム用の最適化をやるデータフレームコンパイラを作ろう!というのがFireDucksプロジェクトのスタートでした.
当初は実行にはSX-AuroraのVector EngineやGPUを使うことを考えていたので,カーネルとしてはそれぞれ用のデータフレームライブラリであるFrovedis(正確にはFrovedisの一部にデータフレームがある)やcudfを使う予定で,中間言語からライブラリに落とし込むラッパーだけを作ればよいと考えていました.ところがやはりCPUでの高速化のニーズが高いので,まずはCPUからやるかとなったのですが,調べたところCPU用の良いデータフレームライブラリが見つからず,しょうがないから自分たちで開発することにしました.その時にデータ構造として使うことにしたのが,FireDucks支える技術その3のApache Arrowです.
Apache Arrowは,テーブルデータ(カラムナー)のためのデータ構造を提供するライブラリで,非常に多くのプロジェクトで利用されているので,ほぼ一択と言う感じでした.pandasの作者のWes McKinneyさんが開発されていて,pandasの悪いところを踏まえてやっているプロジェクトなのも良い点だと思いました.そのあたりはご本人のブログ記事があります.
Wes Mackinneyさんは,このpandasの悪い点の解決策としてpandasにこだわらないという選択をされたのかと思いますが,pandas APIから中間言語を生成してApache Arrowで作られたカーネルで実行するというコンパイラレイヤーを入れれば,padansとAPI互換でありながら欠点の解決もできると思ったわけです.
Apache Arrowは3ヵ月にいちどメジャーバージョンが更新されるペースで非常に活発に開発されており,我々がカーネルと呼ぶ演算部分(arrow::compute)は,現在もどんどん整備が進んでいるところの様です.ということで,Apache Arrowがサポートしてる演算を使いつつ,サポートしてない演算だったり(開発当初はgroupbyも無かったんですよね),並列化されてない演算などを補ったりしながら開発しているのが現在FireDucksのCPU用のバックエンドです.
Apache Arrowはデータ型(整数,浮動小数点,時刻など)や複数の部分領域(チャンク)に分割された配列(ChunkedArray)などを抽象化してくれるので,ある程度それらを意識しないでカーネルを書くことができます.メモリ管理もやってくれますし,並列処理のための枠組みもカーネル開発に便利です.当初はカーネル開発は余り考えてなかったのですが,こういうカーネルレベルの高速化も楽しいところなので,Apache Arrowの助けを借りながら頑張ってやっております.
おわりに
本記事ではFireDucksが中で使ってる技術として,MLIR, TFRT, Apache Arrowを取り上げました.これらの技術のおかげで,FireDucksの主な開発はfrontend(pandas APIから中間言語を生成するところ)とカーネルとコンパイラ最適化です.特にfrontendは,結局modinを利用することはあきらめて独自開発しており,pandasとの互換性を上げるためのfallback機能だったり,コンパイラ最適化を支援するための依存解析だったりと独自機能を実装しています.この辺りもいずれ解説したいと思います.
frontend周りはいずれは汎用化して独立させたいと考えているのですが,そうするとfrontend, 中間言語(MLIR),runtime(TFRT)と一通り道具がそろうので,実行時コンパイラを内蔵しAPIと実装が分離されたアーキテクチャを採用するソフトウェアが増えるといいなと思っています.
最後になりましたが私はこちら.