こんにちは、低レイヤーが好きなsoichiroです。
この記事では現在開発中のFlunoという言語について紹介していきます。
Fluno は、現実世界と融合するような次世代AIの実行基盤を目指して設計している言語です。
Flunoは主にRustを使用して開発している言語で、特徴は確率的プログラミング、リアクティブプログラミング、非同期処理、自動微分、などの意味論を統合し、さらには実行前に型を確定させることで同時に使用不可能な意味論が発生した場合にはエラーを出し実行時エラーを防ぐ仕組みになっています。
また、いくら意味論を統合して便利に書けるとしても速くなければ意味がないと考えているため、中間表現の段階から最適化できる部分を最適化し、また、数値計算などではJITコンパイルすることで一部のタスクではRustと並ぶ水準のパフォーマンスを実現しています。
ただ、もちろん全ての実装が終わっているわけではないので、まだ実行できないような構文もありますし、エラーが出ることもあると思います。
ちなみにFlunoの名前の由来はfluxとunoを合体させたことから来ています。
全て一から説明するととても長くなってしまうので、少し難しい用語なども出てくるかと思います。そういったものはAIに聞いたり質問したりしてもらえればと思います。
Flunoの開発動機
一応なぜFlunoという言語を作ろうと思ったのか書いておこうと思います。
まず、僕はAIにそこそこ興味を持っていて、自分でAIの簡易的なモデルを作成していました。
でもやっぱり既存のアーキテクチャを使っているようなAI(いわゆるANN:アテンションニューラルネットワーク)は多くの人が研究しているし、自分がやっても車輪の再発明のようになるだけだなと思い他の面白いアーキテクチャはないかなと思って考えていたところ、SNN(スパイキングニューラルネットワーク)というものに出会い、面白そうだなと思い触れてみました。
そしてこの時、同時に割とまともな言語を作ってみたいなと思っていたのと、「確率型みたいなのあったら面白そうじゃね」と思っていたので(自分が無知なだけですでに確率型というものはありましたが)、SNNを実行するのに適している言語を作ってみようと思いました。
ちなみに、この時の用途はSNNを想定していましたが、開発しているうちにFlunoの用途は結構広く、例えばエッジデバイスでのAIワークロード、ロボティクス、HFTなどにも向いているんじゃないかと思っています。
Flunoのアーキテクチャ
とても大雑把に分けるとFlunoのアーキテクチャはフロントエンドとバックエンドの2層に分類されます。フロントエンドの役割は型の確定と実行時エラーの排除で、バックエンドの役割は最適化と高速化です。
フロントエンド
フロントエンドの流れは大体下図のような感じです。
Flunoでは型とエフェクトという2つを明確に区別し、それぞれについて矛盾がないかフロントエンドで確かめます。エフェクトとは、その変数を扱う計算や関数がランタイムに対して引き起こす効果のことです。
一般的な静的型付け言語では、型自体にエフェクトの意味が含まれているため型とエフェクトを分離する必要がありませんが、Flunoは特殊なエフェクトをいくつも持っているため、分離する必要があります。例えば自動微分される変数は表面上の型はFloat型ですが、内部的にはdiffというエフェクトを保持しています。また、エフェクトは非同期処理の伝播などを思い浮かべてもらえるとわかりやすいと思いますが、関数の呼び出しによって伝播していきます。そのため型とエフェクトを分離して考える必要があります。
バックエンド
バックエンドの流れは大体下図のような感じです。
Flunoの最適化はいくつかのレイヤに分けて行われています。
1,ASTからCore IRへの変換
まずフロントエンドから渡されたAST(抽象構文木)をCore IR(中間表現)に変換します。
2,MLIRによる最適化
数値計算などのコードはMLIRのFluno dialectという独自のdialectに変換され、機械語レベルまで落とされます。また、この時AOTコンパイルとJITコンパイルの2つが組み合わされます。将来的にはこのMLIRによる最適化の範囲を拡大していきたいと考えています。
- AOTコンパイル:静的にループ条件やサイズが確定しているものは、MLIRからLLVM IRを経由して、ターゲットハードウェアに最適化されたネイティブバイナリへ事前コンパイルされます。
- JITコンパイル:実行時の入力データの形状を見て初めて最適化できるものはLLVM/MLIRベースのJITコンパイルが行われます。
3,BytecodeVMへのフォールバック
全てのコードをMLIRに落とし込むことはできないため、現状対応している数値計算の部分以外はFluno BytecodeにコンパイルしてからBytecode VMでインタプリタ的に実行されます。
メモリ管理
個人的にはFlunoを設計していた時に一番面白かったところはメモリ管理です。
Flunoは単一のメモリ管理方式を採用するのではなく、計算の性質に応じて異なるメモリ管理方式を使い分けています。
そして、一番の特徴としては静的に確定するものに関してはコンパイラが自動で必要なメモリの量を確保し、計算が終わった瞬間にその部分のメモリを一括で破棄することで安全性と高速性を確保しています。それが適用できない部分ではものによって参照カウントとマークアンドスイープの2種類のGCを使い分けています。
なぜFlunoか
世の中にはRustやPython,そして数々のDSL(ドメイン特化言語)があります。その中でなぜFlunoが必要なのかを話します。つまりFlunoのアピールをします。
まず、Flunoの主戦場は開発動機の章でも述べた通り、時間変化と確率的な数理モデルが必要とされ、さらにリアルタイムな処理が必要とされるような領域です。
このような領域においての現在の主要な3つの選択肢(Python,Rust,既存の数理DSL)との比較をします。
1,Python
Pythonのエコシステム（PyTorchなど）を使えば、高度な数理モデルを直感的に記述することは可能です。しかしこれらはPythonのランタイムからC++などで作られたバイナリを呼び出しているため、メモリ競合などの予期せぬエラーの発生や遅延が生じます。
しかしFlunoは数理的計算を言語のコアに入れ込んでいるため、直感的な記述はそのままに速度を向上することを可能にしています。
つまり、FlunoはPythonの書きやすさはそのままに速度を向上させ、遅延もなくした言語と言えます。
2,C++
C++とFlunoの決定的な違いはコンパイラによって最大化できる範囲の違いです。一般的なC++コンパイラではドメイン固有の意味論（確率モデルやリアクティブな依存関係など）を考慮した高レベルのドメイン特化最適化が難しいのに対し、Flunoでは言語仕様としてこれらの意味論をコンパイラが把握しているため、よりドメインに特化した最適化（ドメイン固有の表現の簡約など）が可能になります。
3,Rust
おそらくRustの一番の特徴はメモリ安全であることです。しかしその代償として所有権を記述してコンパイルを通すのがとても面倒で、記述も複雑になり開発速度も遅くなると思います。また、C++と同様に最適化が一般的な式の最適化までしかできません。ちなみにアーキテクチャのメモリ管理の部分を見てもらえると少しわかってもらえるかもしれませんが、Flunoは結構メモリ安全に作れていると思います。Rustほどじゃないですが
4,DSL
DSLは例えばJAXやStanなどがあります。これらは単体ではとても速いですが、実用的ワークフローで使用するには異なるDSLを繋ぐいわゆるグルーコードのようなものが必要になり、これらの間のデータのやりとりだけで相当な遅延が発生してしまいます。Flunoは一つの言語内で完結するためこれらの問題は起こりません。
まとめ
つまり言いたいことは、Flunoは速くて記述しやすくて安全ということです。
ちなみにこれは一部のワークロードにおいては現在でも言えるかもしれませんが、まだ開発中のため一般には言えません。
ベンチマーク結果
Flunoのベンチマークでの結果をお見せします。ただ、これはマイクロベンチです。
あと、これらのベンチマークは値を動的に与えることで、AOTコンパイルで定数畳み込みなどはできないようになっています。
まずこれはプロセス内部の純粋な計算関数の実行時間だけを切り出して測定したものです(カーネル実行時間と呼ぶことにします)。
| ワークロード | Fluno (primary) | Rust | Node.js | Python |
|---|---|---|---|---|
| 確率的FRP・逐次自動微分 | 0.082 ms | 14.082 ms | 67.724 ms | 171.151 ms |
| テンソル要素ごとの連鎖演算 | 0.083 ms | 15.105 ms | 79.309 ms | 179.646 ms |
| 1次元畳み込み演算 | 0.136 ms | 14.081 ms | 90.087 ms | 170.997 ms |
| 行列乗算 | 1.138 ms | 14.235 ms | 78.275 ms | 308.721 ms |
| 配列の簡約化 / 総和 | 0.065 ms | 14.096 ms | 79.589 ms | 198.560 ms |
| 1次元ステンシル計算 | 0.118 ms | 14.156 ms | 67.336 ms | 226.514 ms |
ちなみにprimaryとは変換可能なプログラムをMLIRに落とし込む実行パスのことです。
次に外部プロセス起動から実行、検証までのプロセスの総時間を測定し、それをループ回数で割った時間です(プロセス総時間と呼ぶことにします)。
ただ、全て載せると長くなるので確率的FRP・逐次自動微分だけ載せます。
| 実装言語 | repeat = 1 (ms/call) | repeat = 100 (ms/call) | repeat = 10000 (ms/call) |
|---|---|---|---|
| Fluno (primary) | 252.152 ms | 2.942 ms | 0.137 ms |
| Rust | 8.502 ms | 0.095 ms | 0.002 ms |
| Python | 228.605 ms | 2.233 ms | 0.080 ms |
| Node.js | 76.494 ms | 0.714 ms | 0.011 ms |
これらの表を見て分かることとして、Flunoはカーネル実行の方ではとても速いですが、プロセス総時間ではPythonよりも遅くなってしまっていることがわかります。ちなみにrepeat=1でRust以外が遅いのが他の言語はJITコンパイルするコストがかかっているからです。RustのAOTコンパイルの時間は含まれていないので。そして、本題のFlunoがrepeat=10000などでも遅いことに関してですが、これはFlunoがまだMLIRを広げきれてないということが原因です。現状のFlunoはループの部分までMLIRに落とし込むことができないため、いちいち1回のループごとにMLIRとコンパイラ基盤でセットアップを行わなければいけないので遅くなります。
つまり、ループの部分などもMLIRに落とし込むことが出来ればFlunoはカーネル実行時間のように、Rustさえ凌駕する実行時間をたたき出すことができるかもしれないということです。
補足:測定条件について
- Rust側は、cargo build --releaseでビルドされたネイティブバイナリを使用しています。
- すべてのベンチマークは、値を引数や外部ファイルから動的に与えることで、RustやFlunoがコンパイル時に計算を終わらせてしまう定数畳み込みなどが起きないように対策した上で測定しています。
ロードマップ
Flunoの今後の展望について書きたいと思います。
Flunoはまだコアの部分しかできていません。そのコアの部分も満足な性能が出せているわけではないので、今後Primary中心に高速化を目指していきたいと思います。また、Primaryの適用範囲も狭いため、もっと広げていくこと、そしてゆくゆくはSNNを動作させるのに適したニューロモルフィックハードウェアで動作できるようにDialectを変換することもできるようにしていきたいと考えています。また、今後実際の現場のワークロードでもFlunoを使うことができるようにエコシステムも磨き上げていきたいと思います。
おわりに
最後までお読みいただきありがとうございます。
疑問点やフィードバック等ございましたら是非コメントや、メールでも何かしら送っていただけると助かります。
リポジトリは以下です。スターなどしていただけると励みになります。
このリポジトリ内にはAI作成のドキュメントも多いので、順次更新していくつもりです。
あとコア部分は徐々に公開していくかもしれません。