WMMA APIのfragmentについて

CUDA Advent Calendar 2019 1日目の記事です.
あわよくば今書いている修論に転記してしまおうという寸法です.

WMMA APIのfragmentって何?

皆さんもTensorコアを毎日使っていると思うので知っていると思いますが,WMMA APIはTensorコアを使うためのAPIです.
このWMMA APIでどの様にTensorコアを使うかというと,

  1. メモリからレジスタに行列をコピー
  2. レジスタ上のデータをTensorコアに食わせて行列積和を計算(出力もレジスタ)
  3. レジスタからメモリに結果をコピー
といった感じです.
この手順で出てきたレジスタこそfragmentと呼ばれるものです.
まぁ,これは普通に行列積計算を行うときもレジスタブロッキングをしていると思うので今となっては普通のことですよね.
で,この記事はこのfragmentの構造についてのものです.

fragmentはいやらしい

VoltaからTuringでTensorコア自体の構造が変わったという話を聞いたことがあるかもしれません.

これに伴いfragmentの構造もVoltaとTuringでは異なります.

そもそもfragmentについてですが,コレ自体は単なる構造体です.
メンバ変数は行列要素をいくつかと,その個数の定数となります.
fragmentには行列積和計算 D=A x B + C のどの行列のfragmentかにより種類が異なります.
Aであればmatrix_a,Bであればmatrix_b,C, Dであればaccumulatorと呼ばれます.
また,HMMAすなわち浮動小数を対象としたWMMA APIではA, BはFP16,C, DはFP32 / FP16である必要があります.
Tensorコアでは内部の積和計算をFP32で行うため精度の劣化が少ないという特長があります.
で,このfragmentをどの様に使うかというとこんな感じ.
何をしているかは関数名を見れば一目瞭然ですよね.

で,ここからが本題で,このfragmentの構造が気になりませんか???
上で言ったとおりfragmentは普通の構造体なので中身をすべてprintfすれば構造が分かってしまいます.
で,やってみましたというのが私が昔やったインターンの内容です.
この当時はまだTuringが発表されていなかったため平和でした.
Turingが出てfragmentの構造を解析したところVoltaとは全く異なる構造が出てきてうわーって気持ちです.
WMMA APIは絶賛開発中なAPIなため機能が少なく,自分でfragmentを手動でいじる関数を作りながらでないと私がやりたいことは性能が出ないのですが,VolatとTuringの双方に対応しようとすると別々に対応しなくてはならずしんどいです.
まぁ,性能が良くならいいんですけどね,良くなるなら.

どう違うかを視覚的に表現するとこんな感じ.
たとえばこれはmatrix_aでcol_majorなVoltaとTuringのfragmentです.
見方としては,縦方向にWarp内のスレッド番号(一番上が0,下が31),横方向にfragment内にある行列要素配列となっています.

おわり

悪口を言いましたが,Tensorコアはいいやつです.
個人的には足しこみを倍精度にしてほしいのとWMMA APIで無駄な0埋めなしにベクトルを読み込む関数がほしい(これは自分で作った)のとさらなる高速化を望みます.

カテゴリー:CUDA
記事作成日:2019-12-01