前回の記事でRust版のComposableThreads（crates.ioにも公開しています）をリリースしたことを書きましたが、 その後も引き続き改良中なのでその過程を整理しておきます。

前回も書いていた通り、ベンチマークを取っていて rayon の join() が想像以上に速いと分かりました。 最初は、stackless (async) の実装をちゃんと詰めれば rayon の性能に追いつけるだろうと思っていたのですが、 asyncをチューニングするだけでは限界があると分かってきました。 よく考えると、同期的にタスクを一時中断する仕組みを作る以上、 タスクの状態を管理するための情報や、スケジューラに一度戻ることで発生する分岐のコストは避けようがありません。 async は「ゼロオーバーヘッド」と謳われていますが、この継続を取り出すコストは避けようがなく、 API を修正する以外で rayon の性能に追いつくのは難しいとわかりました。 そこで、この rayon の API をもう一つのモデルとして取り入れるという方向で実装を修正してみました。

  • stackful: 普通の（スタックフルな）ユーザーレベルスレッド。ミューテックスやバリアでブロックできる
  • stackless: async/awaitベース。Futureとしてポーリングで駆動するのでスタック確保がない
  • scoped: 継続を取り出さず、並列に呼びたい計算をクロージャとしてそのまま渡して呼ぶだけの、2分岐専用モデル。タスクディスクリプタもヒープ確保も無い、一番軽いモデル

このscopedモデル、調べてみると意外と HPC や Cilk 系の fork/join の研究の文脈ではちゃんと名前が付いていない気がしています。 継続を取り出さずにfork/joinするという書き方自体は、 よく考えるとTBBのparallel_invokeやOpenMPのparallel forのように世の中では普通に存在している方法ですが、 それを一つの独立したプログラミングモデルとして名指しした呼び方はなさそうでした。 今回 ComposableThreads の中では便宜上「scoped」という名前を割り当てました。 理由として、継続相当の処理をクロージャとしてスコープに収めているからです。

また、スレッドローカルストレージ（TLS）を使わずに「自分が今どのワーカーか」を特定できないか、という検証もしていました。 スタックポインタが特殊な領域に置かれていることを利用して、そこからワーカーを逆算して求めるというやり方です。 実装自体はできることを確認したのですが、実際にはTLSのコストはそこまで高くなく、 むしろ逆算のための計算コストの方が大きくなってしまい、少なくとも手元のmacOS環境では性能向上には繋がりませんでした。

ベンチマークも新しく取り直しました。マイクロベンチマークを手元のマシンで動かして、 stackful・stackless・dual (stackful/stackless両方使える)・scopedそれぞれ単体のシステムを、rayonのjoinと比較しています。

まず、stacklessとscopedでは明らかに性能差があることがわかります。 stackless の実装は私が AI と一緒にアセンブリを眺めながらあれこれチューニングしましたが、そこそこ限界に近くなってきています。 そして、なぜかscopedの結果がrayonより速くなったのですが、 この実装自体はAIにほぼそのまま書かせたもので中身を細かくは追えておらず、正直よくわかりません。 rayonが持っている何らかの抽象化の機構を省いてしまっている可能性もあるのですが、 いずれにせよ、プログラミングモデルに「継続を取り出さない」という制約を課すだけでここまで性能が変わるというのははっきり確認できたかと思います。 stacklessの場合はstackfulに必要なcontext switch等が不要で、コールスタックの管理もしなくていいので、stackfulよりは多少速くなります。 dualの場合はstackful/stacklessの分岐が入るので遅くなります。この辺は概ね想定通りです。

色々まだまだ試作品に近くて、インターフェイスがまだ固まっていないのですが、 もう少しチューニングを繰り返して使いやすいものにしていければと思います。