🧠

Tensorコアを使ってみた

Techblog管理者

2018/10/10に公開

GPU

tech

こんにちは。アルバイトの大友です。

TensorコアのWMMA APIを使っている人があまりいなかったため、6月中はインターンとして、7月からはアルバイトとしてその使い方や性能を調べていました。

この記事はその成果をまとめたものです。

Tensorコアを使うことでFP16のSIMD計算(f16x2)に比べ密行列積を5倍程度高速化できました。

 TensorコアとはNVIDIA Voltaアーキテクチャから採用されたTensorコアは2つの4 \times 4のFP16行列の積を1サイクルで計算し、その累積和をFP16/FP32で取ることができる計算ユニットです。

cuBLAS, cuDNNなどのライブラリではCUDA 9からTensorコアを利用できます。

 WMMA APIを用いた行列積計算CUDA 9ではWMMA (Warp Matrix Multiply Accumulate) と呼ばれるTensorコアを使用してGEMM計算を行うためのC++ APIが用意されています。

このAPIでは16 \times 16など決められた大きさの行列をfragmentと呼ばれる構造体にスレッドごとに分割し、1ワープ(32スレッド)が協調してその行列積を計算します。

 行列計算の流れ行列A,B,C \in R^{16 \times 16}に対しC \Leftarrow A \times Bを計算する流れは次のようになります。
各スレッドがメモリからA,Bそれぞれの一部をfragmentとして読み込む
各スレッドのfragmentを用いて行列積を計算 (計算結果Cは同じくfragmentとして各スレッドが一部ずつ保持)
各スレッドがCのfragmentをメモリに書き込む

 WMMA APIを使用したプログラム16 \times 16の2行列A,Bの積Cを計算するCUDAのカーネル関数は次のように書けます。(CUDA 9.2.148)
__global__ void matmal_16x16(const half* const a_ptr,const half* const b_ptr,half* const c_ptr){
    // A,B,Cのfragment
    nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, 16, 16, 16, half, nvcuda::wmma::col_major> a_frag;
    nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, 16, 16, 16, half, nvcuda::wmma::col_major> b_frag;
    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, 16, 16, 16, half> c_frag;

    // Cのfragmentを0で初期化
    nvcuda::wmma::fill_fragment(c_frag, __float2half(.0f));

    // A,Bをメモリからfragmentに読み込み
    nvcuda::wmma::load_matrix_sync(a_frag, a_ptr, 16);
    nvcuda::wmma::load_matrix_sync(b_frag, b_ptr, 16);

    // C ← A x B + C
    nvcuda::wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

    // Cのfragmentの中身をメモリに書き出し
    nvcuda::wmma::store_matrix_sync(c_ptr, c_frag, 16, nvcuda::wmma::mem_col_major);
}
このAPIでは1ワープで16 \times 16の行列の積を計算するため、カーネル関数は次のように呼び出します。
constexpr unsigned int warpSize = 32;
matmal_16x16<<<1,warpSize>>>(dA, dB, dC);

 nvcuda::wmma::fragment構造体各スレッドが保持するfragmentは
template<typename Use, int m, int n, int k, typename T, typename Layout=void> class fragment;
と定義されており、それぞれのテンプレート引数は次のような役割を担っています。

Use : GEMM計算 D \Leftarrow A \times B + C のA, B, C, Dどの行列のfragmentか

Aの場合nvcuda::wmma::matrix_a

Bの場合nvcuda::wmma::matrix_b

C, Dの場合nvcuda::wmma::accumulator


m, n, k : Tensorコアで計算する行列積の行列の大きさ

ただし、(m, n, k)は (16, 16, 16), (32, 8, 16), (8, 32, 16)のいずれか

A : m \times k

B : k \times n

C, D : m \times n


T : fragmentの型

A, B : half

C, D : half / float


Layout : 列優先か行優先か
列優先 : nvcuda::wmma::col_major
行優先 : nvcuda::wmma::row_major


 メンバ変数
x : fragmentの要素配列

num_elements : fragmentの要素数

 nvcuda::wmma::fill_fragment関数void fill_fragment(fragment<...> &a, const T& v);
nvcuda::wmma::fragment aの全要素にvを代入

 nvcuda::wmma::load_matrix_sync関数void load_matrix_sync(fragment<...> &a, const T* mptr, unsigned ldm);
void load_matrix_sync(fragment<...> &a, const T* mptr, unsigned ldm, layout_t layout);
fragmentをメモリから読み込む

 引数
a : 読み込み先fragment

mptr : 読み込み元ポインタ

ldm : 行列全体のLeading dimension

layout : 列優先の場合はnvcuda::wmma::mem_col_major, 行優先の場合はnvcuda::wmma::mem_row_major

 制約
mptrが128-bit境界である必要あり (Alignment制約)

ldmが16 bytesの倍数である必要あり (halfでは8, floatでは4) (Leading dimension制約)

 nvcuda::wmma::store_matrix_sync関数void store_matrix_sync(T* mptr, const fragment<...> &a, unsigned ldm, layout_t layout);
fragmentをメモリに書き出す

 引数
mptr : 書き出し先ポインタ

a : 書き出し元fragment (nvcuda::wmma::accumulatorのみ)

ldm : 書き出し先行列のLeading dimension

layout : 列優先の場合はnvcuda::wmma::mem_col_major, 行優先の場合はnvcuda::wmma::mem_row_major

 制約nvcuda::wmma::load_matrix_sync と同様の制約と未定義動作あり

 nvcuda::wmma::mma_sync関数void mma_sync(fragment<...> &d, const fragment<...> &a, const fragment<...> &b, const fragment<...> &c, bool satf=false);
Tensorコアを用いたGEMM計算

 引数
d, a, b, c : GEMM計算 d \Leftarrow a \times b + c の各fragment

satf : fragmentの要素が+-Infinity, NaNとなった場合に有限値に修正するか否か

 任意の大きさの行列積計算WMMA APIでは決められた大きさの行列積しか計算できませんが、行列積を分解して考えることで任意の大きさの行列積を計算することができます。

行列A,Bの積Cを計算する流れは次のようになります。
行列A,B,Cを16 \times 16行列のブロックA_{i,j},B_{i,j},C_{i,j}に分割する。(端数部分は0埋め)
上図ではC_{1,1} = A_{1,0} \times B_{0,1} + A_{1,1} \times B_{1,1}と計算できる。

このようにC_{i,j} = \sum_k A_{i,k} \times B_{k,j}と計算することができる。

A_{i,k} \times B_{k,j}は16 \times 16の2行列の積のため、Tensorコアを用いて計算する。
2をCのすべてのブロックに対して行う。

 WMMA APIを使用するにあたって上述したとおり、nvcuda::wmma::load_matrix_sync関数とnvcuda::wmma::store_matrix_sync関数にはメモリのAlignment制約とLeading dimension制約があり、

Globalメモリにある任意の大きさの行列のGEMM計算を行うにはこの制約に対応しなければなりません。 そこでSharedメモリを用いることで対応します。
fragmentとして読み込むGlobalメモリの領域をSharedメモリにコピー
コピーしたSharedメモリからnvcuda::wmma::load_matrix_sync 関数でfragmentに読み込み

nvcuda::wmma::mma_sync 関数でGEMM計算
計算結果のfragmentをnvcuda::wmma::store_matrix_sync 関数でSharedメモリに書き出し
書き出したSharedメモリからGlobalメモリに書き出し

 注意点SharedメモリであればAlignment制約が満たされるわけではないので、必要ならば__align__([n byte])で境界を指定しなければならない。

 性能調査
 実験方法行列A,B,C \in \mathrm{half}^{N \times N}に対しC = A \times Bを計算
Tensorコアを使用した場合としなかった場合(f16x2を用いた場合)で計算速度を比較
それぞれ5回計算を行う
実験コードはtensorcore/matmul_evalにあります

 実験環境CPU : Intel Core i9-7900X
GPU : NVIDIA Titan V
RAM : 64GB
OS : Ubuntu 16.04

 実験結果 
Tensorコアを使用した場合、使用しなかった場合に比べてN \geq 512で5倍程度高速化されました。

 nvcuda::wmma::fragment構造体の調査行列がどのようにfragmentとしてワープ内で保持されているのかをprintfですべて標準出力して調査しました。

 nvcuda::wmma::matrix_a, nvcuda::wmma::matrix_bの場合行列M \in \mathrm{half}^{16 \times 16}を(m, n, k) = (16, 16, 16),nvcuda::wmma::col_majorなnvcuda::wmma::matrix_a,nvcuda::wmma::matrix_bそれぞれのfragmentにloadする場合を考えます。

MをM_{i,j} \in \mathrm{half}^{4 \times 4}のブロックに分割し

\left(\begin{matrix}
M_{0,0} & M_{0,1} & M_{0,2} & M_{0,3} \\
M_{1,0} & M_{1,1} & M_{1,2} & M_{1,3} \\
M_{2,0} & M_{2,1} & M_{2,2} & M_{2,3} \\
M_{3,0} & M_{3,1} & M_{3,2} & M_{3,3}
\end{matrix}\right)
と表すとthreadIdx.x=i の

nvcuda::wmma::matrix_a のfragmentは

\left(\begin{matrix}
M^T_{0,0} & M^T_{1,0} & M^T_{2,0} & M^T_{3,0} \\
M^T_{0,0} & M^T_{1,0} & M^T_{2,0} & M^T_{3,0} \\
M^T_{0,2} & M^T_{1,2} & M^T_{2,2} & M^T_{3,2} \\
M^T_{0,2} & M^T_{1,2} & M^T_{2,2} & M^T_{3,2} \\
M^T_{0,1} & M^T_{1,1} & M^T_{2,1} & M^T_{3,1} \\
M^T_{0,1} & M^T_{1,1} & M^T_{2,1} & M^T_{3,1} \\
M^T_{0,3} & M^T_{1,3} & M^T_{2,3} & M^T_{3,3} \\
M^T_{0,3} & M^T_{1,3} & M^T_{2,3} & M^T_{3,3}
\end{matrix}\right)

nvcuda::wmma::matrix_b のfragmentは

\left(\begin{matrix}
M^T_{0,0} & M^T_{0,1} & M^T_{0,2} & M^T_{0,3} \\
M^T_{2,0} & M^T_{2,1} & M^T_{2,2} & M^T_{2,3} \\
M^T_{0,0} & M^T_{0,1} & M^T_{0,2} & M^T_{0,3} \\
M^T_{2,0} & M^T_{2,1} & M^T_{2,2} & M^T_{2,3} \\
M^T_{1,0} & M^T_{1,1} & M^T_{1,2} & M^T_{1,3} \\
M^T_{3,0} & M^T_{3,1} & M^T_{3,2} & M^T_{3,3} \\
M^T_{1,0} & M^T_{1,1} & M^T_{1,2} & M^T_{1,3} \\
M^T_{3,0} & M^T_{3,1} & M^T_{3,2} & M^T_{3,3}
\end{matrix}\right)
で表される行列のi+1行目となります。
これを可視化すると 
となります。

 nvcuda::wmma::accumulatorの場合行列M \in \mathrm{half/float}^{16 \times 16}は(m, n, k) = (16, 16, 16),nvcuda::wmma::col_majorなnvcuda::wmma::accumulatorのfragmentでは

MをM_{i,j} \in \mathrm{half/float}^{4 \times 1}のブロックに分割し

\left(\begin{matrix}

M_{0,0} & M_{0,1} & M_{0,2} & M_{0,3} & M_{0,4} & M_{0,5} & M_{0,6} & M_{0,7} & M_{0,8} & M_{0,9} & M_{0,10} & M_{0,11} & M_{0,12} & M_{0,13} & M_{0,14} & M_{0,15} \\
M_{1,0} & M_{1,1} & M_{1,2} & M_{1,3} & M_{1,4} & M_{1,5} & M_{1,6} & M_{1,7} & M_{1,8} & M_{1,9} & M_{1,10} & M_{1,11} & M_{1,12} & M_{1,13} & M_{1,14} & M_{1,15} \\
M_{2,0} & M_{2,1} & M_{2,2} & M_{2,3} & M_{2,4} & M_{2,5} & M_{2,6} & M_{2,7} & M_{2,8} & M_{2,9} & M_{2,10} & M_{2,11} & M_{2,12} & M_{2,13} & M_{2,14} & M_{2,15} \\
M_{3,0} & M_{3,1} & M_{3,2} & M_{3,3} & M_{3,4} & M_{3,5} & M_{3,6} & M_{3,7} & M_{3,8} & M_{3,9} & M_{3,10} & M_{3,11} & M_{3,12} & M_{3,13} & M_{3,14} & M_{3,15}
\end{matrix}\right)
と表すとthreadIdx.x=iでは

\left(\begin{matrix}
M_{0,0} & M_{0,1} & M_{0,2} & M_{0,3} & M_{0,4} & M_{0,5} & M_{0,6} & M_{0,7} \\
M_{2,0} & M_{2,1} & M_{2,2} & M_{2,3} & M_{2,4} & M_{2,5} & M_{2,6} & M_{2,7} \\
M_{0,8} & M_{0,9} & M_{0,10} & M_{0,11} & M_{0,12} & M_{0,13} & M_{0,14} & M_{0,15} \\
M_{2,8} & M_{2,9} & M_{2,10} & M_{2,11} & M_{2,12} & M_{2,13} & M_{2,14} & M_{2,15} \\
M_{1,0} & M_{1,1} & M_{1,2} & M_{1,3} & M_{1,4} & M_{1,5} & M_{1,6} & M_{1,7} \\
M_{3,0} & M_{3,1} & M_{3,2} & M_{3,3} & M_{3,4} & M_{3,5} & M_{3,6} & M_{3,7} \\
M_{1,8} & M_{1,9} & M_{1,10} & M_{1,11} & M_{1,12} & M_{1,13} & M_{1,14} & M_{1,15} \\
M_{3,8} & M_{3,9} & M_{3,10} & M_{3,11} & M_{3,12} & M_{3,13} & M_{3,14} & M_{3,15}
\end{matrix}\right)
で表される行列のi+1行目となります。
これを可視化すると 
となります。

 行優先の場合Globalメモリに行列A,Bが行優先で置かれている場合、単純にfragmentに読み込む際に転置して読み込んでいるわけではありません。 行優先か列優先かでfragmentの中身が異なるため、nvcuda::wmma::mma_sync関数に対応するPTX命令であるwmma.load命令は
wmma.mma.sync.alayout.blayout.shape.dtype.ctype{.satfinite} d, a, b, c;

.alayout = {.row, .col};
.blayout = {.row, .col};
.shape  =  {.m16n16k16, .m8n32k16, .m32n8k16};
.ctype   = {.f16, .f32};
.dtype   = {.f16, .f32};
という構造となっており(3)、fragment a,bが行優先か列優先かを指定する必要があります。
C,Dに関しては行優先か列優先かを指定する必要はなく、実際列優先か行優先かでfragmentに差は見られませんでした。

 nvcuda::wmma::load_matrix_syncの調査Warp内の各スレッドでのnvcuda::wmma::fragment構造体の中身がわかったので、nvcuda::wmma::load_matrix_sync関数を使わずに自前でfragmentを読み込んだ場合と速度を比較しました。
行列M \in \mathrm{half}^{16 \times 16}をnvcuda::wmma::matrix_aとして読み込むだけのカーネルを実行
カーネル内で2^{30}回Globalメモリから読み込みを実行
nvprofでカーネルの実行時間を測定
実験コードはtensorcore/load_evalにあります

 結果

関数
実行時間



load_matrix_sync 関数
91 us

自作load関数
77670 us

高速化を余り考えずに書いたと言え、自作load関数に比べてnvcuda::wmma::load_matrix_sync 関数が850倍程度高速という結果になりました。考察NVIDIA Visual Profilerで実行されたSASSコードを見たところ、nvcuda::wmma::load_matrix_sync関数でも汎用的なメモリ読み込み命令であるLDG命令が使われているようでした。

読み込みアドレスの計算と実際の読み込み命令の実行順などが工夫されているのかもしれません。

 まとめTensorコアはWMMA APIを用いることで簡潔に利用することができました。

WMMA APIのload_matrix_sync,store_matrix_sync,mma_sync関数はほとんど単純にPTXの命令に置き換えられるだけなためレイヤーは低く、使用の自由度は高いと考えられます。

性能面ではTensorコアを使用することでf16x2を使用した場合に比べFP16密行列積を高速に計算できることが確認できました。

 謝辞吉藤さんにはインターン及びアルバイトでCUDAやコードの書き方について指導していただきました。 ありがとうございました。

 参考文献NVIDIA Developer Blog - Programming Tensor Cores in CUDA 9
CUDA Toolkit Document - CUDA C Programming Guide (Warp matrix functions)
CUDA Toolkit Document - Parallel Thread Execution ISA (Warp Level Matrix Multiply-Accumulate Instructions)
GitHub - parallel-forall/code-samples
VOLTA AND TURING: ARCHITECTURE AND PERFORMANCE OPTIMIZATION - Akira Naruse, Developer Technology, 2018//14

 利用許諾・ライセンス
本記事に含まれるあらゆる文章および画像は、クリエイティブ・コモンズ 表示-継承 4.0 国際(CC-BY-SA 4.0 International)ライセンス（帰属表示/attribution: Fixstars Corporation）の下で利用可能です。