Blog Dice

Everyone Should Know SIMD

Mitchell Hashimoto

誰もが知っておくべきSIMD

原文は Mitchell Hashimoto により に公開されました。 このブログを購読する

SIMDは複雑だというイメージを持たれている。私が出会った優秀なソフトウェアエンジニアの中にも、学ぶには複雑すぎる、あるいは最高レベルのパフォーマンスが求められるソフトウェアだけのためのニッチな最適化であり、日常的なプログラミングでは役に立たないと切り捨てる人が多くいる。

それは間違いだと思う。SIMDはシンプルに理解できる1し、素朴なforループを高速化するための典型的な「N個の値をまとめて処理する」SIMDコードは、ほぼ常に同じ基本形に従っている。基礎を覚えてしまえば、SIMDを書くのはforループを書くのとほとんど同じくらい簡単だ。そうでない場合は、今は見送るべきだというサインであることが多い。

すべての開発者は、少なくともその程度のSIMDは知っておくべきだ。

この記事では例としてZigを使っているが、内容はどのプログラミング言語にも当てはまる一般的なものだ。SIMD命令のサポート状況は言語によって異なるが、将来的にはより多くの言語がこうした汎用的な概念を公開してくれることを期待している！

毎回こんなことを書かなければならないのは本当に嫌なのだが、この記事はAIの支援を一切受けず、完全に手書きで書かれたものだということも付記しておきたい。

背景：SIMDとは何か

すでにSIMDが何か知っているなら、このセクションは飛ばしてほしい。

SIMDを使うと、CPUは複数の値を並列に操作できる。たとえば、1バイトずつ比較する代わりに、CPUは1つの命令で4バイト、8バイト、あるいはそれ以上のバイトを一度に比較できる。

もしコードの中にこんなループがあったら：

for (byte in bytes) { /* ... */ }
for (character in string) { /* ... */ }
for (value in array) { /* ... */ }

そこはSIMDを使うチャンスだ。SIMDはそれらをこう変える：

for (8 byte chunk in bytes) { /* ... */ }

これは並列度に直接対応した局所的な高速化をもたらす。データを4倍、8倍、あるいはそれ以上の速さで処理できるのだ。

これが効果を発揮するための本当の条件は、十分な量のバイトを日常的に処理していることだけだ。データが常に数バイトや数十バイトしかないようなforループであれば、やる価値はない。しかし、数百、数千、数百万バイトにわたって繰り返し処理しているのであれば、効果は絶大だ。

これが基本だ。simdutfやsimdjsonのようなプロジェクトはこれを極限まで推し進め、理解が難しいSIMDテクニックを使っている。しかし、SIMDの恩恵を受けるためにあのようなアルゴリズムを書く必要はない。一般的なケースは劇的にシンプルだ。

共通の形

よくある「N個の値をまとめて処理する」SIMDコードは、同じ5つのステップに従う：

  1. 必要な定数をブロードキャストし、必要であればベクターアキュムレータを初期化する。
  2. 入力をベクター幅ごとに1チャンクずつループする。
  3. すべてのレーンで並列に比較や演算を実行する。
  4. 必要に応じてベクターの結果をリデュースまたは格納する。
  5. 残りの要素をスカラーテールで処理する。スカラーテールとは、ベクター化する前の通常のループそのものだが、ベクターに収まりきらなかった余りだけを処理する。

これを繰り返すうちに、どんなforループも自然とこの5つのステップに分解できるようになり、SIMDを書くことはスカラーループを書くのとほぼ同じくらい自然になる。

実例

Ghosttyからの実例を見てみよう。スカラー実装、SIMD実装を見たあと、それを上記の共通の形にマッピングしてみる。

デコードされたコードポイントのスライスがあり、0xF以下の値（C0制御文字）が現れるまで消費したいと考えている。2ターミナルはほとんどが表示すべき通常の文字なので、それらをできるだけまとめてバッチ処理しようとしている。つまりこのループは、次の印字可能な連続領域の終端をできるだけ速く見つけるものだ。

スカラーループは1行だ：

while (end < cps.len and cps[end] > 0xF) end += 1;

一度に1つのコードポイントを処理する。理解するのは簡単だ。

こちらが、CPU固有のintrinsicを一切使わない汎用的なベクター版で、コメントも付けていない。3後で詳しく説明する。

if (simd.lanes(u32)) |lanes| {
    const V = @Vector(lanes, u32);
    const threshold: V = @splat(0xF);
    while (end + lanes <= cps.len) : (end += lanes) {
        const values: V = cps[end..][0..lanes].*;
        const greater_than_threshold = values > threshold;
        if (@reduce(.And, greater_than_threshold)) continue;
        const mask: std.meta.Int(.unsigned, lanes) = @bitCast(greater_than_threshold);
        end += @ctz(~mask);
        break;
    }
}

while (end < cps.len and cps[end] > 0xF) end += 1;

コードが12行増えた。

これにより、ループのスループットはARM NEON（Apple Siliconを含む）で最大4倍、AVX2（最近のほとんどのx86 CPU）で8倍、AVX-512（一部のIntel CPUやAMD Zen 4以降）で16倍まで向上する。

AVX2のIntelデスクトップで、ターミナルプログラムから最終的なターミナルの状態までの実際のエンドツーエンドのスループットでは、これは5倍程度の高速化だった。SIMDコード周辺の他の処理のせいで理想的な高速化からはどうしても目減りするが……それでも5倍だ！

さて、この12行が概念に馴染みのない人にはとても異質に見えることは理解している。そこで一度立ち戻って、先ほど触れた形に直接マッピングしながら、ステップバイステップで説明していこう。

ステップ1：定数をブロードキャストする

まず最初の3行から始めよう：

if (simd.lanes(u32)) |lanes| {
    const V = @Vector(lanes, u32);
    const threshold: V = @splat(0xF);

simd.lanes(u32)はGhosttyのヘルパーで、ターゲットCPUが一度に処理できるu32の数を返す。これらの個々の値はレーンと呼ばれる。ARMでは4、AVX2では8、AVX-512では16を返す。ターゲットが使いたいベクターサイズを持っていない場合はnullを返し、このコード全体をスキップしてSIMD処理は一切行わない。

@Vector(lanes, u32)はベクター型を作る。lanesが8なら、VはCPUが並列に操作できる8つのu32値を含む単一の値になる。以下同様だ。

最後に、すべての値を0xFと比較する必要がある。ベクター同士の比較では両辺がベクターである必要があるため、@splat(0xF)が0xFをすべてのレーンにコピー、つまりブロードキャストする。結果は次のようなベクターになる：

{ 0xF, 0xF, 0xF, 0xF, 0xF, 0xF, 0xF, 0xF }

これがステップ1だ。ベクター型を準備し、定数をブロードキャストする。一部のアルゴリズムではここでベクターアキュムレータも初期化するが、このアルゴリズムでは必要ない。

ステップ2：一度に1ベクターずつループする

次に、一度に1つの完全なベクターずつループする：

while (end + lanes <= cps.len) : (end += lanes) {
    const values: V = cps[end..][0..lanes].*;

lanesが8なら、少なくとも8つの値が残っているときにのみループに入る。ループ内では、その8つの値をベクターvaluesにロードする。ループの最後では、end += lanesで1つではなく8つ分進む。

完全なベクターであるという要件は重要だ。もし5つの値しか残っていなければ、8レーンのベクターをロードすることはできない。これを処理するには様々なトリックがあるが、ここでは簡単な方法を取り、スカラーテールで処理する。これについてはステップ5で後述する。

これがステップ2だ。入力をベクター幅ごとに1チャンクずつロードしてループする。ここでレーン数分の高速化が効いてくる！

ステップ3：SIMD演算を実行する

ここで比較を実行する：

const greater_than_threshold = values > threshold;

valuesもthresholdもベクターなので、これはベクター演算（文字通りのベクターCPU命令）にマッピングされる。1つの>が、valuesの各レーンをthresholdの対応する各レーンと比較する。8レーンあれば、これはスカラー比較のcps[end] > 0xFを8回実行するのと同等だが、それを1つのCPU命令でやってのける。4

結果はレーンごとに1つの真偽値を持つ別のベクターになる。概念的には、こんな感じだ：

values:                 { 0x41, 0x42, 0x43, 0x0A, 0x44, 0x45, 0x46, 0x47 }
threshold:              {  0xF,  0xF,  0xF,  0xF,  0xF,  0xF,  0xF,  0xF }
greater_than_threshold: { true, true, true, false, true, true, true, true }

これが実際のSIMD演算だ。明示的な内部ループはない。>演算子がすべてのレーンに並列に適用される。

比較は一例にすぎない。これは加算、乗算、最小値、最大値、あるいはベクター型がサポートする他のどんな演算でもあり得る。重要なのは、コードが依然として同じ形をしているということだ。

ステップ4：ベクターの結果をリデュースする

これで真偽値のベクターが得られたが、元のループが知りたいのは0xF以下の最初の値の位置だ。

まずは、すべての値が0xFより大きいという一般的なケースを処理しよう：

if (@reduce(.And, greater_than_threshold)) continue;

@reduce(.And, ...)はすべての真偽値をandで結合し、1つの真偽値を返す。すべてのレーンがtrueならcontinueして次のベクターを処理する。今回の例ではレーン3がfalseなので、@reduceはfalseを返し、どのレーンが失敗したかを正確に見つける処理へ進む。

いずれかのレーンがfalseなら、どのレーンが失敗したかを正確に見つける必要がある：

const mask: std.meta.Int(.unsigned, lanes) = @bitCast(greater_than_threshold);
end += @ctz(~mask);
break;

@bitCastは真偽値のベクターを、レーンごとに1ビットを持つ整数に変換する。ビットが1なら値は0xFより大きく、0ならそうでなかったことを意味する。マスクを反転させて失敗した比較を1にし、@ctzで最初の失敗までのゼロビットの数を数える。その数が、最初に失敗したレーンのインデックスになる。

そのインデックスをendに加算してbreakする。制御文字が見つかったからだ。

ステップ3と同じ値を使って、レーンごとのこの変換を見てみよう：

values:                 { 0x41, 0x42, 0x43, 0x0A, 0x44, 0x45, 0x46, 0x47 }
greater_than_threshold: { true, true, true, false, true, true, true, true }
mask:                   {    1,    1,    1,     0,    1,    1,    1,    1 }
~mask:                  {    0,    0,    0,     1,    0,    0,    0,    0 }

@ctz(~mask)は最初の1までのゼロビットを3つ数えるので、3を返す。endに3を加えると、最初の制御文字である0x0Aを含むレーン3を指すことになる。

これがステップ4だ。ベクターの結果を、元のアルゴリズムが必要とするものへリデュースする。このステップはアルゴリズム間で最もばらつきが大きい部分でもある。合計を求めるならベクターアキュムレータを1つの数値にリデュースするだろうし、変換ならベクター全体を出力バッファに格納するかもしれない。今回のスキャンでは、特定の1つのレーンを見つけられるようにベクターをビットマスクに変換している。

ステップ5：スカラーテールで仕上げる

ベクターループの後、最初に示したのとまったく同じスカラーループを実行する：

while (end < cps.len and cps[end] > 0xF) end += 1;

入力長がベクター幅のちょうど倍数でない場合、これが残りの値を処理する。たとえば、8レーンのベクターループでは、このループに0から7個の値が残る。これがスカラーテールと呼ばれるものだ。

このループは、simd.lanes(u32)がnullを返すCPUにも対応する。その場合はSIMDコードをすべてスキップし、スカラーループが入力全体を処理する。元の実装はフォールバックとテールの両方を兼ねているのだ。

これがステップ5だ。ただの通常のループである。

まとめ：共通の形

実装全体を5つのステップにマッピングしてみよう：

  1. @splat(0xF)が比較値をすべてのレーンにブロードキャストする。
  2. whileループが一度にlanes個の値をロードする。
  3. values > thresholdがすべてのレーンを並列に比較する。
  4. @reduce、@bitCast、@ctzが最初に失敗した比較を見つける。
  5. 元のスカラーループが余りと非対応CPUを処理する。

ステップ4の詳細は最初は理解に少し時間がかかるが、全体の形はシンプルだ。そしてステップ1、2、3、5は、まったく異なるアルゴリズム間でもほぼ同じ見た目になる傾向がある。

for (byte in bytes)を見かけたら、これがマッピングすべき形だ。

なぜコンパイラがやってくれないのか？

ときにはやってくれる！コンパイラは、特に複雑な制御フローのない規則的な算術ループなど、単純なループを自動ベクター化できる。手書きでSIMDを書く前に、必ずスカラー版を最適化付きでコンパイルし、コンパイラが何を生成するかを確認すべきだ。

しかし、コンパイラが自動ベクター化できる範囲は厳しく制限されており、一般的にはあまり得意ではない。自動ベクター化は何十年もコンパイラ研究の活発な分野だが、最近の研究でさえ、製品版のコンパイラが定期的にベクター化の機会を逃しているという観察から始まっている。これはすぐにはなくならない問題だと思う。

さらに重要なのは、このループが5倍の高速化を気にするほど重要であるなら、ベクター化は明示的で予測可能であってほしいということだ。無関係なコード変更やコンパイラのアップデートで、こっそりスカラーループに戻されてしまうのは避けたい。

誰もがSIMDを知るべきだ

すべての開発者がその機会を見分けられるようになるべきであり、そして何より、SIMDを恐れるべきではない。大量の連続したデータをスキャン、比較、カウント、あるいは変換するホットループを見かけたら、一度にベクター幅ごとのチャンクで処理することをイメージできるようになるべきだ。

この記事で示したように、こうした一般的なケースは非常に規則的なパターンに従っており、すぐに慣れることができる。そして言語のサポートが良ければ、アセンブリやCPU固有の癖を知らなくても簡単に改善が得られる。

誰もが、これができる程度にはSIMDを知っておくべきだ。5

脚注

  1. simdutfやsimdjsonのような非常に印象的なプロジェクトは、目的を達成するために極めて複雑なSIMDトリックを使っている。しかし、これは私が「日常的なSIMD」と考えるものではない。 ↩

  2. C0制御文字は0xFを超えて存在する。これはGhosttyがこの特定のコードパスで使っている閾値であり、ESCやその他の制御シーケンスの処理は別の場所で行われる。 ↩

  3. 汎用ベクターはCPU固有の構文を取り除くものであり、CPU固有のコード生成を取り除くものではない。Zigは依然としてこれらの操作をターゲットで有効になっている命令セットに落とし込む。Ghosttyはサポートされているベクター幅を選択できない場合、スカラーコードにフォールバックする。 ↩

  4. 比較自体は1つのベクター演算だ。ベクターのロード、結果のリデュース、失敗したレーンの特定には追加の命令が必要になる。重要なのは、複数の比較を同時に行っているということだ。 ↩

  5. この記事は私が書いたLobstersのコメントに基づいている。 ↩

この記事は「muse-spark-1.2-contributor」を使用して翻訳されました。

コメント