ハクソク

世界を動かす技術を、日本語で。

GPUにおける非同期/待機

97日前原文(vectorware.com)

概要

  • VectorWareが GPUネイティブRust async/await の実現を発表
  • 従来のGPU並列処理 と warp specialization による課題整理
  • JAX/Triton/CUDA Tile 等の先行事例と比較
  • RustのFuture/async/await が持つGPU並列性への適合性
  • 実装例・エグゼキュータ の説明と今後の展望

VectorWareによるGPUネイティブRust async/awaitの実現

  • VectorWareは GPUネイティブソフトウェア企業 として初のRust async/awaitのGPU実装を発表
  • RustのFutureトレイトとasync/await構文 をGPU上で動作可能にした世界初の事例
  • 開発者が お馴染みのRust抽象化 を使って高性能GPUアプリケーションを記述可能に
  • これにより 並列・非同期プログラミング がGPUでも安全かつ効率的に実現

従来のGPU並列処理モデルとwarp specialization

  • 従来のGPUプログラミングは データ並列処理 が中心
  • すべてのスレッドが同じ処理を異なるデータ領域で実行
  • 例:
    • data[thread_id] = data[thread_id] * 2; のような一括処理
  • warp specialization により、GPU内の一部が異なるタスクを同時並行
    • 例: warp 0がロード、warp 1が計算A、warp 2・3が計算Bを担当
  • この方法は 明示的なタスク並列性 を実現するが、 同期や競合管理が手動で困難

既存の高レベルGPUプログラミング事例

  • JAX: 計算グラフにより依存関係を管理し、最適化されたコード生成
    • Python DSLで記述、複数ハードウェア対応
  • Triton: ブロック単位で独立計算、Python DSLで記述
    • MLIRダイアレクトを経て最適化
  • CUDA Tile: ブロックに加えタイル単位でデータ依存性を明示
    • 既存言語（Python等）からTile IRへ変換しGPU実行
  • これらは 新しいパラダイムやエコシステム が必要で、既存コードやライブラリとの親和性が低い

高レベルアプローチの限界

  • 新しい記述法が 一部用途にしか適合せず、完全なアプリケーション移行は困難
  • 既存CPU/GPUライブラリの 直接再利用が困難
  • 理想 は「既存言語で明示的・構造的な並列性を記述し、既存コードとも容易に統合できる抽象化」

RustのFutureトレイトとasync/awaitの優位性

  • Futureトレイトとasync/await構文 は既存言語内で構造的並列性を表現
  • Futureは 完了していない計算 を表し、実行モデルやハードウェアに依存しない
  • pollメソッドでReady/Pendingを返す 最小限の仕様
  • コンパイラによる状態機械への自動変換 で、手動のタスク管理不要
  • Rustの 所有権モデル によりデータ依存性や共有制約を明示
  • JAX/Triton/CUDA Tileの利点 をRustの型・抽象化で自然に表現

GPU上でのasync/await実装例

  • CPUと同一構文 でGPUカーネル内async/awaitが利用可能
  • 例:
    • async fn async_double(x: i32) -> i32 { x * 2 }
    • let doubled = block_on(async_double(val)); など
  • 複数のasync関数、条件分岐、複数await、asyncブロック、サードパーティ製combinator（futures_util等）もサポート
  • NVIDIA ptxasのバグ など技術的課題も克服

GPU上でのエグゼキュータ設計

  • RustのFutureは 自動で実行されず、エグゼキュータで駆動
  • 最初は block_on によるシンプルな実装で検証
  • Embassy executor をGPU向けに移植
    • Rustのno_std環境向け設計でGPUにも適合
    • 既存OSSライブラリの再利用が容易
  • 複数の非同期タスクを同時スケジューリングする例も提示
    • 共有状態のカウンタを非同期にインクリメントし、ホストからの停止要求にも対応

今後の展望

  • Rust async/awaitモデルのGPU実装 により、従来のGPUプログラミングの複雑さを大幅に軽減
  • 既存Rustエコシステムやライブラリとの 高い親和性
  • 安全性・パフォーマンス・生産性 の全てを両立する新たなGPU開発スタイルの提案
  • VectorWareは今後も GPUネイティブソフトウェアの進化 を牽引予定

Hackerたちの意見

これを見るのはすごくクールだし、自分も興味を持ってこの分野を探求してるんだ。これとNVIDIAのstdexecの違いや共通点について、特にRustでFutureを使ってるっていう点以外で、何か知りたいな。

これの本当の利点がよくわからないんだけど、ワープが異種並列ワークロードを実行する際に、作業の奪取や継続の奪取ができるようになるってこと？でも、それには非同期関数の状態をGPU全体の共有メモリに保持する必要があるから、一般的には貴重なリソースだよね。

└

そう、その通り。データセンター用のカードや統一メモリを持つシステムでは、GPU全体のメモリはそれほど不足してないよ。ローカルな未来を持つローカルエグゼキュータを使って、!Sendのものをより速いアドレス空間に配置することもできる。

└

たくさんのGPUワークロードは、大量のRAMをGPUに常駐させて、CPUからの新しいデータで計算を行う必要があるんだ。

└

これはもうC++で起こってることなんだ。NVIDIAが送信者/受信者の提案を推進していて、これはC++標準ライブラリに追加される可能性のあるコルーチンランタイムの一つなんだよ。

└

神様、GPGPUやコンピュータシェーダーが流行り始めた頃にグラフィックスプログラムの選択科目を取った者として、これを読むと今のモダンGPUのアーキテクチャについてアップデートが必要だと実感するわ。異種ワークロードについてだけど、HPCの友達から、ワープ内での分岐を避ける古いアドバイスはもうあまり問題じゃないって聞いたんだけど、これって本当？

パフォーマンスはどうなの？ストリーミングマルチプロセッサのプログラミングモデルをこれに変えるメリットは何だろう？

└

まだパフォーマンスには焦点を当ててないんだ（それはワークロードやエグゼキュータに依存するし、投稿にもあるように今は非効率なポーリングをしてるからね）。でも、RustのFutureは状態機械にコンパイルされるから、コストゼロの抽象化なんだ。期待される利点は、CPUのasync/awaitの利点に似てるよ：並行コードを書く開発者にとってのエルゴノミクスが向上すること、共有・限られたリソースの利用が良くなること、並行性のバグが減ること。

本当にクールな実験だね（会社全体が）。トレーニングパイプラインは、まずCPUでデータ準備をしてからGPUに移動するデータでいっぱいで、CPUに何を残してGPUに何を置くか、テンソルを作る価値があるのか、それともタイルにすべきかを常に考えてる。君の会社はこういう問題を解決することに賭けてるんだろうね（例えば、GPUで直接推論リクエストを処理するためにはasync-awaitが必要だし）。私の質問は少し違ってて、SIMDの問題をどう扱いたいの？Rustの関数は、32の長い配列をデータ型として持つワープ上で動くべきなのか、それとも常に自動ベクトル化がうまくいくことを「期待」するべきなのか（特にRustのiterライブラリのヘルパーを使う場合）。

└

32幅の配列が良いかどうかもよく分からないな。AMDのワープは64幅だから、完全に自動ベクトル化には行かない方がいいと思う。

俺が心配してるのは、このasync/awaitのアプローチがTritonのように「AOT」じゃないってこと。コンパイル時にどの操作をするか分かってるから、どのワープで計算を効率的にスケジュールするかが分かるんだよね。ここではasync/awaitのアプローチだと、実行時に何が終わったか、何が終わってないかを手動で管理しないといけないみたいで、その後に新しい計算をどのワープに入れるかを考えないといけない。パフォーマンスに違いが出ると思う？

└

コンパイル時やAOTでの処理は、パフォーマンスに関してほぼ常に良いよね。async/awaitやfutureを使うことで、より複雑なプログラムが可能になって、以前はGPUでは簡単にできなかったことができるようになると思ってる。パフォーマンスよりも能力の問題だね（ただし、async/awaitのパフォーマンスがいくつかのケースでは良くなると信じてるけど、時間が経てば分かるだろうね）。

└

'async'計算が終了するかどうかを決定論的に判断できるなら、プログラムが正しくスケジュールされて、まだ計算されていない値への参照を避けるために、型システムのような静的解析を使える可能性が高い。でも、多くの場合、動的スケジューリングの方が好まれるんだよね。

ワープの特化は本当にひどいもので、排除されるべきだと思う。これが代替手段になるのは嬉しいけど、ブックキーピングのコストを最小限に抑えられるといいな。大きなカーネルのパフォーマンスに悪影響を与えるなら、AIでは普及しないと思う。

これ好きだわ。futureが協調的で、GPUに割り込みがないって言ってるけど、GPUのワープにはすでに命令レベルで先取りするハードウェアスケジューラがあるよね。あえてそのレイヤーの上で作業してるの？それとも、ワープスケジューリングにもっと直接的にフックして先取り的な動作を実現する将来のエグゼキュータの道が見えるの？

これ、rapids cudfについてずっと気になってたんだ。10年前にGPUネイティブのエンドツーエンド加速のためにgraphistryスタックを作ったけど、データの読み込み、整形、分析、強化、可視化など、サーバーGPUからクライアントGPUまで全部やってきた。でも、これが大きな障害になってる。小さなワークロードに対する大きな定常オーバーヘッドが避けられそうなのに。基本的には、Nvidiaのカーネルが最適化できるようにバルク指向でスタックを書く問題を解決したんだ。Apache Arrowや純粋なベクトル化されたデータフレームパイプラインみたいな感じ。ただ、cudfは「イager」で、ステップごとのCPU/GPUコントロールプレーンの調整が必要なんだよね。データプレーンはGPUにあるのに。理論的にはPolarsがレイジースケジューリングに移行して、よりバルクなGPU側のコントロールマクロステップのための最適化ができるけど、実際にはそうでもない。NvidiaがマルチテナントなどのフローのためにPythonのasyncioコストを削減しようとした試みも上手くいかなかったし。だから、ここでGPUにもっと移行できるようになるのはすごく興味深い。注目してるよ！

GPUは、ワープ内の全スレッドが同じ実行パスをたどる限り、かなり速くコードを実行するって知ってる。でも、分岐があるとパフォーマンスが落ちちゃうよね。複数のコルーチンに対して全く同じコードを実行するのは、実際にはほぼ不可能な気がするんだけど。そんなアプローチで本当に良いパフォーマンスが出るのかな？

└

短い答えは「はい」です。この投稿にはコルーチンの例が載ってるよ（Cスタイルを考えてみて：可能だけど見た目はイマイチ）。ここでの違いは、どれだけ簡単に書けるかってこと。実現できるかどうかじゃなくて、どのユースケースで使いやすいかが問題だと思う。

└

GPUの「スレッド」はSIMDレーンで、この投稿ではそれについては触れてないよ。複数のGPU「ワープ」、つまり別々の制御フローとコード実行を持つハードウェアスレッドを実行することについての話だね。（それを超えて、単一のコルーチンの複数のインスタンスで「同じコードを実行する」ことは、時々機会に応じて可能なはずだよ。）

とても興味深いね。「async trilema」に基づいたランタイムモデルは何なの？ https://without.boats/blog/the-scoped-task-trilemma/ 俺はtokioみたいな、つまりワークスティーリングだと思ってるんだけど。