协作头脑风暴
brainstorming
在实现前澄清需求、探索设计选择,并形成清晰的开发方向。
pipeline-gpu-kernel【开发与工程】Apply software pipelining (double-buffering) to a tiled GPU kernel to overlap global memory loads with Tensor Core computation. Covers prologue/loop/epilogue restructuring, LDG-register vs cp.async (LDGSTS) variant selection based on compute/load ratio, shared memory budget verification against architecture-specific occupancy cliffs, and SASS-level v…
独立 GitHub 快照 · 自动分类 · 以原作者文档为准
タイル化された GPU カーネルにソフトウェアパイプライニング(double-buffering)を適用し、tile N+1 のグローバルメモリロードを tile N の Tensor Core 計算と重ねる。順次の load-sync-compute-sync K ループを prologue/loop/epilogue 構造に変換し、タイルあたりの compute/load 比に基づいて LDG-register と cp.async(LDGSTS)バリアントの間で選択し、共有メモリがアーキテクチャ占有率クリフ未満に留まるか検証し、最終 SASS でロード/計算の重ねを確認する。
analyze-kernel-bottleneck がタイルあたり低 compute/load 比の memory-bound カーネルを特定したとき.cu)analyze-kernel-bottleneck から; 提供されなければ推定される)カーネルが __syncthreads() で分離された別個のロードと計算フェーズを持つタイル化された K ループを持つことを確認する。倍にされた共有メモリコストを計算しアーキテクチャ占有率クリフ未満に留まるか検証する。
__syncthreads()、共有メモリタイル上で計算(HMMA/IMMA/FFMA)、__syncthreads()。smem_a_size = BM * BK * sizeof(T) と smem_b_size = BK * BN * sizeof(T)。smem_doubled = smem_a_size * 2 + smem_b_size * 2。Single buffer: smem_a[BM*BK] + smem_b[BK*BN] = 2 KB + 2 KB = 4 KB
Double buffer: smem_a[2][BM*BK] + smem_b[2][BK*BN] = 4 KB + 4 KB = 8 KB
8 KB << 50 KB cliff -> 2 blocks/SM -> 8 warps
num_tiles = K / BK。パイプライニングは num_tiles >= 2(少なくとも一つの prologue + 一つのメインループ反復)を要求する。期待結果: 単一バッファと倍バッファコストを示す共有メモリ予算表、倍にされた割り当てがアーキテクチャクリフ未満に少なくとも 2 blocks/SM 占有率で留まることを確認。
失敗時: 倍バッファがクリフを超えるなら、smem_doubled <= 50 KB(GA104)になるまでタイルサイズを減らす(BK または BM を半分に)。代替として、共有メモリを倍にせずレジスタのみのプリフェッチ(LDG バリアント)を使う — プリフェッチデータをレジスタに保存し、__syncthreads() 後に同じ単一バッファに書く。
タイルあたりの compute/load 比に基づいて LDG-register と cp.async(LDGSTS)の間で選択する。
ratio = (2 * BM * BN * BK) / ((BM * BK + BK * BN) * sizeof(T))(multiply-add ごとに 2 FLOPs、タイルあたりロードバイト)。LDG-register バリアント(ratio >= 5 または CUDA < 11.0):
buf[N % 2] 上で計算(未処理 LDG と重なる)。__syncthreads()、それから STS レジスタを buf[(N+1) % 2] へ、__syncthreads()。(BM * BK + BK * BN) / BLOCK_SIZE レジスタ。cp.async(LDGSTS)バリアント(ratio < 5、CUDA >= 11.0):
__pipeline_memcpy_async tile N+1 を buf[(N+1) % 2] に直接(async、レジスタファイルバイパス)。__pipeline_commit()。buf[N % 2] 上で計算。__pipeline_wait_prior(0) + __syncthreads()。#include <cuda_pipeline.h> を要求。期待結果: compute/load 比とターゲットアーキテクチャに基づく根拠付きで選ばれたバリアント。
失敗時: 比が曖昧(5-20:1 範囲)なら、両バリアントを実装してベンチマークする。CUDA バージョンがサポートするとき cp.async バリアントが安全な既定。
順次 load-sync-compute-sync ループをパイプライン prologue/loop/epilogue 構造に変換する。
3 つのセクションを特定: 元のループ本体は 3 つの部分になる:
buf[0] にロード、同期、それからメインループに入る。num_tiles - 1 まで、tile N+1 のロードを tile N の計算と重ねる。LDG-register バリアント構造:
// === LDG-register variant ===
// Prologue: load tile 0 into buf[0]
cooperative_load_tile(smem_a[0], smem_b[0], global_a, global_b, /*k_offset=*/0);
__syncthreads();
for (int tile = 0; tile < num_tiles - 1; tile++) {
int cur_buf = tile & 1;
int next_buf = 1 - cur_buf;
// Phase 1: LDG next tile into registers (non-blocking)
float reg_a[ELEMS_PER_THREAD_A], reg_b[ELEMS_PER_THREAD_B];
prefetch_tile_to_registers(reg_a, reg_b, global_a, global_b,
(tile + 1) * BK);
// Phase 2: Compute on current buffer (overlaps with LDG flight)
tensor_core_mma(smem_a[cur_buf], smem_b[cur_buf], acc);
// Phase 3: Drain registers into next buffer
__syncthreads();
store_registers_to_smem(smem_a[next_buf], smem_b[next_buf],
reg_a, reg_b);
__syncthreads();
}
// Epilogue: compute last tile
tensor_core_mma(smem_a[(num_tiles - 1) & 1], smem_b[(num_tiles - 1) & 1], acc);
// === cp.async variant ===
#include <cuda_pipeline.h>
// Prologue: async load tile 0 into buf[0]
cpasync_load_tile(smem_a[0], smem_b[0], global_a, global_b, /*k_offset=*/0);
__pipeline_commit();
__pipeline_wait_prior(0);
__syncthreads();
for (int tile = 0; tile < num_tiles - 1; tile++) {
int cur_buf = tile & 1;
int next_buf = 1 - cur_buf;
// Phase 1: cp.async next tile into next buffer (async, direct to smem)
cpasync_load_tile(smem_a[next_buf], smem_b[next_buf],
global_a, global_b, (tile + 1) * BK);
__pipeline_commit();
// Phase 2: Compute on current buffer (overlaps with LDGSTS in flight)
tensor_core_mma(smem_a[cur_buf], smem_b[cur_buf], acc);
// Phase 3: Wait for async copies to complete
__pipeline_wait_prior(0);
__syncthreads();
}
// Epilogue: compute last tile
tensor_core_mma(smem_a[(num_tiles - 1) & 1], smem_b[(num_tiles - 1) & 1], acc);
num_tiles - 1 反復走る(tile 0 から num_tiles - 2 がどのタイルを計算するかをインデックス、tile 1 から num_tiles - 1 をロード)。Epilogue が最後の反復でロードされたタイルを計算する。期待結果: 選ばれたバリアント用に明確な prologue、メインループ、epilogue セクションを持つ再構築された K ループソースコード。
失敗時: 最も一般的なバグはバッファインデックスでの off-by-one または epilogue 計算パスを忘れること。検証: prologue は buf[0] にロード、最初のメインループ反復は buf[0] で計算し buf[1] にロード、二番目反復は buf[1] で計算し buf[0] にロード、等。Epilogue は buf[(num_tiles - 1) & 1] を計算する。
倍バッファ共有メモリを宣言しロード関数を実装する。
// Before (single buffer)
__shared__ half smem_a[BM * BK];
__shared__ half smem_b[BK * BN];
// After (double buffer)
__shared__ half smem_a[2][BM * BK];
__shared__ half smem_b[2][BK * BN];
__device__ void cpasync_load_tile(half* dst_a, half* dst_b,
const half* src_a, const half* src_b,
int k_offset) {
// Each thread copies its portion (16 bytes = 8 half values per cp.async)
int tid = threadIdx.x;
int bytes_per_thread = 16; // cp.async.cg supports 4, 8, or 16 bytes
// A tile: BM * BK elements, distributed across BLOCK_SIZE threads
int elems_a = BM * BK / BLOCK_SIZE;
for (int i = 0; i < elems_a; i += 8) {
int idx = tid * elems_a + i;
__pipeline_memcpy_async(dst_a + idx,
src_a + k_offset * BM + idx,
bytes_per_thread);
}
// B tile: BK * BN elements, distributed similarly
int elems_b = BK * BN / BLOCK_SIZE;
for (int i = 0; i < elems_b; i += 8) {
int idx = tid * elems_b + i;
__pipeline_memcpy_async(dst_b + idx,
src_b + k_offset * BN + idx,
bytes_per_thread);
}
}
// Declare register staging (size = elements per thread)
half reg_a[BM * BK / BLOCK_SIZE];
half reg_b[BK * BN / BLOCK_SIZE];
// Prefetch: LDG from global to registers (non-blocking, issued early)
for (int i = 0; i < BM * BK / BLOCK_SIZE; i++) {
int idx = threadIdx.x * (BM * BK / BLOCK_SIZE) + i;
reg_a[i] = global_a[k_offset * BM + idx];
}
// ... similarly for reg_b
// Store: STS from registers to shared memory (after __syncthreads)
for (int i = 0; i < BM * BK / BLOCK_SIZE; i++) {
int idx = threadIdx.x * (BM * BK / BLOCK_SIZE) + i;
smem_a[next_buf][idx] = reg_a[i];
}
__launch_bounds__(BLOCK_SIZE) を保つ。nvcc --cubin -arch=sm_86 -O2 -o kernel.sm_86.cubin kernel.cu。期待結果: 倍バッファ共有メモリと選ばれたロード機構を持つコンパイル可能なカーネル。エラーなしの成功 cubin 生成。
失敗時: パイプライン API 呼び出しでコンパイルが失敗したら、#include <cuda_pipeline.h> が存在し CUDA toolkit が >= 11.0 か保証する。レジスタスピルが起こる(nvcc --resource-usage を確認)なら、BLOCK_SIZE を増やすか BK を減らしてレジスタステージング配列サイズを減らす。
CPU 参照に対してパイプライン化されたカーネルを動かして同一の数値出力を確認する。
nvcc -arch=sm_86 -O2 -o bench bench.cu -lcuda -I../../phase2/common。abs=0.5, rel=0.1abs=1e-2, rel=1e-2abs=1e-3, rel=1e-3期待結果: 非パイプラインベースラインと同一のエラー境界で小・ターゲット問題サイズの両方で PASS。
失敗時: バッファインデックスバグが最も可能性の高い原因。検証: 計算は buf[tile & 1] から読み、ロードは buf[1 - (tile & 1)] に書く。Epilogue がバッファインデックス (num_tiles - 1) & 1、num_tiles & 1 ではなく、を処理するか確認する。cp.async については、__pipeline_wait_prior(0) が __syncthreads() 前に完了することを検証する — そうでないと計算が部分書き込みデータを読みうる。
ターゲット問題サイズで非パイプラインベースラインに対してパイプラインカーネルを計測する。
speedup = pipelined_metric / baseline_metric。| Variant | GFLOPS | Speedup vs Baseline |
|------------------|--------|---------------------|
| Baseline | XXX | 1.00x |
| LDG-register | XXX | X.XXx |
| cp.async (LDGSTS)| XXX | X.XXx |
期待結果: 改善を示す性能比較表。選ばれたバリアントは compute/load 比予測と一致する計測可能なスピードアップを示すべき。
失敗時: 性能がリグレッションするなら、3 つを確認: (1) 予期しない命令オーバーヘッドの SASS(余分な BAR.SYNC、レジスタスピル)。(2) 共有メモリが占有率クリフを越えなかった — nvcc --resource-usage または cuobjdump -res-usage で検証。(3) 問題サイズが prologue/epilogue オーバーヘッドを償却するに十分なタイル(K / BK >= 4)を生む。
コンパイルされた SASS を検査してメインループ本体内でグローバルロードと Tensor Core 命令が重なるか確認する。
cuobjdump -sass kernel.sm_86.cubin | grep -E 'IMMA|HMMA|LDGSTS|LDG|BAR'。LDGSTS または LDG 命令が HMMA または IMMA 命令の 前 に現れる。BAR.SYNC がない(warp スケジューラで重ねるため自由でなければならない)。BAR.SYNC が計算ブロックの 後 に現れ、ロードされたデータの次の反復での使用をゲートする。# Full SASS pipeline verification
cuobjdump -sass kernel.sm_86.cubin | grep -E 'IMMA|HMMA|LDGSTS|LDG|BAR'
# Count compute instructions per loop
cuobjdump -sass kernel.sm_86.cubin | grep -c 'HMMA\|IMMA'
# Check for register spills
nvcc --resource-usage --cubin -arch=sm_86 -O2 kernel.cu 2>&1 | grep -i spill
期待結果: 介在バリアなしで load-before-compute パターンを示す注釈付き SASS 抜粋。レジスタスピルゼロ。
失敗時: コンパイラが計算後にロードを並べ替えた(重ねを破る)なら試す: (1) 過攻撃的な展開を防ぐためメインループに #pragma unroll 1。(2) シーケンシングヒントを作るためロードと計算を別個のインライン関数に分ける。(3) ロードと計算ブロック間のコンパイラフェンスとして asm volatile("" ::: "memory") を使う(最後の手段 — 他の最適化を阻害するかも)。
buf[tile & 1] パターン)buf[0] にロードbuf[(num_tiles - 1) & 1] から最後のタイルを計算BAR.SYNC なし)nvcc --resource-usage を確認)smem_doubled を計算する。28 KB 単一バッファのカーネルが倍で 56 KB に跳ね、クリフを越え占有率を半減する。これがパイプライニング +20% 利得を -50% 占有率リグレッションに変えうる。buf[tile & 1] を、次のロードバッファに buf[1 - (tile & 1)] を使う。一般的な誤りは次のバッファに buf[(tile + 1) & 1] を使うこと。これはバッファ数が 2 のときのみ buf[1 - (tile & 1)] と等価 — しかし計算インデックスに誤って適用されると間違って読む。__pipeline_commit() は計算フェーズの 前 に呼ばねばならない(async コピーのバッチを封印する)。__pipeline_wait_prior(0) は計算フェーズの 後 に呼ばねばならない(コミットされたコピーすべてが完了するまでブロックする)。これらを入れ替えると async コピーが同期的になり、すべての重ね利益を排除する。__syncthreads() が必要(計算が現バッファの読み取りを終えてから上書きされる)。STS ドレイン後にもう一つ __syncthreads() が必要(次の反復が読む前にすべてのスレッドが書き終える)。cp.async バリアントでは、__pipeline_wait_prior(0) の後の __syncthreads() がすべてのスレッドが完了した async コピーを見ることを保証する。__pipeline_memcpy_async はソースアドレスが有効でアラインメントされていることを要求する。K が BK の倍数でない行列エッジでは、最後のタイルが境界外を読みうる。最終タイルには境界チェック付きスカラーロードにフォールバックするか、入力行列を BK の倍数までパディングする。analyze-kernel-bottleneck — カーネルが memory-bound かを特定し、バリアント選択を駆動する compute/load 比を計算する