Vulkan コンピュートシェーダでサブグループ操作を使う方法

Post 62。コンピュートシェーダの「呼び出し毎に1スレッド」モデルの下にはハードウェアの真実がある：呼び出しは サブグループ としてロックステップで動く（NVIDIA では warp、AMD では wavefront—— RDNA では64レーン）。サブグループ操作はレーン間で1命令でデータを交換する——共有メモリもバリアも不要。64値のリダクションは1回の subgroupAdd。

example.sh

$ ./app
subgroup size=64 ops={ Basic | Vote | Arithmetic | Ballot | Shuffle |
  ShuffleRelative | Clustered | Quad | Rotate | RotateClustered }
sum=2080 (want 2080)  count=22 (want 22)  bcast=1  shuffled=33
OK

Example: VulkanExamples/p62_subgroup

サポートを確認し、SPIR-V 1.3+ 向けにコンパイル

main.cpp

auto sg = phys.getProperties2<vk::PhysicalDeviceProperties2,
    vk::PhysicalDeviceSubgroupProperties>()
    .get<vk::PhysicalDeviceSubgroupProperties>();
// sg.subgroupSize       — サブグループあたりのレーン数（ここでは64）
// sg.supportedStages    — どのステージで使えるか
// sg.supportedOperations — どの操作クラスがあるか

GLSL には GL_KHR_shader_subgroup_* 拡張と SPIR-V 1.3 が必要。 shaderc のデフォルトは SPIR-V 1.0 をターゲットにする——コンパイルは requires SPIR-V 1.3 で失敗する。ターゲットを直す：

main.cpp (shaderc options)

shaderc::CompileOptions opts;
opts.SetTargetEnvironment(shaderc_target_env_vulkan,
                          shaderc_env_version_vulkan_1_1); // → SPIR-V 1.3
auto res = cc.CompileGlslToSpv(src, shaderc_compute_shader,
                               "sg.comp", opts);

4つの操作ファミリー

example.txt

uint lane = gl_SubgroupInvocationID;   // 自分のレーンインデックス
uint val  = lane + 1;

uint s  = subgroupAdd(val);                 // ツリー reduce：全レーンが和を得る
uvec4 b = subgroupBallot(lane % 3 == 0);    // 一致するレーンのビットマスク
uint bc = subgroupBroadcastFirst(val);      // 最小アクティブレーンの値
uint sh = subgroupShuffle(val, 32);         // 別レーンの値を読む
uint n  = subgroupBallotBitCount(b);        // ballot の popcount

操作	結果	コスト
`subgroupAdd/Min/Max`	アクティブレーンのリダクション	~1命令
`subgroupBallot` + `BitCount`	「何レーンが通ったか」	vote + popc
`subgroupBroadcastFirst`	正規レーンの値	1命令
`subgroupShuffle`	レーン↔レーン交換	1命令
`subgroupInclusiveAdd`	プレフィックス和	log2(N) shuffle

バグ第1位：分岐した呼び出し

サブグループ操作はすべてのアクティブレーンが参加する必要がある。これ：

example.txt

if (lane == 0) {               // 間違い — サブグループを1レーンに狭める
    out = subgroupAdd(val);    // lane0 の値だけを返す！
}

は 2080 ではなく 1 を返す——コンパイルエラーではなく、静かな間違い。正しいパターン：全レーンで評価し、ストアだけをゲートする：

example.txt

uint s = subgroupAdd(val);     // 全レーンが評価
if (lane == 0) sum = s;        // 1レーンが書き込む

使うべき場面

リダクション：共有メモリのピンポンなしにワークグループ全体の和/最大/最小——64値を約6命令で。
Ballot 駆動コンパクション：subgroupBallotFindLSB で「次の未処理ワークアイテム」——post 61 のインダイレクト描画と素晴らしく相性が良い。
Warp 全体のアトミック：subgroupElect + サブグループ毎1アトミック。
単純な barrier()+shared[] の方が明確な場合は使うな——サブグループ操作が輝くのは既にレーン並列なとき。