Rustにおける安全なSIMD、内部でも

RustのSIMD抽象化は、私が望むほど安全ではありませんでした。これまでは。 生のSIMD組み込み関数（intrinsics）が使いにくいことは周知の事実です。 a + b と書きたいのに、この化け物のようなコードを書くことになります。 ```rust unsafe { #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))] _mm256_add_ps(a, b) #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "sse", not(target_feature = "avx2")))] _mm_add_ps(a, b) #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] vaddq_f32(a, b) } ``` 見てください。醜いでしょう。しかも全体が`unsafe`で囲まれています！ これは簡略化された例です。まだ以下のようなことを扱っていません。 * 他の一般的なプラットフォーム: AVX-512、32ビットARM、WebAssembly * SIMDがないプラットフォームやRISC-Vのような特殊なプラットフォーム * &[f32]のようなデータを各組み込み関数が受け付ける形式にロードすること * 実行中のCPUに最適な実装を選択すること 幸いなことに、Rustはこれらのすべてを処理してくれる多くのSIMD抽象化を提供しており、単純に a + b と書くだけで済みます。 ただ一つ問題があります。内部では、それらは依然として`unsafe`で満ちています。それは消えたのではなく、隠されていただけです。表面下には大量の`unsafe`が潜んでおり、時々問題を引き起こしていました。 あるいは、そうでした。これまでは。 なぜ「unsafe」が必要なのか？ 長い間、各組み込み関数（例: _mm256_add_ps）の呼び出しを`unsafe`で囲むことを避けることができませんでした。なぜなら、実行中のCPUで利用できない場合に呼び出すことは不正だからです。 そのため、各組み込み関数に必要な命令セットと、アクセス可能な命令セットを追跡し、それらを照合して、特定の関数を呼び出すことが安全かどうかを判断するメカニズムが必要でした。 これは手動で行うと面倒で、コードジェネレーターで行うと複雑で、常にエラーが発生しやすく、すべての組み込み関数の周りに`unsafe`が必要でした。 Rust 1.87でコンパイラが命令セット自体を追跡するようになったことで、これが変わりました。これにより、次のように書けるようになりました。 ```rust #[target_feature(enable = "avx2")] fn add_avx2(a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) // これはavx2の組み込み関数です } ``` 見て、`unsafe`なしです！ …まだです。 これでもまだ a + b とは書けません。できる最善のことはこれです。 ```rust unsafe { add_avx2(a, b) } ``` これは`unsafe`を一段階上にずらしただけです。正しい#[target_feature]でアノテートされた関数内で組み込み関数を呼び出すことはできるようになりましたが、チェーンのどこかに依然として`unsafe`が必要です。 もう一つの問題はより根本的です。+演算子の実装に#[target_feature]をアノテートすることはできません。なぜなら、+は常に利用可能でなければならないからです。したがって、このメカニズムを使っても a + b とは書けません。 補題: CPU機能トークン 最終的なソリューションがどのように機能するかを理解するには、まずCPU機能検出がどのように機能するかを理解する必要があります。 通常、AVX2のようなCPU機能をチェックするには、実行時にis_x86_feature_detected("avx2")を使用します。しかし、2つの数値を加算するたびにこのチェックを実行したくはありません。パフォーマンスが完全に低下してしまいます。一度だけチェックし、それ以降AVX2命令の使用が安全であることをコンパイラに証明したいのです。 代わりに、この証明を型システムにエンコードするために、偽造不可能なトークンを使用できます。これは、プライベートな内部フィールドを持つゼロサイズの型です。このトークンを取得する唯一の方法は、CPU機能チェックを実行する関数を呼び出すことです。チェックが成功すると、関数はそのトークンを渡します。 ```rust pub struct Avx2(()); fn detect_avx2() -> Option<Avx2> { if is_x86_feature_detected!("avx2") { Some(Avx2(())) } else { None } } ``` そして、これはゼロサイズの型なので、このトークンを渡しても実行時のオーバーヘッドはありません。それは純粋にコンパイル時の証明として存在します。 その結果、Avx2構造体のインスタンスを持っている限り、そのシステムでAVX2命令が利用可能であることを確信できます。 重要な洞察 コンパイラはそれを知りませんが、この関数は安全に呼び出すことができます。 ```rust #[target_feature(enable = "avx2")] fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } ``` この関数は、Avx2トークンを持っている場合にのみ呼び出すことができます。そして、Avx2トークンは、そのシステムでAVX2命令が利用可能な場合にのみ取得できます。 コンパイラにこれが有効であることを（`unsafe`を使用して）説明できれば、その`unsafe`を一度だけ記述して、どこでも再利用できます。 必要なのは、安全に呼び出せるマクロです。 ```rust with_avx2!( fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } ) ``` しかし、舞台裏では次のように展開されます。 ```rust fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { // SAFETY: トークンによるとAvx2が利用可能であり、 // 内部関数が `unsafe fn` でないことを確認しました。 unsafe { inner(token, a, b) } #[target_feature(enable = "avx2")] fn inner(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } } ``` これで、AVX2に含まれていない組み込み関数を使用すると、コンパイラが拒否します！ カスタムのターゲット機能追跡なしに、SIMD組み込み関数への安全なプログラミングインターフェースを提供することに成功しました！ 内部にはまだ1つの`unsafe`ブロックがありますが、それは健全なAPIにカプセル化されているため、メモリ安全性のバグを引き起こすために誤用することはできません。その意味で、println!のように`unsafe`コードを安全に抽象化しています。 このように、数百、数千ものカスタム`unsafe`ブロックではなく、この1つのマクロをレビューして監査するだけで済みます。そして、実装で誤って設定する可能性のある唯一のことは次のとおりです。 * トークンを間違った#[target_feature]にマッピングする * `unsafe fn` が安全なコンテキストから呼び出されるのを許可する そして、これらの失敗モードはどちらも非常に簡単にチェックできます。 これで、`unsafe`なしで add_avx2(token, a, b) を呼び出すことができますが、それでも a + b には到達しません。どうすれば解決できるでしょうか？ ジェネリクスが救世主 +演算子の実装に#[target_feature]をアノテートすることはできません。なぜなら、どこからでも安全に呼び出せる必要があるからです。また、関数にトークンを渡すこともできません。なぜなら、それはaとbは受け取りますが、トークンは受け取らないからです。 しかし、たとえそれができたとしても、それはかなり醜いAPIになります。a + b が常に機能し、ユーザーがトークンをいじることなく、最適なSIMD命令を自動的に使用できるようにしたいのです。 ジェネリクスを使用して、これら2つの問題を同時に解決できます。利用可能な命令セットに対してジェネリックなf32x8型を定義することで、トークンを内部に密かに埋め込み、SIMD命令セットごとに個別の実装を作成する加算を実装できます！ これがその様子です。 ```rust pub trait Level {} #[derive(Clone, Copy)] pub struct Avx2(()); impl Level for Avx2 {} pub struct f32x8<L: Level> { // 簡単のため、この例では配列でバックします。 // 本番コードでは、レベルごとにネイティブSIMD型を使用します。 data: [f32; 8], // 密かに埋め込まれたトークン！ token: L, } /// Avx2に対する`a + b`の実装 impl std::ops::Add for f32x8<Avx2> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { // (型変換は省略) // Avx2トークンを使用して安全なラッパーを呼び出します let result = add_avx2(self.token, self, rhs); Self { data: store_m256(result), token: self.token } } } ``` そして、他の命令セットや、SIMDが利用できない場合にも、同様に簡単に機能させることができます。 ```rust #[derive(Clone, Copy)] pub struct NoSimd(()); impl Level for NoSimd {} /// SIMDが利用できない場合の`a + b`の実装 impl std::ops::Add for f32x8<NoSimd> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { let result = std::array::from_fn(|i| self.data[i] + rhs.data[i]); Self { data: result, token: self.token } } } ``` これで、安全性と実行時の命令選択を同時に解決しました！ システムで利用可能な最適なLevelを提供する便利な関数を追加すれば、SIMDにほぼ完璧なAPIが得られます！ ABIが口を挟む 残念ながら、a + b を書いてSIMD命令に低下させることには根本的な問題があります。関数呼び出しのオーバーヘッドです。 関数呼び出しは無料ではありませんが、非常に安価です。ほんの数命令です。しかし、数命令というのは、加算の実装に使用した1命令よりもはるかに多い数です！ そのため、関数呼び出しが途中にあると、加算のパフォーマンスは急落します。そして、パフォーマンスこそがSIMDを使用する最初の目的です！ コンパイラは