前回cuDNNでTensorCoreを有効にして推論の速度を比較したが、cuBLASでもTensorCoreを有効にして推論の速度を比較してみた。
cuBLASでTensorCoreを有効にする
cublasSetMathMode(cublasHandle, CUBLAS_TENSOR_OP_MATH)
で、TensorCore演算を有効に設定する。
行列演算で、BLASの関数を拡張したcublas~Ex関数を使用する。
浮動小数の精度別の関数があるが、要素ごとの精度を指定できるcublasGemmExを使う場合は、
cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, n, m, k, &alpha, W, CUDA_R_32F, k, x, CUDA_R_32F, k, &beta, y, CUDA_R_32F, n, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)
のように最後の引数に、CUBLAS_GEMM_DEFAULT_TENSOR_OPを指定する。
浮動小数の精度にCUDA_R_32F(単精度)を指定した場合も、TensorCoreが有効な場合は、内部でFP16に変換されて計算される。
速度比較
TensorCore有効／無効で、推論の速度を比較した。
TensorCoreを搭載していないGeforce 1080 Tiでも測定した。
将棋AI用の10ブロックのResNetのモデルで、cuDNNのTensorCoreは有効にしている。
cuBLASは、価値ネットワークの出力の全結合層(中間ノード数256)で使用している。
結果は以下の通り。
10万局面推論の平均時間
バッチサイズ128
単位：ms
| GPU | 有効 | 無効 | 比率(無効/有効) |
|---|---|---|---|
| Titan V | 7.041 | 6.850 | 1.03 |
| 1080 Ti | 16.178 | 16.183 | 1.00 |
精度比較
| 無効 | 有効 | |
|---|---|---|
| 指し手一致率 | 0.43799 | 0.43799 |
| 勝敗一致率 | 0.708287 | 0.708287 |
| 評価値誤差(MSE) | 0.0199915 | 0.0199915 |
考察
cuBLASでTensorCoreを有効にしてもほとんど速くなっていない（誤差の範囲）。
1080 TiではTensorCoreを有効に設定しても、結果は変わっていない。
TensorCoreが有効になっているかプロファイルで確認したところ、
TensorCore無効／有効で同じアルゴリズム(volta_sgemm_32x32_sliced1x4_tn))が選択されていた。
CUBLAS_GEMM_DEFAULT_TENSOR_OPを指定しても、必ずTensorCoreが使われるわけではないことが確認できた。
FP16の場合は、TensorCoreが使われる可能性があるので、次はFP16にして試す予定。
追記
FP16にして測定してみた。
前回書いた方法でFP16化したが、計算結果がおかしくなってしまった。
バグとれていないので結果があやしいが、とりあえず速度だけ測定すると、8.438msとなり、FP32よりも遅くなった。
プロファイラで確認すると、
TensorCoreが使われているように見えない。
FP16化は苦労しそうなので、とりあえずFP32でもcuDNNでTensorCoreが使えているので調査はここまでにする。
