Rustで書かれたスレッドライブラリをリリースしました

Rustでいわゆるユーザーレベルスレッドを扱えるライブラリとして、ComposableThreads（のRust版）をリリースしました。

今回は初めて crates.io にもパッケージを公開してみました。

ComposableThreads は元々、私が博士課程時代に C++ で書いたライブラリで、博論のメインテーマだった HPC 用途の分散共有メモリ (DSM) を作るため、片手間でスレッドライブラリを作ったのが始まりです。 DSM を作ろうとすると、アプリケーション・システムの双方にマルチスレッドが深く関与してきます。学生時代に所属していた田浦研究室では MassiveThreads というライブラリを長年開発していて、 Cilk 的な fork/join プログラムを高速実行しつつ、システム的にはスタックフルなユーザーレベルスレッドとして一通りの機能を備えていて、業界ではそれなりに有名な使いやすいライブラリです。普通に考えたら MassiveThreads を使えばよかっただけなのですが、どうしても DSM のような変なライブラリを開発しようとすると MassiveThreads だけではカスタマイズ性がネックになります。そこで、MassiveThreads のコードをきちんと理解した上でそれと遜色ない性能を出しつつ、各種ライブラリのプロトタイピングにも適したカスタマイズ性を備えたライブラリを作ろうということで、当時 ComposableThreads というライブラリを作りました。元々は DSM のおまけ的な形で作ったものだったのですが、結果的には論文にもなりましたし、博士時代までに作ったものの中で一番出来が良かったもののように思います。

最近少しまとまった時間ができたのと、AIでコーディングするのも簡単にできるようになったという事情もあり、以前からやりたかった ComposableThreads の Rust 化をやってみました。基本的なアイディアは C++ 版と変わらず、スタックフルなスレッドを一通り実装すること、一度だけ消費できる継続 (suspended threads) を第一級オブジェクトにすること、ジェネリクスによってコンポーネントをユーザーが差し替えられるようにすること、そして性能面では MassiveThreads に近づけることです。加えて、Rustだとasync/awaitによる非同期プログラミングも一般的なので、Wakerについてのサポートも入れることにしました。これは若干の分岐のオーバーヘッドにもつながっているので、ここも差し替えられるようにすべきとは思っています。

ベンチマークを取るにあたって、Rustで同じようなことをやっているライブラリを探したのですが、意外とないことに気づきました。スレッドライブラリで有名なのは Tokio ですが、主にI/Oの高速化に特化しているらしく、HPC用途には厳しい印象です。見つけた中で一番性能が良かったのはrayonでしたが、 rayonはデータ並列に特化しているようで、スタックフルなスレッドを実装しているわけではなく、例えばmutexやyieldのようなものはありません。

上図が手元の MacBook Pro で ComposableThreads, MassiveThreads, rayon の3つの実行時間を測ったグラフです。なぜか私の環境では MassiveThreads の性能があまり出なかった印象で、詳しく検証できていないのですが、私が研究室にいた当時はほぼ x86-64 のマシンばかりだったので、ARM64にあまり最適化されていないのかもしれません。 rayonはタスク生成がかなり速く、spawn/join 1回あたり 10 ns ぐらいでした。ただし、rayonはそもそも実装しているものが全然違うことには注意が必要です。 ComposableThreads は rayon と比べるとあんまり速くなさそうに見えますが、 spawn/join 1回で 50 ns ぐらいだったので、ユーザーレベルスレッドとして機能させるオーバーヘッドとしてはそこそこ削れているんじゃないかと思います。もっとも、コードは主に AI に書かせている部分が多く、私もアセンブリを眺めるところまでしていないので、もう少し限界までチューニングできるかもしれません。

今後はちゃんとコア数の多いマシンで計測したり、ライブラリとして成熟させるためにインターフェイスを整えたり、また暇を見つけて色々やりたいとは思っています。