GPU上でのRust SIMD

VectorWareでは、初のGPUネイティブソフトウェア企業を構築しています。本日、RustのポータブルSIMD（core::simd）をGPU上で成功裏に使用できるようになったことを発表できることを嬉しく思います。このマイルストーンは、開発者が使い慣れたRustの抽象化を使用して、GPUハードウェアの全能力を活用する複雑で高性能なアプリケーションを記述できるようにするという私たちのビジョンに向けた重要な一歩となります。 スレッド以下の並列処理 以前、GPUにRustスレッドを導入した際には、各std::threadをGPUのワープにマッピングしました。これにより、GPU上で多数の同時実行スレッドを実行できましたが、各スレッド/ワープ内の並列レーンは使用していませんでした。CPUでは、スレッド内の並列処理の抽象化はSIMDです。単一の命令がベクトルユニットにパックされた複数のデータ要素に作用します。スカラーコードが2つの数値を加算するのに対し、SIMD加算は例えば8つのf32値の2つのベクトルを取り、一度に8つの合計を生成します。このデータ並列性は、オペレーティングシステムが何かをスケジューリングするレベルよりも下の、単一スレッド内で行われます。 CPUスレッド SIMDオペレーション 012N⋯ SIMDレーン CPUスレッド RustのポータブルSIMD 歴史的に、RustでSIMDを記述することは、x86-64の_mm256_add_psやArmのvaddq_f32のような、core::archのアーキテクチャ固有のベンダー組み込み関数に頼ることを意味していました。これらの組み込み関数は単一の命令セットに固有であるため、複数のアーキテクチャで実行されるプログラムには、それぞれ個別の実装が必要でした。RustのポータブルSIMDは、これらの組み込み関数の上に抽象化レイヤーを追加します。それは、T型のN個の要素のベクトルを表す単一の汎用型Simd<T, N>を提供します。プログラムは、算術演算、比較、削減、レーンシャッフルをSimdに対して一度記述し、コンパイラはそれをターゲットCPUが持つ任意のベクトル命令に低下させます。VectorWareでは、GPUはポータブルSIMDがターゲットとするもう一つのベクトルハードウェアに過ぎないと認識しました。ボーナスとして、ポータブルSIMDはstdではなくcoreにあり、GPUに導入したstdサポートさえ必要としません。 SIMTはSIMDである GPUはNVIDIAがSIMT、つまりSingle Instruction, Multiple Threadと呼ぶモデルで実行されます。ワープは1つの命令を発行し、その32レーンのそれぞれが自身のデータに対してその命令を実行します。多くのデータ要素に作用する単一の命令は、まさにSIMDが意味するところであり、SIMTが追加するレーンごとのアドレッシングはその意味を変えません。ワープは幅広のベクトルユニットであり、ポータブルSIMDベクトルは直接そのユニットにマッピングされます。 CPUスレッド 012N⋯ SIMDレーン ≈ GPUワープ 012N⋯ ワープレーン 例えば、Simd<i16, 32>は、ワープの32レーンそれぞれに1つのi16要素を与え、そのような2つのベクトルの加算は、すべてのレーンが一度に要素を加算する単一のワープ命令にコンパイルされます。 CPU let a: Simd<i16, 32> = [1, 1, 1, ..., 1]; let b: Simd<i16, 32> = [2, 2, 2, ..., 2]; let c = a + b; compiles to vpaddw %zmm2, %zmm1, %zmm0 a0+b0レーン 0 a1+b1レーン 1 a2+b2レーン 2 a31+b31レーン 31⋯ println!("{c:?}"); GPU let a: Simd<i16, 32> = [1, 1, 1, ..., 1]; let b: Simd<i16, 32> = [2, 2, 2, ..., 2]; let c = a + b; compiles to add.s16 %rs3, %rs1, %rs2; a0+b0レーン 0 a1+b1レーン 1 a2+b2レーン 2 a31+b31レーン 31⋯ println!("{c:?}"); この新しいマッピングは、以前の作業からの並列処理階層を完成させます。CPUでは、スレッドはSIMDレーンを含み、GPUでは私たちのstd::threadは、そのハードウェアレーンが同じ役割を果たすワープです。どちらの場合も、core::simdがそれらのレーンを駆動します。 CPU ⋯ スレッド 0 012N⋯ スレッド 1 012N⋯ スレッド N 012N⋯ SIMDレーン ≈ GPU ⋯ ワープ 0 012N⋯ ワープ 1 012N⋯ ワープ N 012N⋯ ワープレーン 世界初：GPU上でのcore::simd 以前の投稿と同様に、コードは通常のRustであるため、これを視覚的に示すのは困難です。x86-64 SIMDに低下する同じcore::simd型がラップトップ上で動作し、ソースを変更せずにGPU上のワープ操作に低下します。ここでは、小さなポータブルSIMDルーチンを定義し、mainから呼び出します。これは、モデルのコア機能、つまり要素ごとの算術演算、レーンマスクを生成する比較、そのマスクによって駆動される選択、およびレーン全体にわたる水平削減を実行します。 #![feature(portable_simd)] use core::simd::cmp::SimdPartialOrd; use core::simd::num::SimdFloat; use core::simd::{Select, Simd}; // ポータブルSIMD。この正確な関数はCPU上でもコンパイルおよび実行され、 // ターゲットに応じてx86-64、Arm、またはスカラーコードに低下します。 fn relu_dot(a: Simd<f32, 32>, b: Simd<f32, 32>) -> f32 { // 要素ごとの乗算：一度に32個の積が計算されます。 let products = a * b; // レーンごとの比較は、レーンごとに1つのブール値を持つマスクを生成します。 let positive = products.simd_gt(Simd::splat(0.0)); // 正の積を保持し、残りをゼロに置き換えます。 let clamped = positive.select(products, Simd::splat(0.0)); // すべてのレーンを単一のスカラーに削減して合計します。 clamped.reduce_sum() } fn main() { // 通常のRustで構築された2つの32幅ベクトル。 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)); // 要素ごとのオペレーション、比較マスク、選択、および削減： // すべて通常のポータブルSIMDであり、すべてGPU上で実行されます。 let result = relu_dot(a, b); // 私たちのstdサポートを使用してGPUから印刷されます。 println!("relu_dot = {result}"); } エントリポイントは、GPU固有のアノテーションのない通常のfn mainです。私たちのツールチェーンはそれをGPUカーネルにコンパイルし、結果はデバイスから私たちのstdサポートを使用して印刷されます。以下は、GPU上でプログラムを実行し、CPU上で実行した場合とまったく同じ出力を生成した記録です。 実装 前述のように、マッピングは単一の観察に基づいています。ワープは、レーンが個別にアドレッシング可能なベクトルユニットであるということです。Simd<T, N>がレーンごとにレイアウトされると、各オペレーションファミリーは直接ワープレベルの対応物を持つことになります。SIMD要素ごとのオペレーションは簡単なケースです。加算、乗算、比較、その他のレーンごとの演算子は、Simd上の通常のRustトレイト実装（Addなど）から取得されます。GPUはそれらをネイティブに実行します。reduce_sumやreduce_maxのようなSIMD削減は、すべてのレーンを単一のスカラーに結合します。これらはGPUのワープシャッフル命令を使用してレーン間で値を交換および結合し、すべてのレーンで同じスカラー結果を生成します。simd_swizzle!やrotateのようなSIMDクロスレーンシャッフルは、レーン間で要素を移動します。SIMDレーンはGPUワープレーンであるため、これらはGPUレーンがデータを交換するのに非常に適しているのと同じワープシャッフルプリミティブにマッピングされます。SIMDマスクも同様にきれいにマッピングされます。Mask<T, N>は各SIMDレーンに1つの述語を与えます。Mask::selectは、すべてのワープレーンで選択を実行します。anyやallのような水平マスククエリは、GPUのvoteおよびballot命令を使用します。ループカウンターや定数のような周囲のコードのスカラー値は、すべてのレーンによって同一に計算されるため、通常のCUDAのようにワープ全体に単純に複製されます。これは、ISPCのようなデータ並列言語が明示するのと同じユニフォーム対バリアントの区別ですが、ここではRust自体の型から派生しています。プレーンなf32はユニフォームであり、Simd<f32, 32>はバリアントです。 レーンでの作業 抽象化とハードウェアが一致しない唯一の場所は、レーン数です。CPU上ではSimd<T, N>は1から64までの任意のNを許可しますが、GPUハードウェアには固定幅があります。NVIDIAでは32レーン、AMDでは32または64レーンです。マッピングは、Nがその幅と一致する場合にのみ1対1になります。より小さいNは一部のレーンをアイドル状態にし、より大きいNは一部またはすべてのレーンに複数の要素を処理させます。ワープよりも多くの作業がある場合、どのレーンが何をするかを指定する方法が必要です。ワープをそれ自体の小さな「マシン」と考えるのに役立ちます。レーン間でデータを移動および結合するための固定セットのプリミティブと、どのレーンがアクティブで、それぞれがどのくらいのデータを保持しているかについての不変条件があります。「プログラミング」することは、それらのルール内でレーンに作業を配置することを意味します。VectorWareでは、そのマシンにIRを提供します。スタンドアロンのデータ構造ではなく、型、ジェネリクス、定数ジェネリクス、およびトレイト境界を使用してRustの型システムにエンコードします。プログラムは、投票、シャッフル、削減、スキャン、ギャザー、スキャッター、アトミック、およびワープよりも広いベクトル用のストリップマイニングなどの型付けされたオペレーションで構成されます。オペランド、実行形状、および容量も型付けされています。なぜなら