概要

GPU向けネイティブソフトウェアを手がけるVectorWareは8月10日、Rustの標準ライブラリが提供するポータブルSIMD(core::simd)をGPU上でそのまま実行できるようにしたと発表した。これにより、開発者はCPUとGPUの両方で修正なしに動作する、アーキテクチャに依存しない高性能なコードを記述できるようになる。VectorWareはこれまでRustのスレッドやasync/awaitをGPU環境に実装してきた実績を持ち、チームにはRustコンパイラチームのメンバーも参加している。

技術的な仕組み

今回の実装の核となる洞察は、GPUの実行単位である「warp」(32レーンで構成される並列実行グループ)が、本質的には単なるベクターユニットとみなせるという点だ。SIMT(Single Instruction Multiple Thread)の実行モデルはSIMDと本質的に同じであるとし、たとえば32個のf32を含むSimd<f32, 32>は、warpの32レーンそれぞれに1要素ずつ配置される。

要素ごとの加算・乗算といった操作は、Rust標準のAddなどのトレイト実装を通じて自動的にwarp命令へマッピングされる。reduce_sum()のようなリダクション処理はGPUのwarp shuffle命令でレーン間のデータ交換を行い、simd_swizzle!によるクロスレーンシャッフルも同じシャッフルプリミティブを利用する。またMask<T, N>はwarpレーン内の述語を表現し、vote/ballot命令に対応させることで、ループカウンターのようなスカラー値は全レーンで同一に計算される「uniform」として扱われる。これはISPCなど既存のデータ並列言語の概念を、Rustの型システムに組み込んだ格好だ。

VectorWareはwarpを独自の小さな「マシン」と捉え、Rustのジェネリクスやconst generics、トレイト境界を活用して中間表現(IR)を実装した。このIRはインタープリターを介さず、手書きPTXと同等のコストでネイティブ命令へ直接コンパイルされる点が特徴で、同じIRはCPU上でも実行可能であり、決定的な参照実装インタープリターとしても機能する。

対応アーキテクチャと制限事項

現時点ではNVIDIA GPUを主なターゲットとしているが、AMDのwavefrontやVulkanのsubgroupも同様のプリミティブを公開しており、この仕組みはCUDA固有のものではないとされる。IRはアーキテクチャに依存しない設計を志向している。

一方で、ポータブルSIMD自体がRustではまだ不安定な機能であり、利用には#![feature(portable_simd)]が必要となる。またwarp幅より狭いベクターを使うとレーンが遊んでしまい、逆に幅の広いベクターではより多くの命令が必要になる。任意の順列によるシャッフルは複数命令やshared memoryへのアクセスを要する場合があるなど、パフォーマンス面での制約も残る。「抽象化のコストはゼロ」とうたう一方で、warp幅に一致しない場合はコストが増加するとも注記されている。

今後の展望

VectorWareは、スレッド(warp間の並列性)、core::simd(レーン内の並列性)、非同期処理(並行性)を組み合わせた複合的な抽象化や、行列形状のSIMDを扱うテンソルコア対応、スカラーループを自動的にSimd操作へ変換する自動ベクトル化、CPU/GPUで共通のベクター表現の確立などを今後の課題として挙げている。将来的には複数のプログラミング言語やランタイムへの対応も視野に入れているが、高性能かつ信頼性の高いGPUネイティブアプリケーションの構築においてはRustが最適だとの見解を示している。