Stable SIMD

2018/3/4のnightlyよりstdsimdプロジェクトの成果として、標準ライブラリとしてSIMDを提供する機能(stable SIMD)が実装されました。この記事ではStable SIMDを用いてRustでSIMDのコードを書く方法について解説します。

(2018/8/9追記) std::simdが消えたことによって不要になった説明を削除・修正


前回までのあらすじ

上述の記事が詳しいですが簡単にまとめると


  • RustでSIMDを使う方法は2つある


    • LLVMの最適化に任せる

    • intrinsics を使う



今回のStable SIMDは両方の用法に対して大きな変更点になります



  • #[target_feature]属性が追加され、関数単位でアーキテクチャを指定できるようになった


  • std::{simd,arch} module の導入により、直接intrinsicsを叩かなくて良くなった



    • std::simdは消えました(´・ω・`) (下の追記参照)




  • std::archが導入されてターゲット毎のintrinsicsをstableで使えるようになった


    • 以前は#![feature(platform_intrinsic)]が必要でnightlyでしか使えなかった




target_feature attribute

RustはMIRと呼ばれる中間表現を経由してLLVM IRにコンパイルされますが、LLVMによる自動ベクトル化機能により、単純なfor文等をSIMDを用いて計算することができます。しかしこの最適化を実行するにはSIMDの拡張命令を明示的に教える必要がありました。

RUSTFLAGS='-C target-feature=+avx' cargo run --release

これまではこの指定子 target-feature はコンパイル毎に指定されるため、複数のCPUで動くような実装をこのようにコンパイルすることは不可能でした。今回stable SIMDで追加された#[target_feature]属性はこの指定を関数単位で指定できるようにしたものです。

#[cfg(target_feature = "avx")]

fn foo() {
// implementation that can use `avx`
}

#[cfg(not(target_feature = "avx"))]
fn foo() {
// a fallback implementation
}

(RFC2325より)このように関数の属性としてSIMDの有効・無効を指定することができます。これにより関数fooはAVX命令を用いて最適化することが許可されるため、コンパイラにtarget-feature=+avxを渡した場合と同様のアセンブラが出力されることが期待できます。また同時にcfg!マクロの引数としても使えます:

if cfg!(target_feature = "avx") {

println!("this program was compiled with AVX support");
}

(RFC2325より)

加えて、実行時に検出することも出来ます:

if is_target_feature_detected!("sse4.1") {

println!("this cpu has sse4.1 features enabled!");
}

初回のis_target_feature_detected!呼び出し時にCPUを見て判定し、以降の呼び出し時にはコストがかからないようになります。これにより一つのバイナリに複数のCPU向けの最適化を同居させ、実行時に切り替えることが可能になります。


std::arch modules

RustはLLVMのintrinsicsを#![feature(platform_intrinsic)]経由で呼び出すことができるため、特定のSIMD命令を実行するためにこの機能を用いていました1simd crateや上述の記事にあるx86intrin crate はこれらをラップしたRust関数を提供していました。これはnightlyの機能なのでstableでは使えませんでした。

今回のStable SIMDでこの点が大きく整理されました。まず二つのモジュールがstdに追加されました:



  • simd: ポータブルなSIMD計算のための型定義 i32x4




  • arch: プラットフォーム固有な関数のラッパー _mm_setr_epi32等、および固有の型__m128i

これらは #![feature(stdsimd)] で有効化されます。これは #![feature(platform_intrinsics)]とは別個の機能で、恐らく比較的すぐに安定化されると期待されます2

Rust 1.27で#![feature(stdsimd)]は安定化されました。

Tracking issue for stable SIMD in Rust #48556

なお、AVX512は記事執筆段階でまだ実装中のため入っていません。

この構成を見ると、例えばstd::arch::x86_64::_mm_xor_si128のような関数は引数としてstd::simd::i32x4のような汎用なSIMD型を取れるのかな、と期待されますが残念ながらプラットフォーム固有のstd::arch::x86_64::__m128iを取ります。これについてはだいぶ議論があったようですが、非常に多数あるSIMD命令全てをこのように「型付きの」関数にするにはコストがかかり過ぎるという事で、引き続きstdsimdプロジェクト側で開発が続けられstable SIMDからは外れることになったようです。__m128iからi32x4に変換するには "either via transmutes or via explicit functions"と言っていますが、現状明示的な変換関数が見当たらないのでtransmuteを使う形になりそうです。


最後に

ちょうどモンテカルロシミュレーション用の高速な疑似乱数生成アルゴリズムであるSIMD-oriented Fast Mersenne Twister (SFMT)をRust/stdsimdで実装していたところだったので、勢いでstable SIMDに移植しました:

https://github.com/termoshtt/rust-sfmt

Nightlyでしか動作しませんが、XOR Shiftと同等の速度で周期 $2^{19937} -1$の疑似乱数が生成できます。


追記(2018/7/7)

2018/6/21リリースのRust 1.27よりstdsimdは安定化し、#[feature(stdsimd)]無しでコンパイルできるようになりました


追記(2018/8/9)

1.27でのstd::simdの安定化が無かったことになりました

std::simdとして安定化された部分はpacked_simdとして再度検討されるようです

std::simd互換のインターフェースを提供する目的でpacked_simd crateが出来ており、例えばrandservoもそちらを使っているようです。このRFCの安定化まではpacked_simdはnightlyでしか動きません。

ややこしいですが、stdsimd全体が無かったことになったわけでは無く、std::archは安定化されているのでターゲット固定のSIMD機能はstableで使用することができます。

余談ですが、Embedded-WGではRustでインラインアセンブラをstableで使うためにasm!機能を安定化させるのでなく、core::arch以下に実装していく方針のようです





  1. RustでCUDAカーネルを書く 



  2. SIMDの安定化はRust 2018 Roadmapでも Custom Allocator/Macros2.0 と共に言及されています。