Everyone Should Know SIMD

Mitchell Hashimoto

誰もが知っておくべきSIMD

SIMDは複雑だというイメージがあります。非常に優秀なソフトウェアエンジニアの中にも、学ぶには複雑すぎる、あるいは最高水準のパフォーマンスが求められるソフトウェアのためだけのニッチな最適化で、日常的なプログラミングには役立たないと切り捨てる人を何人も見てきました。

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

すべての開発者が、少なくともそのくらいのSIMDは知っておくべきです。

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

毎回こう書かなければならないのは正直気が進みませんが、この記事はAIの支援を一切受けず、すべて手書きで執筆したことも付記しておきます。

背景:SIMDとは

SIMDが何かすでにご存知の方は、このセクションは読み飛ばしてください。

SIMDを使うと、CPUは複数の値を並列に処理できます。たとえば、1バイトずつ比較する代わりに、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倍、あるいはそれ以上の速さで処理できるようになります。

これが効果を発揮するための唯一の実質的な条件は、ある程度まとまった量のバイトを日常的に処理していることです。扱うデータが常に数バイトや数十バイト程度であれば、やる価値はありません。しかし、数百、数千、数百万バイトにわたってループしているのであれば、効果は絶大です。

以上が基本です。simdutfsimdjsonのようなプロジェクトでは、これを極限まで推し進め、理解が難しい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回に1つのコードポイントを処理します。とても分かりやすいコードです。

こちらが、CPU固有のイントリンシックを使わない汎用的なベクター版です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の個数を返します。この1つひとつの値をレーンと呼びます。ARMでは4、AVX2では8、AVX-512では16を返します。使いたいベクターサイズが対象CPUにない場合はnullを返し、このコード全体をスキップしてSIMD処理は一切行いません。

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

最後に、すべての値を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;

valuesthresholdもベクターなので、これはベクター演算(文字通りのベクター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なので、@reducefalseを返し、どのレーンが失敗したかを特定する処理に進みます。

いずれかのレーンが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を返します。end3を加えると、最初の制御文字である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の使いどころを見分けられるようになるべきですし、何よりもSIMDを恐れないことが大切です。大量の連続したデータをスキャン、比較、カウント、あるいは変換しているホットなループを見かけたら、ベクター幅1つ分ずつ処理する姿を思い描けるようになるべきです。

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

誰もが、このくらいはSIMDを知っておくべきです。5

脚注

  1. simdutfやsimdjsonのような非常に優れたプロジェクトでは、目的を達成するために極めて複雑なSIMDのテクニックが使われています。しかし、これは私が考える「日常的なSIMD」ではありません。

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

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

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

  5. この記事は、私が書いたLobstersのコメントを基にしています。

原文は Mitchell Hashimoto により に公開されました。

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