[CUDA-Samples] cudaTensorCoreGemm

CUDA Advent Calendar 2019 6日目の記事です．
ネタのつなぎのCUDA Samples読み回2回目です．
読んでいたもののバージョンが変わってしまうといけないのでこちらにforkしてあります．

何このサンプル?

Tensorコアを使ったGemmのサンプルコードです．
実行するとこのように行列積を計算しFlop/sを測って表示してくれます．

Initializing...
GPU Device 0: "Volta" with compute capability 7.0

M: 4096 (16 x 256)
N: 4096 (16 x 256)
K: 4096 (16 x 256)
Preparing data for GPU...
Required shared memory size: 64 Kb
Computing... using high performance kernel compute_gemm
Time: 3.100640 ms
TFLOPS: 44.33

このサンプルはざっと見ただけでも

• Sharedメモリブロッキング
• レジスタブロッキング
• SharedメモリのBank Conflict回避

などの工夫がなされているようです．
各種ブロッキングについてだと，ブロッキングサイズの決め方などについてこんな論文があります（ちょっと古いですが）．
Guangming Tan, Linchuan Li, Sean Triechle, Everett Phillips, Yungang Bao, Ninghui Sun - Fast Implementation of DGEMM on Fermi GPU SC11
コメントによく出てくるCTAはCompute Thread Arrayでblockと同義です．

このサンプルの面白いところ

サンプルコード中の

extern __shared__ half shmem[][CHUNK_K * K + SKEW_HALF];

がSharedメモリのBank Conflict回避に当たる部分です．
SKEW_HALFが8と定義されており，

extern __shared__ half shmem[][CHUNK_K * K ];

とした場合に起こりうるBank Conflictを減らしています．
もっともWMMA APIのload_matrix_sync関数には読み込む行列のポインタにアラインメント制約があるため，ずらす値をSKEW_HALF=8としてこのアラインメント制約に対応するようになっています．
このようにWMMA APIとBank Conflict回避を組み合わせて使うためのちょっとした工夫も必要なようです．

おわり

WMMA APIの使い方というよりは行列積和の高速化のための工夫を知るためのいいサンプルコードだと思います．
どうでもいいことですが，CUDAではwarpSizeというWarpの大きさの組み込み変数が用意されているにも関わらずWARP_SIZEという定数が定義されているのが興味深かったです．
やはりNVIDIAの方々もこの変数は使い勝手が悪いのでしょうか? - CUDAのwarpSizeについて - 天炉48町