CPU向けRust SIMDコードがそのままGPUで動く——「ワープはベクタユニット」という発想と実装の詳細

8月11日、VectorWareが「Rust SIMD on the GPU」と題した記事を公開した。RustのポータブルSIMD（core::simd）向けに書いたコードが、GPU固有のアノテーションを一切加えることなくGPU上でそのまま動作する——そんな成果を、同プロジェクトが実現したという。

CPU向けとGPU向けのコードを別々に保守するコストは、HeterogeneousコンピューティングがAI・数値計算の分野で当たり前になった今、多くのエンジニアにとって現実の痛点だ。VectorWareはRustのGPU対応コンパイラバックエンドを開発するオープンソースプロジェクトで、これまでにスレッドやasync/awaitのGPU動作も実現してきた。今回のSIMD対応は、その集大成ともいえるマイルストーンだ。

「SIMTはSIMDである」——概念的な突破口

この取り組みの核心は、シンプルな観察から始まる。

NVIDIAのGPUはSIMT（Single Instruction, Multiple Threads）モデルで動作する。1つのワープ（warp）が1つの命令を発行し、32本のレーン（lane）がそれぞれ自分のデータに対してその命令を実行する。「1つの命令が複数のデータに作用する」——これはSIMD（Single Instruction, Multiple Data）の定義そのものだ。

VectorWareはここに着目した。ワープは幅広のベクタユニットであり、ポータブルSIMDのベクタ型はそのユニットに直接マッピングできる、と。

用語補足

ワープ（warp）: NVIDIAのGPUにおける実行単位。32スレッドが束になって同一命令を実行する

SIMD: 1命令で複数データを並列処理するCPUの仕組み。x86-64のAVX2、ArmのNEONなどが代表例

ポータブルSIMDとは何か

従来、RustでSIMDを書くにはcore::archのアーキテクチャ固有なベンダー組み込み関数（intrinsics）を使う必要があった。x86-64なら_mm256_add_ps、Armならvaddq_f32という具合に、アーキテクチャごとに別実装が必要だった。

RustのポータブルSIMD（core::simd）はこれを抽象化する。Simd<T, N>という単一の型でN個のT型要素からなるベクタを表現し、算術演算・比較・リダクション・レーンシャッフルをひとつのコードで記述できる。コンパイラがターゲットのアーキテクチャに応じて適切な命令へ変換する仕組みだ。なお、ポータブルSIMDは現時点でもnightly限定機能（#![feature(portable_simd)]）であり、stable Rustへの安定化はまだ完了していない。

今回VectorWareはこのターゲットにGPUを追加した。しかもポータブルSIMDはstdではなくcoreに属しているため、同プロジェクトが別途整備したstdサポートすら不要で動作する。

実際のコード——変更なしで動く

以下が実際に動作したコードだ。GPU固有のアノテーションが一切ないのがポイントである。

#![feature(portable_simd)]
use core::simd::cmp::SimdPartialOrd;
use core::simd::num::SimdFloat;
use core::simd::{Select, Simd};

fn relu_dot(a: Simd<f32, 32>, b: Simd<f32, 32>) -> f32 {
    // 32要素の積を一度に計算
    let products = a * b;
    // レーンごとの比較でマスクを生成
    let positive = products.simd_gt(Simd::splat(0.0));
    // 負の値をゼロに置換
    let clamped = positive.select(products, Simd::splat(0.0));
    // 全レーンを横断してスカラーに集約
    clamped.reduce_sum()
}

fn main() {
    let a = Simd::<f32, 32>::splat(2.0);
    let b = Simd::<f32, 32>::from_array(std::array::from_fn(|i| i as f32 - 16.0));
    let result = relu_dot(a, b);
    println!("relu_dot = {result}");
}

このコードはCPU上でも同一の出力を返す。ソースを変えずにGPUとCPUの両方で動く、というのが今回の主張の核だ。

実装の内訳

各SIMD操作は次のようにGPUの命令へマッピングされる。

要素ごとの演算（加算・乗算・比較）: ワープが各レーンでネイティブに実行
リダクション（reduce_sum、reduce_max）: ワープのシャッフル命令でレーン間の値を交換・集約
クロスレーンシャッフル（simd_swizzle!等）: 同じくワープシャッフル命令へマッピング
マスク操作（select、any、all）: NVIDIAのvote/ballot命令を使用

コード中のスカラー値（ループカウンタや定数）は全レーンで同じ値を持つため、通常のCUDAと同様にワープ内で複製される。f32は「一様（uniform）」、Simd<f32, 32>は「レーンごとに異なる（varying）」という区別が、Rustの型システムそのもので表現されている点も興味深い。

正直なトレードオフ

記事はデメリットも明示している。

ポータブルSIMDはまだnightly限定（#![feature(portable_simd)]が必要）
ベクタ幅がワープ幅（NVIDIAは32）と一致する場合のみゼロコスト抽象。ベクタが小さければ一部レーンがアイドル、大きければ1操作が複数命令に分解される
任意のレーン置換（arbitrary permutation）は安価なシャッフル命令で表現できない場合があり、shared memoryを経由する可能性がある
リダクションやall/anyはワープ内の同期ポイントになる

また、他のRust機能との相互作用でコンパイラ自体を変更する必要があり、「未踏領域のため全ケースをカバーできている自信はない」と率直に述べている点は、実用化に向けた課題として念頭に置きたい。

今後の方向性

VectorWareはすでにGPU上でRustのスレッドとasync/awaitも動作させており、今回のSIMDでその並列性の階層が揃った。次のステップとして、これら3つを組み合わせた構成（スレッドがワープをまたいで処理を分散し、core::simdがワープ内のレーンに展開し、asyncが並行性を制御する）を目指すとしている。

さらに、行列形状のSIMDをNVIDIAのTensor Coreへ落とし込むことや、通常のスカラーRustループを自動ベクタ化してSIMD操作に変換することも視野に入れている。現状はNVIDIA向けだが、AMDのwavefrontやVulkanのsubgroupも同様のプリミティブを持つため、アーキテクチャを問わない拡張も検討中だ。

詳細はRust SIMD on the GPUを参照していただきたい。