CUDA SharedMemory

CUDAのSampleにimmaTensorCoreGemmというのがあり、CUDA 10から導入されたWarp Matrix Multiply and Accumulate (WMMA) APIを使ってuint8_tのGEMM演算を行っている。

github.com

SharedMemoryのサイズが一定以上なら compute_gemm_imma という kernel を呼び出し、そうでないなら simple_wmma_gemm_imma を呼び出す。自分が今使っている環境だと compute_gemm_imma を実行する場合は 110TOPS ぐらい出るけれど、simple_wmma_gemm_imma を実行する場合は 30TOPS ぐらいしか出ない。という事で shared memory を使わないで global memory だけだと演算器の性能を引き出せていない事が確認出来た。

async-copy を使えばもっと性能を引き出せるんだろうけど、CUDA Toolkit 11.2 時点では使われていない。bf16TensorCoreGemm では使われているので是非 immaTensorCoreGemm でも使うようにしてほしい。なお bf16TensorCoreGemm では以下のような処理性能になった。

コマンドライン引数 selected_kernel TOPS
kernel=0 bf16mma_shmem_gemm_async_copy 52.68
kernel=1 bf16mma_shmem_gemm 49.48
kernel=2 simple_bf16mma_gemm 14.64

実行するとグラフィックボードのファンが勢い良く回転するのでブーストクロックによって数値は少し変動しそう。

実行結果から判断できるのはやはり Shared Memory を使わないと性能は引き出せないという事と、async-copy を使う事で僅かに性能を改善出来る事。async-copy の効果が想像していたより少なく感じるけど、消費電力を改善する効果があるという事なので使えるなら使った方が良いだろう。

wccftech.com

の記載によると Peak INT8 Tensor TOPS は 284/568 という事だけれど実際に自分の環境でサンプルを動かすと110TOPSくらいしか出ないので演算器のピーク性能を引き出すプログラムを書くのは簡単ではなさそう。なおサンプルコードのGEMMのTOPSの算出式は M×N×K×2÷処理時間の秒÷1e12