博士

やあ、ロボ子。今日はGPUのテンソルコアを使ったGEMMの高速化について話すのじゃ。

ロボ子

GEMM、つまり行列積の高速化ですね。CUDAの知識が必要とのことですが、私にも理解できるでしょうか？

博士

大丈夫だぞ！この記事では、warpレベルで動作するMMA命令、つまり`Matrix Multiply Accumulate`命令を使うのじゃ。NVIDIAのコンシューマーGPUではこれが一番速いらしい。

ロボ子

warpレベル…32スレッドで動作するんですね。具体的にはどういう命令を使うんですか？

博士

`mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32`みたいなのじゃ。A行列が16x16、B行列が16x8で、データレイアウトとかスレッドのデータのロード方法が決まっているのじゃ。

ロボ子

なんだか呪文みたいですね…。共有メモリからレジスタへのロードには`ldmatrix`命令を使うと。

博士

そう！`ldmatrix.sync.aligned.m8n8.x4.shared.b16`とかじゃな。これで高速化できるのじゃ。

ロボ子

Mが16で割り切れ、Nが8で割り切れる場合、Kが16であれば複数のブロックを起動してGEMMを計算できるんですね。

博士

その通り！Kが16の倍数なら、KをK/16個に分割して結果を累積するのじゃ。

ロボ子

最適化の基本原則は、データ移動を減らして計算を増やすこと、データ再利用、データ局所性、非同期データ転送などを活用してメモリ階層を最適化するんですね。

博士

よくできました！でも、3090 GPUでのベンチマーク結果を見ると、自作のテンソルコアカーネルはcuBLASに比べてかなり遅いのじゃ。

ロボ子

えっ、そうなんですか？M=1024, N=1024, K=32の場合、自作カーネルが35.20 GFLOPsで、cuBLASが234.90 GFLOPs…6.67倍も違うんですね。

博士

そう、まだまだ改善の余地があるのじゃ。でも、テンソルコアの可能性は無限大！

ロボ子

なるほど。奥が深いですね。私ももっと勉強して、博士のお役に立てるように頑張ります！

博士

期待してるぞ！ところでロボ子、GEMMって何の略か知ってるか？

ロボ子

えっと…General Matrix Multiply、ですよね？

博士

ブー！正解は「ご飯めっちゃ盛って満腹」じゃ！

ロボ子

…博士、それ、今考えましたね？