- Rustの
std::simdで作られた vb64 base64コーデックは、手続き的なループをそのままベクトル化するのではなく、データ配置と演算の流れを回路のように再設計してこそ、高速で移植性のあるSIMDコードになる
- 中核となる最適化は、分岐とメモリアクセスによるstall を減らすことにあり、比較・マスク・select・shuffleによって入力に依存せず同じ演算を行うbranchless構造を作る
- base64デコードでは、ASCII文字をsextetに変換するために
byte >> 4 と / 補正を使った perfect hash を作り、SIMDベクトル内のlookup tableとshuffleでoffsetを求める
- 4つの6ビットsextetを3バイトにパックする際は、laneを
u16に拡張してshiftした後、low/high byteを分離し、rotate_lanes_left とORで隣接laneのバイト断片を結合する
- ベンチマークでは、
-Zbuild-std、-Ctarget-cpu=native、N = 32 の組み合わせとremainderロード最適化の後、crates.ioのbaseline base64実装に対してほぼ全区間で 約2倍の性能 を示した
SIMDが必要となる物理的背景
- コンピュータの性能向上は理論的なCSだけでなく、物理的制約 と直接結び付いている
- Moore’s lawは2023年時点でもなお維持されているように見えるが、この15年で Dennard scaling の効果が崩れ、より高密度なトランジスタが電力消費密度の増加につながるようになった
- クロック周波数を継続的に上げることが難しくなった後、2000年代初頭から性能向上の主な方向は、より多くのコアを使うことへ移った
- マルチスレッディングではコア間の協調が必要なため同期コストが生じ、ジャンプ・仮想呼び出し・同期のような制御フローはstallを引き起こす
- stallの主な原因は2つある
- 分岐:
if、ループ、関数呼び出し、関数リターン、Cのswitchのような制御フロー
- メモリ操作: load/store、特にキャッシュフレンドリーでないアクセス
手続き的コードと命令レベル並列性
- 現代のCPUコアはコードを1行ずつ実行するのではなく、互いに依存しない演算を同時に発行する
a = x + y と b = x ^ y のように相互依存しない演算では、add回路とxor回路を同時に使える- この方式が 命令レベル並列性 であり、これを妨げる依存関係はdata hazardと呼ばれる
- CPUがfunctional unitをよりよく飽和させるほど、単位時間あたりに処理できる演算量は増える
- 分岐は次の命令を取り込む前に条件計算の完了を待つ必要があり、メモリ操作はデータが物理的にCPUまで到達しなければならないため、stallが発生する
- GPUは画像をベクトル形式のピクセルとして扱い、局所性の高い演算を多く実行するため、バッチ処理と制限された制御フローに合わせて設計されたSIMDマシンに近い
- SIMDは single instruction, multiple data であり、1つの命令が複数のデータlaneに対して並列演算を行う方式である
lane単位の考え方
- SIMDとvectorはしばしば同じ意味で使われ、SIMD命令の基本単位は固定サイズの数値配列であるvectorである
- vectorの各構成要素は lane と呼ばれる
- SIMDベクトルはレジスタに収まる必要があるため、たいてい小さい
- 例の環境における最大ベクトル幅は256ビット
- これは
u8x32の32バイト、またはf64x4のdouble 4個に相当する
- 小さなベクトルであっても、パイプライン飽和の負担を4分の1にできるなら、その分だけレイテンシ改善につながり得る
popcntに見る分割統治
- 最も単純なベクトル演算はbitwise and/or/xorである
- 通常の整数も、bitwise演算の観点では1ビットlaneのベクトルと見なせる
popcnt は整数内の1ビットの個数を数える演算であり、i32 を i1x32 と見ればreduce演算である- 32個のビットを配列として取り出して足し合わせる単純な実装は、質の悪いコードになりがちである
- より良い方法は、隣接するビット対を足し、さらに対の対を足すという形でlane幅を広げながら合計していくことだ
0x55555555、0xaaaaaaaa マスクで偶数/奇数ビットを分離- shiftでlaneを揃えた後に加算
- その後、2ビット、4ビット、8ビット、16ビット単位で繰り返す
- この実装は
popcnt命令には最適化されないが、そのような命令がないシステムでは小さく高速なコードになる
u64 にもreduction段階を1つ追加するだけで適用でき、64ビット全体の加算は不要である- このような 分割統治 アプローチはSIMDプログラミングの中核パターンである
SIMD命令セットの主要ツール
- 実際のSIMDベクトルはスカラーよりも複雑な意味論を提供し、特に遅い制御フローを置き換えるための機能が重要である
- 利用できる命令はアーキテクチャに大きく依存する
- x86の多くの高性能コアはAVX2を実装している
- AVX2は256ビットの
ymmベクトルを提供する
- レジスタ自体にはlane数という概念はなく、命令がlaneの解釈方法を決める
- たとえば
vpaddbはymmをi8x32として解釈する
- 一般的に利用できる演算は次のとおりである
- bitwise演算: lane幅は常に1ビットとして暗黙に扱われる
- lane-wise算術: 加算、減算、乗算、除算、整数shift、min/maxなど
- lane-wise比較:
m[i] = a[i] < b[i] のようなmask vectorを生成する
- select: maskを使って2つのベクトルのどちらから値を取るかをlaneごとに選ぶ
- shuffle/swizzle: 1つのベクトルをlookup tableのように見なし、index vectorでlaneを並べ替える
- mask vectorのtrue/falseは通常、all-onesまたはall-zerosのビットパターンを使う
- 比較とselectは、SIMDコードを branchless な状態に保つための中核ツールである
- branchlessコードは入力に関係なく同じ演算を行い、
x * 0 = 0、a ^ b ^ a = b のような性質で不要な結果を捨てる
shuffleでデータ位置を揃える
- shuffleはSIMDにおいて、データを「正しい位置」に持っていくための中核ツールである
- broadcastまたはsplatは、すべてのlaneが同じscalarを持つベクトルを作り、
[0, 0, ...] のindex shuffleで表現できる
- interleaveまたはzip/packは、2つのベクトル
a、b のlaneを交互に配置する
c = [a[0], b[0], a[1], b[1], ...]- shuffle2で実装できる
- deinterleaveまたはunzip/unpackはinterleaveの逆である
- rotateは
b[i] = a[(i + j) % n] の形でlaneを回転させ、これもshuffleである
- SIMDプログラミングでは、整数より大きいデータブロックをさまざまなサイズの小さなブロックとして再解釈し、再配置することが多い
intrinsics、target feature、portable SIMD
- SIMDで使用できる演算は、アーキテクチャや instruction set extension によって異なる
- x86 には ARM にない演算が存在する場合があり、Intel AVX-512 のように、同じベンダー内でも高性能サーバーチップにのみ提供される拡張もある
- ツールチェーンはこれらの拡張を target feature として一般化している
- Linux の
lscpu は CPU が認識する feature を表示する
- LLVM は feature 設定に応じて命令選択を変える
+avx2 がないと LLVM は ymm を使うコードを生成できない
-march=native や -Ctarget-cpu=native はビルドしたマシンに最適な良いコードを生成できるが、他のプロセッサへの移植性は低くなる可能性がある- ランタイム feature detection は、CPU がサポートする機能を確認してどの関数バージョンを呼び出すかを決める方式であり、暗号化ライブラリのようにさまざまなデバイスへ配布されるコードで使われる
- C++ の SIMD コードは通常
_mm256_cvtps_epu32 のような intrinsics を使う
- 特定の instruction set の低レベル演算を表す
- 必ずしも単一命令にマップされるわけではない
- コンパイラは統合、重複除去、命令選択の最適化を行える
- 複数の instruction set 向けに似たコードを繰り返し書くことになると、assembly と比べて保守性の面で大きな利点がない場合もある
- portable SIMD ライブラリは、ライブラリレベルで命令選択の一部を処理し、残りはコンパイラに任せるアプローチである
vb64 の実装は、Rust の portable SIMD が競争力のあるコードを生成できるかを確認するための実験である
base64 デコードを SIMD に置き換える
- base64 は任意のバイナリデータを ASCII にエンコードする方式である
- 入力バイト列をビットベクトルと見なし、6 ビット単位の chunk である sextet に分ける
- sextet の値は次の文字に対応する
0..25 → 'A'..'Z'26..51 → 'a'..'z'52..61 → '0'..'9'62 → +63 → /
- base64 には複数の変種があるが、複雑さの大部分は共通している
- 注意点は 2 つある
- base64 はバイト内のビットが big endian の形式である
- 入力長が 4 で割り切れないことがあり、原則としては
= padding で 4 の倍数に揃えるが、padding が正しくないメッセージも処理できる必要がある
- decoded length は
input / 4 * 3 に input % 4 に応じた余りの長さを加えて計算する
branchless に向けた基本リファクタリング
- 単純な base64 デコーダには複数の分岐がある
- 入力を chunk ごとに走査するループ
- chunk 内の byte ループ
- ASCII 文字ごとの
match
- エラー時の
return Err
decoded_len 内部の matchVec::extend_from_slice と allocator 呼び出しの可能性
- 最適化の指針は すべての分岐を取り除く ことである
decoded_len の match は、input % 4 の値 0, 1, 2, 3 を 0, 1, 1, 2 にマップする- これを
mod4 - mod4 / 2 に置き換えると branchless 版になる
- LLVM は元の
match を switch table に畳み込めるが、この領域では不要なメモリアクセスが性能を落とす
最も熱いループの分離
- SIMD の強みは、一度に多くのデータを処理してループを大きく unroll し、branchless に近づけられる点にある
- hot loop の目標は、最大 4 バイトを読み、最大 3 バイトのデコード結果を作り、さらに文法エラーの有無も知らせることである
- 利用できる事実は 3 つある
- 出力長は branchless な
decoded_len() で計算できる
- 不正な base64 は非常にまれな経路とみなし、エラー位置が必要なら後から再走査できる
- base64 では
A は 0 なので、切り詰められた chunk を A で padding しても値は変わらない
decode_hot() は 4 つの入力バイトを処理し、デコード結果と成功可否の bool を返す形に分離されるOption<[u8; 3]> の代わりに bool を別で返すと、後続の if !ok 分岐を取り除きやすい- SIMD 版では、
Simd<u8, 4> を入力として受け取り、出力も power-of-two の lane 数に合わせて Simd<u8, 4> とする
- 実際に必要な出力は 3 バイト
- 最後の lane は使用しない
ASCII を sextet に変換する方法
- ASCII 文字を sextet に変換する
match の大部分は byte - C の形で表せる
'A'..'Z' → byte - 'A''a'..'z' → byte - 'a' + 26'0'..'9' → byte - '0' + 52'+' → byte - '+' + 62'/' → byte - '/' + 63
- lane ごとの offset ベクトルを作って
ascii - offsets を行えばよい
- 最初のアプローチは compare-and-select である
A-Z、a-z、0-9、+、/ に対する mask を作る- どの mask も選ばれなかった lane は invalid と判断する
- 各 mask に対応する offset を splat し、OR で結合する
- この方式は洗練されていて競争力のあるコードを生成できるが、比較が合計 8 回必要で、生きている値が多く register pressure が発生する可能性がある
SIMD hash table と perfect hash
A-Z、a-z、0-9 の byte 範囲はそれぞれ 0x41..0x5b、0x61..0x7b、0x30..0x3a であり、high nibble が異なる+ と / は 0x2b、0x2f なので、byte >> 4 だけで大半を区別できる/ の場合に 1 つ引くと、範囲に対する perfect hash になる(byte >> 4) - (byte == '/') のマッピングは次の通り
A-Z → 4 または 5a-z → 6 または 70-9 → 3+ → 2/ → 1
- この値は小さいため、offset lookup table を SIMD ベクトルの中に入れ、shuffle で lookup できる
- この perfect hash のアイデアは GitHub issue の匿名ユーザーが提案した
Simd::swizzle_dyn() には、index 配列と lookup table の長さが同じでなければならないという制約がある- perfect hash 方式では、sextet 計算の過程で validation を副作用として得られないため、同じ GitHub issue の exact bloom filter を使って byte の妥当性を検査する
- 実装例は vb64 の simd.rs にある
4 つの sextet を 3 バイトにパックする
- 4 個の 6 ビット sextet を 3 バイトに結合する段階はさらに厄介である
- 特定の入力 sextet 1 つを all-ones にして、出力でビットがどこへ移動するかを確認すると、配置関係を追跡できる
- バイト単位の shuffle だけでは不十分である
- 移動対象がバイト片だからである
- shift だけでも足りない
- overshift されたビットが隣接 lane に移動しなければならない
- 解決策は lane をより大きくすること である
sextets を u16 ベクトルに cast してから lane ごとに shift する
input[0] は 2 ビット shiftinput[1] は 4 ビット shiftinput[2] は 6 ビット shiftinput[3] は 8 ビット shift に調整
- shift 結果から low byte と high byte のベクトルを分離する
hi.rotate_lanes_left::<1>() で high byte 側の断片を隣接 lane に合わせ、その後 lo | hi_rotated で結合する- この方式は hardware primitive を積極的に活用するため、コードが小さく効率的である
lane数の拡張と garbage lane の除去
Simd<u8, 4> は x86 の最小 128 ビットベクターレジスタよりも小さいため、decode_hot() を lane 数 N に対してジェネリックにするLaneCount<N>: SupportedLaneCount 制約によって、小さい 2 の累乗の lane 数を保証する- lookup table と shift table は
tiled() helper で繰り返しパターンのベクトルを作る
N = 4 では最後の lane の garbage 値を無視すればよかったが、N が大きくなると 4 つおきの lane に garbage が混ざる- これを除去するために shuffle を使う
- 望ましい関係は
shuffled[i] = output[i + i / 3]
- 4 番目のインデックスごとに飛ばして garbage lane を削除する
- オーバーフローする部分は最終出力ベクトルの上位 1/4 なので無視される
- こうすると
decode_hot::<32>() で 32 個の base64 byte を並列デコードできる
outer loop の最適化
decode() も内部 lane 数 N に対してジェネリックに変える- 残っているコストは次のとおり
for chunks in ... の長さ比較分岐[T]::copy_from_slice の可変長 memcpy- 各 loop iteration の
ok 分岐
Vec::extend_from_slice の allocator 呼び出しの可能性と、もう 1 回の memcpy
- 出力長が分かっているので
out.reserve(final_len + N / 4) であらかじめ領域を確保する
- さらに slop 領域を設けて、可変長 memcpy の代わりに完全な SIMD store を行う
- 各 iteration は SIMD ベクトル全体を書き込み、次の write は
3/4 * N だけ進んで前の garbage byte を上書きする
- 最後の garbage byte は最終
Vec::set_len() に含まれないため、削除されたものとして扱われる
if !ok のために early return しても、set_len() で commit していないので out は未変更の状態のまま残る
エラー処理を hot loop の外へ遅らせる
- 各 iteration ごとに
if !ok で return せず、error |= !ok で累積する
- 最終
set_len() の直前に 1 回だけエラー有無を確認する
- ほとんどの base64 blob は valid だという前提では、エラーパスは hot loop の外に押し出される
- 文法エラーがあっても、その後の SIMD 演算が任意に誤動作するわけではないため、garbage write は commit されず消える
- その後の
Vec::push() のような呼び出しが同じバッファ領域を上書きできる
unroll and jam と remainder 処理
copy_from_slice の可変長 memcpy を減らすために unroll and jam を適用する- ループを 2 つの部分に分ける
- hot vectorized loop: 常に長さ
N の入力だけを処理する
- cold remainder part:
i < N の入力を最大 1 回処理する
- Rust の
Iterator::chunks_exact() を使って hand-rolled な unroll-and-jam を実装する
- hot loop では
Simd::from_slice() を呼び出して単一の vector-sized load を行う
- bounds check はコンパイラが除去しやすい形になる
ベンチマークと手動ロード最適化
- ベンチマークは長さ 0 から約 200 または 500 バイトまでのメッセージをデコードし、crates.io の baseline base64 実装と比較する
- コンパイルオプションには
-Zbuild-std と -Ctarget-cpu=native を使う
- チューニングの結果
N = 32 が最も良く、hot loop の各 iteration で YMM レジスタ 1 本を使う
- 当初は baseline に勝っていたが、
data.len() % 32 と強く相関する heartbeat 形状の性能変動が現れた
- assembly を確認した後、
copy_from_slice がバイト単位の load loop として inline/unroll されたようだと判断した
Simd::gather_or() も試したが、より悪い assembly を生成したため使わなかった- 代わりに可変長データ向けの手動 loading 関数を書いた
- hot part では可能な限り大きい scalar load である
u128 load をループで行う
- LLVM は 16 バイト chunk を XMM load に落とし込む
- remainder では重なり合う
u64、u32、u8 load を使う
- 15 バイトを読むときは、
p から u64、p + 7 から u64 を読んで 1 バイトを重複させ、OR で結合する
- 4〜7 バイトでは重なり合う
u32 load を使う
- 1〜3 バイトでは
p、p + len/2、p + len - 1 から読んで一部の byte を重複ロードすることはあるが、分岐数を減らす
- 新しい loading code 適用後は variance が非常に小さくなり、baseline 比でほぼ全域で 2 倍の性能を示した
encoding と web-safe base64
- encoding 関数は
decode_hot() の演算を逆に行う encode_hot() を実装すればよい
- デコードで使った perfect hash は encoding には合わないため、新しい hash が必要になる
- encoder 周辺の loading/storing code も decoder と少し異なる
vb64 は効率的な encoding routine も実装している- web-safe base64 は
+ と / を - と _ に置き換える変種である
- web-safe base64 の perfect hash 構成はさらに厄介で、例として
(byte >> 4) - (byte == '_' ? '_' : 0) のような方式が必要になることがある
vb64 はまだ web-safe base64 をサポートしていない
結論
vb64 は重要なボトルネックを解消しようとするライブラリではなく、base64 decoding が実際にボトルネックになる場所も分からないと述べている- branchless code はしばしばやりすぎだが、コンパイラができることとできないことを理解する助けになる
- Rust の
std::simd は全体として良好で、優れたコードを生成する
- SIMD code をもっと単純にするために改善されてほしい rough edge はあるが、現時点の成果には満足していると評価している
- SIMD と性能最適化は多くのトリックとハードウェア知識を要する複雑な分野であり、そのかなりの部分は文書化されていない
1件のコメント
Hacker Newsのコメント
portable SIMDを実際に使っているのを見るのは興味深く、Zen 3システムでベンチマークを再現してみたところ、同じ速度向上が得られた
M1 MacBook Proでは入力長110バイトで性能向上が1.4倍から始まり、徐々に2倍まで上がっており、x86_64よりは低いものの目標は達成しているように見える
ただしコードを見ると、RustはSIMDやポインタ関連の作業、より広くは性能エンジニアリングにおいてエルゴノミクスがかなり悪いという自分の経験を裏付けている
それでもRustのportable SIMDはC++に比べるとまだ良い話ではなく、生バイト領域・ポインタ・バッファ操作のレイヤーに降りるには
PinやMaybeUninitなどに慣れる必要があるportable_simdとallocator_apiは何年も不安定なままで、参入障壁も高く、さらにぎこちないが、その大半は意図された設計であるただし、自分のプログラム内でより使いやすくする抽象化を自作したり、サードパーティ製クレートを使ったりすることを妨げるものはない
C++のSSE intrinsicはアンダースコアも見苦しく、名前も覚えにくいので、はるかに悪い
古典的なC++で最善を尽くして実装したのに、誰かがSIMD版で10倍以上速くしてくるのを見ると、ときどき本当に驚く
その代わり、このコードは移植性が低い
コンパイラの自動ベクトル化がもっと良くなってほしいし、一部の演算の並べ替えを局所的に許可する言語レベルの注釈のような支援もあってほしい
コンパイラがごく局所的な文脈の外ではデータを代わりに直せないため、自動ベクトル化は本当に難しくなる
例えば
for(double v: vec) sum+=vでは浮動小数点加算は結合法則が成り立たないため、値を順番に足す場合と、8個間隔で足してから残りを合算するSIMD方式は同じではないコンパイラから見ると明白な最適化のように見えても、特定の保証を緩めるよう指示されない限り、最適化よりも逐次意味論の保証を優先する
そのため面倒なことになり、janwasが言うようにホットパスにはライブラリ、特にGoogle HighwayやIntel ISPCのようなものを使うほうがよいと思う
できるだけ移植性を保って効率的であろうとしつつ、必要なときにはターゲット特化のプログラミングを容易にする
自動ベクトル化はFORTRANコンパイラのほうが明らかに優れているが、それはエイリアシングが許されないためである
C++はCのメモリモデルに従っているせいで足を引っ張られている
CUDAは今日の究極のSIMDマシンであるGPUのために設計されたC++であり、ROCmは事実上AMD向けCUDAに近い
個人的にはMicrosoftのC++AMPが好きだった。入門が最も簡単だったと思う
ただし結局定着しなかったのは残念だ
また、SIMDラッパーライブラリを使えば、実際にはかなり移植性のある形にできる
ちょっとした参考として、コンパイラはその
popcount実装を単一命令には最適化できなかったが、別の実装では可能であるもちろん、かなり扱いは難しい: https://godbolt.org/z/T69KxWWW8
_mm256_cvtps_epu32は特定の命令セットの低レベル演算を表すものだとされ、AVX2のfloat-to-intキャストだと説明されていたが、その命令はAVX-512に属するAVX2にはfloat-to-intキャストはなく、AVX1では整数結果がsignedで、命令は
_mm256_cvtps_epi32であるfastbase64[0]と比べるとどうなのか気になる
記事は素晴らしく、こういう内容をオンラインで見られるのはうれしいが、portable SIMDライブラリに対する著者の楽観論までは共有しにくい
[0]: https://github.com/lemire/fastbase64
SIMDをC++やRustに付け足すより、ISPCのほうが単純に優れていると思う
動的ディスパッチもサポートしており、自分で実装しようとするとつらい機能である
そうすればC++へインライン呼び出しを戻せるし、SIMDコードでテンプレートやクラスを使え、複数のSIMDコード領域をまとめてインライン化することもできる
動的ディスパッチの実装が難しい点には同意するが、Highwayがその部分を処理してくれる
素晴らしい記事で、「自分は絶対にここまで賢くなれないだろう」という感覚が強く残る
普通の人がソフトウェアエンジニアや物理学者ではないのと似ている
数か月集中して勉強すれば、同じようなレベルでできるはずだ
結局は興味と必要性の問題である
自分も性能最適化や、よりシステムに近いベアメタルエンジニアリングを個人プロジェクトで行き来して試しているが、仕事でもっと必要になればいいと思う
ただし、業界の仕事の大半が要求するのはそちらではない
慣用的なPythonではなく、すべてをnumpyで解くようなやり方である
面白いし、この種の賢さを学べる。記事の多くの部分は、そうした言語で問題を解く考え方ではごく自然に感じられる
時間が経つと、問題をそのような形で見るようになる
興味深い記事。
冒頭の最初の例では、ベクトル化されていない
popcnt実装が「正直、滑稽なほどひどいコード」を生成すると述べているが、release モードで native ターゲット CPU を使うと、この関数はかなりうまくベクトル化されているように見える。https://godbolt.org/z/WE1Eq65jY
pub fn popcnt(mut x: u32) -> u32 { x.count_ones() }これは
popcnt eax, edi; retにコンパイルされる。大きなビットベクトルでは、AVX2 実装が
POPCNTより速い場合がある。“Faster Population Counts Using AVX2 Instructions” を参照: https://academic.oup.com/comjnl/article/61/1/111/3852071
32ビットは十分に大きくなく、Rust が生成するコードは実際に滑稽なほどひどい。
popcnt命令に落とし込まれるべきだと思う。最近、ベクトル演算結果のマスクのビット数を数えるコードを書いたが、これは
popcntにうまく変換された。https://godbolt.org/z/zT9Whcnco
「これはひっかけ問題っぽいけど……単に add じゃないの？」というような部分があるため、普通は中間ベクトル表現を対象にして、細部はコンパイラに決めさせたくなる。
たとえば Haswell チップにはコアごとに複数の浮動小数点実行ユニットがあり、CPU はパイプライン化された浮動小数点演算を複数同時に実行できたが、そのうち
add命令は 1 つだけだった。直前の結果に依存しない加算が多く、レイテンシを避けられるなら、乗算項が 1 の融合積和命令も一緒に発行して、加算のスループットを 2 倍にできた。
その命令は通常のベクトル浮動小数点加算と同時に実行できた。