TadaoYamaokaの日記

山岡忠夫Homeで公開しているプログラムの開発ネタを中心に書いていきます。

cuBLASでTensorCoreを有効にする

前回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に変換されて計算される。

参考:
cuBLAS :: CUDA Toolkit Documentation

速度比較

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が有効になっているかプロファイルで確認したところ、
f:id:TadaoYamaoka:20180929122105p:plain
TensorCore無効/有効で同じアルゴリズム(volta_sgemm_32x32_sliced1x4_tn))が選択されていた。

CUBLAS_GEMM_DEFAULT_TENSOR_OPを指定しても、必ずTensorCoreが使われるわけではないことが確認できた。
FP16の場合は、TensorCoreが使われる可能性があるので、次はFP16にして試す予定。

追記

FP16にして測定してみた。
前回書いた方法でFP16化したが、計算結果がおかしくなってしまった。
バグとれていないので結果があやしいが、とりあえず速度だけ測定すると、8.438msとなり、FP32よりも遅くなった。
プロファイラで確認すると、
f:id:TadaoYamaoka:20180929133854p:plain
TensorCoreが使われているように見えない。
FP16化は苦労しそうなので、とりあえずFP32でもcuDNNでTensorCoreが使えているので調査はここまでにする。