技術
約2分で読めます
Rustのasync/awaitがGPU上で動作 - VectorWareが初の実装を発表
VectorWareが、RustのFutureトレイトとasync/awaitをGPU上で動作させる初の実装を発表した。GPU上での並行処理に構造化されたプログラミングモデルを持ち込むという、かなり野心的な試みだ。
GPUプログラミングの現状の課題
GPUは伝統的にデータ並列性に特化している。より複雑なプログラムでは「warp specialization」を使って異なるワープが異なるタスクを並行実行するが、並行処理と同期の手動管理が必要でエラーが起きやすい。
JAX、Triton、CUDA Tileといった既存フレームワークは「計算グラフ」や「ブロック」の概念を導入しているが、新しいプログラミングパラダイムの習得が必要で、採用障壁が高い。
なぜRustのasync/awaitなのか
RustのFutureトレイトには、GPU上の並行処理に適した特性がある。
- 最小限の設計:
pollメソッドがReadyかPendingを返すだけのシンプルな設計で、異なる実行環境に移植しやすい - 遅延評価: プログラムを値として構築してから実行でき、コンパイラが依存関係を分析可能
- 所有権モデル: データの共有と転送が型レベルで保証され、GPU上でも安全な並行処理が実現される
実装の詳細
初期段階では単純なblock_onエグゼキューターを実装した。単一のFutureを繰り返しポーリングして完了させる素朴な実装だ。
その後、組み込みシステム向けに設計されたEmbassyエグゼキューターをGPU環境（no_std環境）に適応させた。変更は最小限で済んだとのことで、Embassyの設計の汎用性が活きた形だ。
デモでは以下が動作している。
- シンプルなasync関数
- 複数のawaitポイントを含む処理
- 条件分岐内でのasync
- asyncブロック
futures_utilライブラリとの組み合わせ- 3つの独立タスクの並行実行
現時点での制約
記事は以下の制限事項を正直に認めている。
- 協調的スケジューリング: Futureが譲歩しないと他の処理が飢餓状態になる
- 割り込みの欠如: GPUにはハードウェア割り込みがないため、スピンループが必要
- レジスタ圧力: スケジューリング状態管理がレジスタを消費し、占有率が低下する可能性がある
- 関数色分け問題: CPUと同様の制約がGPU上でも発生する
制約だらけだが、Futureトレイトのミニマルな設計がGPU上でもそのまま動くことを見せたのは面白い。AIワークロードが複雑化する中で、GPUの並行処理モデルを再考する実験としては意義がある。