カーネル最適化の体系

この章でわかること:

• GPUで「全要素を1つの値に集約する」計算(reduction)の実装の定石3段階
• atomic競合、バリアのコスト、スレッド粗粒度化というカーネル設計の基本概念
• 教科書の定石が実測で覆される例と、その解釈
• subgroup命令という次の選択肢
• なぜ体系に必須か: 12章の行列積は「要素ごとに 独立」な計算でした。もう1つの基本形である「全体を1つの値に集約する」 計算を書けなければ、GPUカーネルの基本形の半分しか扱えません

この章の実験はリポジトリのexamples/ch23-reductionです。

cd examples
cargo run --release -p ch23-reduction

計測環境は筆者のMac(Apple M4)です。共有環境ではないものの 実行ごとに数十%のばらつきがあるため、概数で示します。

題材: reduction

reduction(リダクション、集約)は、配列全体を1つの値に集約する 計算の総称です。総和、最大値、平均は、どれも同じ形をしています。 CPUならiter().sum()の一行ですが(4章で見たとおり、 そこでもSIMDとILPが使われています)、GPUでは 数百万のスレッドの結果を1つに集めるという、9章のモデルには そのまま当てはまらない問題になります。各要素の計算は独立に 並列化できますが、「集める」処理は本質的に共有と同期 (5章)を必要とするからです。

題材は1677万要素(64MB)のu32の総和です。CPU(1コア)の実測は 約2〜3ミリ秒で、自動ベクトル化された実装が比較の基準です。

v1: 全スレッドがatomicに加算する

最も単純な設計では、スレッド1本が1要素を担当し、グローバルな 合計変数にatomicAddします。

@compute @workgroup_size(256)
fn reduce_atomic(...) {
    let i = /* 自分の担当添字 */;
    if (i < arrayLength(&input)) {
        atomicAdd(&result, input[i]);
    }
}

1677万回のatomic操作が1つの変数に集中します。 17章で見たRMW操作の競合の、 1677万スレッド版です。古典的なGPUの教科書では 「絶対にやってはいけない設計」の筆頭に挙げられます。

実測は約2〜4ミリ秒で、CPU1コアとほぼ同等です。 理由は後で考察します。

v2: 共有メモリの木で集約する

教科書の定石では、ワークグループ(256スレッド)内で 共有メモリ(10章)に値を置き、 「256→128→64→…→1」と半分ずつ足し合わせる木(ツリー)で集約し、 グループの代表1本だけがatomicAddします。atomic操作は 65,536回(グループ数)まで減ります。

var<workgroup> partial: array<u32, 256>;

@compute @workgroup_size(256)
fn reduce_shared(...) {
    partial[lid.x] = /* 担当要素 */;
    workgroupBarrier();
    var stride = 128u;
    while (stride > 0u) {
        if (lid.x < stride) {
            partial[lid.x] += partial[lid.x + stride];
        }
        workgroupBarrier(); // 段ごとに全スレッドを同期する
        stride = stride / 2u;
    }
    if (lid.x == 0u) { atomicAdd(&result, partial[0]); }
}

実測は約5〜15ミリ秒で、どの実行でもv1より2倍以上遅くなりました。

定石と逆の結果です。このカーネルは 1スレッドが1要素しか処理しないのに、8段の木の各段で workgroupBarrier()を実行します。バリアは「グループ内の全スレッドが 到達するまで待つ」操作なので、ワープ(9章)の 実行がそのたびに止まります。処理(ロード1回)に対して 同期(バリア8回)が多すぎます。12章で見た 「1スレッドあたりの処理量が少なすぎる」問題の同期版です。

v3: レジスタで部分和を取ってから集約する

対策は12章と同じで、スレッドの粗粒度化(thread coarsening)です。 起動するスレッドを26万本に減らし、1スレッドがまず64要素を レジスタ上で足します(同期もatomicも不要な、純粋な逐次加算)。 そのあとで初めて共有メモリの木とatomicを使います。

var sum = 0u;
var i = /* 自分の開始位置 */;
while (i < n) {
    sum += input[i];
    i += 総スレッド数; // グリッドストライドループ
}
// ここまでで加算の 98% が終わっている。残りだけツリーで集約する

添字を「総スレッド数」ずつ進めるグリッドストライドループ (grid-stride loop)にしているのは、隣のスレッドが常に隣の要素を 読む形、つまりコアレッシング(10章)を 保つためです。

実測は約1〜2ミリ秒(最速0.98ミリ秒)で、v1の約3倍、v2の約8倍、 CPU1コアの2倍前後の速さです。64MBを約1ミリ秒で処理しており、 実効帯域は毎秒50〜65GBです。ここで初めてメモリ帯域律速 (10章のルーフラインの帯域側の上限)に 達しました。総和の算術強度は ベクトル加算並みに低いので、これがこの問題の上限近くです。

考察: 定石はハードウェアのコスト比に依存する

この実験で最も重要な点は、「v1が遅くない」「v2が遅い」という結果が 古典的なGPUの常識の逆だという事実であり、順位そのものではありません。

古典的な常識は「グローバルatomicの衝突は数百倍の直列化を生む」 「共有メモリの木こそ正解」でした。それはNVIDIAの古い世代などで 実測された事実です。一方、筆者のApple M4では、同一アドレスへの atomic加算をハードウェアが極めて効率よく併合するらしく (毎秒数十億回を処理しています)、むしろバリアの相対コストが 目立ちました。

つまり、カーネル最適化の定石は、特定ハードウェアの コスト比を前提にしています。移植先ではコスト比が変わり、 順位が入れ替わります。それでもv3が最速だったことに注目してください。 「1スレッドの処理量を増やし、同期と共有を減らす」方向は、 多くのコスト比で成立する、最も頑健な定石です(それでも万能では なく、粗粒度化しすぎると並列度と占有率が下がります。 グループ数を変えて計測し、最適点を決めます)。 検証の方法は実測だけです(26章で GPU側の正確な計測法を導入します)。

まとめ

• reductionは「並列に計算し、同期して集める」というGPUの もう1つの基本形です
• 3つの実装の実測(M4)では、全atomicが約2〜4ms、共有メモリの木が 約5〜15ms、粗粒度化+木が約1〜2msでした。教科書の定石(v2)が 最下位という結果です
• 1スレッドの処理量が少ないと、バリアと同期のコストが支配的になります。 レジスタで部分和を取ってから集約する粗粒度化が、最も頑健な定石です
• 定石はハードウェアのコスト比に依存します。atomicが速いGPUも、 バリアのコストが大きいGPUもあります。移植したら測り直してください

次章は、カーネルの外側にあるCPUとGPUの間の転送と同期を 最適化します。11章で見た往復のオーバーヘッドを、 設計でどこまで隠蔽できるかを実測します。