カーネル最適化の体系

この章でわかること:

GPUで「全要素を1つに畳む」計算(reduction)の実装の定石3段階
atomic競合、バリアのコスト、スレッド粗粒度化というカーネル設計の語彙
「教科書の定石が目の前のGPUで覆される」実測と、その読み解き方
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スレッドの仕事を厚くし、同期と共有を減らす」方向は、多くのコスト比で成立する、最も頑健な定石です(それでも万能ではなく、粗粒度化しすぎると並列度と占有率を失います—— グループ数を振って測るのが仕上げです)。そして答え合わせの方法は常に1つ、実測です(26章で GPU側の正確な計測法を導入します)。

まとめ

reductionは「並列に計算し、同期して集める」というGPUのもう1つの基本形です
3つの実装の実測(M4): 全atomicが約2〜4ms、共有メモリの木が約5〜15ms、粗粒度化+木が約1〜2ms。教科書の定石(v2)が最下位という結果でした
読み解き: 1スレッドの仕事が薄いとバリアと同期のコストが支配します。「レジスタで束ねてから畳む」粗粒度化が最も頑健な定石です
定石はハードウェアのコスト比の関数です。atomicが速いGPUも、バリアが重いGPUもあります。移植したら測り直してください

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