Mojoは「C言語のように速いPython」なのか

LLVMやSwiftを作ったChris LattnerがCEOをやっている会社が、Pythonの使用感とC言語並の性能を併せ持つ言語としてMojoをアナウンスした。 まだ手元で試せる状態でリリースされてはいないが、最大35000倍Pythonより速いという。

僕は大学院のディープラーニングの授業や仕事などでPythonはある程度使った経験があり、Pythonに似た性能特性を持つRubyJITコンパイラを書く仕事をしているので、Mojoの何がC言語並に速いのか興味を持った。Mojoローンチ動画ドキュメントを見てわかったことをまとめておく。

どんな言語なのか

Pythonとの互換性

MojoにはPythonにはない機能がいろいろとあるが、ローンチ動画では全く同じコードがCPythonとMojoの両方で動くことが何度か紹介された。これはつまり、単にMojoPythonの資産を使えるだけでなく、MojoPythonのスーパーセット言語であることをアピールしたいように見える。Pythonに対するMojoの関係は、Javaに対するKotlinではなく、JavaScriptに対するTypeScriptのようなものということになる。

一方、Why Mojoのあたりを見ていると、現状の互換性は散々で、クラスもサポートしてないみたいな状態らしい。彼らはスタートアップ企業であるから、お金や人を継続的に惹きつけるために多少の誇大広告が必要なのは理解できるが、現時点でリリースに至ってないのはまだ実際にはPythonのスーパーセットにはほど遠い状態であるからと推察できる。Python3との互換性のためにCPythonを活用しているというよくわからない説明があるが、Pythonモジュールのインポートが普通のimport文ではなく Python.import_module なのはこれがCPythonそのものを呼び出して相互に動作させているからかなと思った。

とはいえ、Python2からのPython3移行の例を引き合いに出して、Pythonのサブセットを実装することでPythonの弱点を解消するのは既存のPython資産という強みを生かせないのでダメなんだと力説しているので、少なくともPython3のライブラリが使えて、運が良ければ本当にPythonのスーパーセット言語として出てくるのではないかという雰囲気はある。

Rust風の言語機能

Rustのようにモダンな言語仕様かつ速い動的言語が欲しい、みたいな言説をTwitterで時おり見かけるのだが、Mojoはこの領域を攻めているように感じる。MojoはキャッチフレーズがAI開発向けの言語であるのと同時に、FAQではMojoは汎用プログラミング言語であると説明されている。

Rustは所有権といったセマンティクスが常に強制されるので、書いてコンパイルを通すのが面倒なことがまあまああるのだが、それらの手間をかけてスレッドセーフティや性能を追求しなくとも、アプリの特性上それが問題にならない場面はよくある。そういう場面では何の型アノテーションもない普通のPythonを書き、局所的に性能が必要な部分だけMojoの高速な書き方を使いたいみたいなユースケースは存在するように思う。

let でイミュータブルな変数を宣言したり、所有権や参照の概念があったりする。def で関数を宣言するといままで通りのPythonのコードが動くのに対し、fn で関数を宣言すると型宣言が強制され、全てのローカル変数や例外の明示的宣言が必須になる。書くのは面倒になると思うが、おそらくこれを使うとコンパイラは高速化がしやすくなるので、書きやすさと保守性や性能のバランスをユーザーが選択できることになる。

デプロイ容易性

素のPythonではダメだったモチベーションの一つに、モバイルやサーバー環境でのデプロイが困難であったことが挙げられている。サイズは何十MBにもなってしまうが、ランタイムを内包したバイナリを生成できるらしい。最近Node.jsがそれをやっている話が流れてきたが、これに関してはCPythonもやればできるのではという気もしないでもない。僕はpercolよりはpecoxkeysnailよりはxremapの方が便利だと思っているが、それはPythonを独立してインストールして維持しないといけないのが面倒くさいからで、シングルバイナリのデプロイが生活を少し楽にしてくれることはよくわかる。

なぜ速いのか

巷ではMojoがまるでC言語並みに速いPython処理系であるかのような受け取り方をしている人が見受けられるが、ローンチ動画を見た限りでは、単にPython処理系として使うとCPythonの数倍速い程度に留まり、C言語の性能には程遠い印象を受ける。しかし、何も書き変えずに数倍速くなるのはそもそも偉いし、Python風の高級言語にちょっと書き直すと35000倍速くなるのも普通にすごい話なので、それに貢献していると思われるポイントをまとめてみた。

MLIRとSIMD、並列化

端的にいうと、Mojoというのは高速なPythonなのではなく、MLIRという並列計算に長けたコンパイラ基盤を言語機能として表出させ、Python風の高級言語で簡単に扱えるようにすることでCやC++より速いコードが書ける言語、という感じに見える。マンデルブロットが35000倍速くなるのはMojoSIMDサポートを使っているからだとローンチ動画で説明されており、このSIMDMojoがMLIRを直接使えることで実現されているとドキュメントに書いてある。

CGOというコード最適化に関するカンファレンスの論文を眺めている時に、Chris Lattnerが書いたMLIRの論文を見かけたことがある。LLVMの作者である彼は、最近はLLVMよりMLIRの方に注力していると伺える。GoogleでMLIRを作ったが、MLIRを使いやすい言語がないから起業して言語作っちゃうぞ、ということだったらカッコイイなと思う*1

MLIRのMLはMachine LearningではなくMulti-Levelの略なのだが、MLIRには特定のアーキテクチャや環境に特化した最適化を可能にするdialectという概念があり、それを使ってSIMDやスレッドで並列計算することができ、LLVMと違ってMachine Learningといったドメイン固有の最適化も可能にしていると論文に書いてある。

動的な機能が制限されたstruct

次に重要そうだなと思ったのがstruct。言語機能的にはdataclassのようなものだし特段驚くようなものではないが、このstructを使うと動的なディスパッチやモンキーパッチが不可能になると書いてある。このstructを使った変数を宣言して、その変数に対してstructが実装していない処理をしようとするとコンパイル時エラーになる例が示されており、メソッドの呼び出し先が静的に解決できるのは最適化目線では非常に素晴らしいことだなと思う。

AOTとJITをサポート

僕はJIT屋さんなので、Mojoの話を聞いた時にPython部分はJITしてるのかなというところが最初に気になったが、FAQを見るとAOTとJITを両方サポートしていると書いてある。PythonのコードをそのままMojoで実行して8.59倍速になってるところは、MLIRの基盤の寄与が大きいのか言語特有のJITサポートががんばってるのかどっちなのかはわからないが、structではない(動的ディスパッチが可能な)既存のPythonライブラリのクラスをこねくり回すところでまともな性能を出したかったら、実行時情報に特化して最適化を行なうJITは必須であろうと思われる。CPythonもそろそろ本気を出して欲しい。

所感

MojoがCPythonの用途を全て置き換えていくかというとそうはならないと思うが、Pythonの既存資産が使えて、AI開発のために高速なコードが書けるというのはいいものだと思うし、何よりLLVMやSwiftを作ったChris Lattnerがやっているというのがアツいところなので、正式リリースに期待している。

言語の使用感としてはRubyの方が好きなので僕はRubyを速くして使っていくが、最適化のところで真似できるアイデアがあったら参考にしていきたい。

*1:実際にはどういうモチベーションで起業したのかは調べてない