データ並列カーネルのための小さなコンパイラ

多くの高速なコードは、退屈なループとして始まります。現代のハードウェアは、同じ操作を一度に複数の値に対して実行できます（例: SIMDやSIMT）。私たちは直接これらの実行モデルのためにコードを書くこともありますが、多くの場合、コンパイラが通常のコードを受け取り、複数のループイテレーションを一緒に実行できるように書き換えます。私は、この変換がどのようなものかを理解するために、小さなコンパイラ（Pythonで約180行）を作成しました。 私のコンパイラはカーネルを「ロウワー」（データ並列性が可視化される、よりシンプルで明示的な形式に書き換え）します。入力は手書きの小さなASTで、出力はPython風のコードとして出力されるロウワーされたIRです。ソースコードから命令まで全てを網羅するのではなく、このコンパイラはより大きなコンパイラの中間ステップとして考えてください。 例を見てみましょう。オーディオのスケーリングは並列化が容易ですが、依然として以下のような明示的に並列ではないコードを書くことが一般的です。 kernel scale_audio(samples, out, n, volume): for i in range(n): out[i] = samples[i] * volume 私のコンパイラはこれを以下のように変換します。 kernel scale_audio(samples, out, n, volume): vector_for base in range(0, n, LANES): let i = (base + lane_id) let active = (i < n) masked_store(out, i, (masked_load(samples, i, active) * volume), active) 目標は、`for`ループを`vector_for`ループに置き換えることです。これにより、ループの複数のイテレーションを並列に実行できます。このグループ化された実行内の各位置は「レーン」と呼ばれます。 レーンとマスク レーンは、グループ化された実行内の一つの独立した要素位置です。例えば、グループ化された操作が一度に4つの値を処理する場合、4つのレーンを持ちます。 [ 10 | 20 | 30 | 40 ] ^ | lane0 ^ | lane1 ^ | lane2 ^ | lane3 各スロットがレーンであるため、グループ全体で乗算のような操作を実行すると、次のようになります。 [10 | 20 | 15 | 30] * [ 3 | 3 | 3 | 3 ] = [30 | 60 | 45 | 90] 各レーンでコードが実行されるとき、どのデータを操作するかを知るための一意のオフセットを持ちます。各レーンは異なる論理要素を処理し、グループ化された操作はそれらを並行して実行します。 データがレーンスロット数で割り切れないケースを処理するために、範囲外のロードとストアをスキップするために、レーンごとのブールマスクが適用されます。マスクを`[true, true, false, false]`のように考えることができます。最初の2つのレーンは実行が許可され、最後の2つは無視されます。これは配列の最後のチャンクで、すべてのレーンを満たすのに十分な要素が残っていない場合に役立ちます。 次のステップでは、`vector_for`の本体を特定のアーキテクチャをターゲットとする命令に変換します。理想的なケースでは、これによりパフォーマンスをレーン数倍に向上させることができますが、メモリアクセスやその他のオーバーヘッドも影響します。 何がロウワーできるかを決定する ロウワーパスは2つの質問に答える必要があります。 このループの異なるイテレーションは独立して実行できるか？ ループ内の各値について、それは均一か、それとも変化するか？ 「均一な値」とは、すべてのレーンで同じ値を持つものです。「変化する値」とは、各レーンで異なる可能性がある値です。この区別は重要です。なぜなら、均一な値はすべてのレーンで共有されますが、変化する値はレーンごとの計算を必要とするからです。 オーディオの例では、`volume`は均一であり、すべてのレーンが同じ値で乗算します。しかし、`i`は変化します。各レーンが異なるインデックスを処理するため、`samples[i]`も変化します。 value kind why --------- ------- ---------------------------- volume uniform same value in every lane i varying each lane gets a different index samples[i] varying each lane reads a different sample 私のコンパイラの核心は、小さなASTウォーカー分類器です。 # (疑似コードだが、現実とそう遠くない) def kind(expr, env): if expr is a literal: return UNIFORM if expr is a variable: return env[expr.name] if expr is a load: return VARYING if kind(expr.index, env) == VARYING else UNIFORM if expr is a binary expression: left = kind(expr.left, env) right = kind(expr.right, env) return VARYING if VARYING in (left, right) else UNIFORM ループの前に、カーネルパラメータは均一であると仮定されます。ロウワーされたループ内では、コンパイラはループインデックスを変化するものとしてマークします。 env = {param: UNIFORM for param in kernel.params} env[i] = VARYING そこから、変化性は式全体に伝播します。`i`が変化するため、`samples[i]`も変化します。`samples[i]`が変化するため、`samples[i] * volume`も変化します（`volume`自体は均一であっても）。 この分類によって、コンパイラは何を出力すべきかを知ることができます。ループがロウワーされると、`i`は`base + lane_id`になります。`base = 8`の4レーングループの場合、レーンは`[8, 9, 10, 11]`を保持します。したがって、`samples[i]`は`samples[8]`、`samples[9]`、`samples[10]`、`samples[11]`を要求します。これは連続したメモリ領域です！ コンパイラは、`i`がループの連続インデックスであることを記録します。そのインデックスを使用するロードは、マスク付きロードになります。マスクは、いくつかのレーンインデックスが配列の外部にある可能性がある最後のグループでのみ重要です。 samples[i] -> masked_load(samples, i, active) しかし、任意のレーンごとのインデックスからの変化するロードは「ギャザー」になります。ギャザーは、「各レーンが独自のアドレスから読み取る」というグループ化されたロードバージョンであり、アドレスが異なる、隣接していない可能性があります。たとえば、この最適化されていないコードでは次のようになります。 kernel color_by_number(color_number, colors, out, n): for i in range(n): number = color_number[i] out[i] = colors[number] 各レーンは異なる`number`をロードするため、`number`は変化します。これは、`colors[number]`が連続したロードではないことを意味します。 colors[number] -> gather(colors, number, active) ギャザーは、各レーンが独自の`address`から読み取りながらも、単一のグループ化された操作として実行することを可能にします。これは依然として並列作業ですが、メモリアクセスパターンは、レーンが隣接する要素を読み取る`samples[i]`よりも規則性が低いです。そのため、ギャザーは遅くなることが多いですが、コストはアーキテクチャに依存します。 上記の`color_by_number`の例の完全なコンパイル出力です。 kernel color_by_number(color_number, colors, out, n): vector_for base in range(0, n, LANES): let i = (base + lane_id) let active = (i < n) let number = masked_load(color_number, i, active) masked_store(out, i, gather(colors, number, active), active) 実際のコンパイラは、ループのイテレーションが安全に独立して実行できるかを判断することを含め、これよりもはるかに多くのこと（型、エイリアシング、制御フロー、ターゲット固有の命令）を処理する必要があります。私のコンパイラは、`out[i]`に書き込む単純なカーネルに限定することで、ほとんどそれを前提としていますが、それでも私が理解したかった部分には十分でした。 なぜこのステップが重要なのか このロウワーパスの後、グループ化された実行がプログラム内で明示的になります。どのループのイテレーションが一緒に実行されるか、どのレーンがアクティブであるか、そして各ロードが隣接するアドレスを使用できるか、またはレーンごとのアドレスを必要とするかが記録されます。 これにより、後のコード生成パスは、より良い命令を出力するために必要な構造を得られます。`masked_load`はマスク付きベクトルロードになり、`gather`はギャザー命令になり、`vector_for`はそれらの操作を囲むループ構造になります。このような分析がなければ、これらの事実はプログラムに存在しないため、通常のスカラ操作のままで、より高速な命令を使用する機会を逃してしまいます。 まとめ 私はこの依存関係分析と、いつどのようにロウワーするかという問題に取り組むことを楽しみました。必要ないときにパーサーを書くなど、他のプロジェクトで取った非効率な手順は避けました。 コンパイラのソースコードを確認すると、入力はハードコードされたツリーになっています。 AST = Kernel( "color_by_number", ["color_number", "colors", "out", "n"], For( "i", Lit(0), V("n"), Let( "number", Load("color_number", V("i")), Store("out", V("i"), Load("colors", V("number"))), ), ), ) 当初はさらに進んでC SIMDコードを生成し、コンパイル済み/未コンパイルのコードでベンチマークを実行しましたが、それは私が学びたかった核心的なアイデアから焦点を外れてしまいました。そこで、より純粋で高レベルの抽象化に戻しました。 おまけとして、このウェブサイトにカーネルスニペットをサポートするためのカスタム構文ハイライトを追加しました。