前回の記事（行列乗算の最適化入門（マルチコア編） - よーる）では、CPU上でマルチスレッドの理論性能の89%以上（1.1 TFLOPS）を出せる簡単なコードを紹介しました。今回はさらに高速化するため、GPGPUを使ってみます。 GPGPUとしては、前回セットアップしたTITAN Vを使いました（Ubuntu 24.04 LTSでTITAN Vを使えるようにした - よーる）。

試行錯誤した結果、TITAN Vの理論性能（約7.45 TFLOPS）の68%以上である5.1 TFLOPSを達成する、非常に簡単なソースコードを作ることができました。理論性能に対して7割未満というのは低いようにも見えますが、cuBLASを使った場合の性能は4.9 TFLOPSであり、それより高い性能です（同じ計算をしているわけではないのでcuBLASの性能と単純に比較することはできません）。かなり頑張っていると思われるコードでも理論性能に対して60%程度（参考文献[1]。全く条件が違うので比較できません）とかなので、GPGPUでは理論性能近くの性能を出すことはほぼ不可能なのでしょう。（同じ計算をしているわけではないのでこのような推論は危ういですが、）同じ程度の性能となっていることから、現実的に可能な範囲でほぼ限界まで性能を出しているとは言えそうです。

mixbench

mixbench（GitHub - ekondis/mixbench: A GPU benchmark tool for evaluating GPUs and CPUs on mixed operational intensity kernels (CUDA, OpenCL, HIP, SYCL, OpenMP)）は、GPUのベンチマークです。これを用いてTITAN Vの性能の限界を調べたところ、以下のようになりました。

FP32: 13.1 TFLOPS, 639 GB/s
FP64: 6.7 TFLOPS, 645 GB/s
Int: 13.1 TIOPS, 639 GB/s

理論性能の7.45 TFLOPSを出すのは無理みたいです。ただ、定格周波数で動作していると6.144 TFLOPSしか出ないはずなので、多少なりともブースト周波数で動いていることがわかります。

また、バンド幅は理論性能（3072 bit×850 MHz×2 (DDR) = 652.8 GB/s）に非常に近い値が出ています。メモリインタフェースは公称値で動いているようです。

実験コード

main.cu↓

#include <iostream>
#include <chrono>
#include <random>

#include <cuda_runtime.h>
#include "kernel.cu"

int main( int argc, char* argv[] ) {
    cudaEvent_t device_start; cudaEventCreate( &device_start );
    cudaEvent_t device_finish; cudaEventCreate( &device_finish );

    std::mt19937_64 mt;
    std::uniform_real_distribution dist( -1.0, 1.0 );

    double* a = (double*)malloc( N*N*sizeof(double) );
    double* b = (double*)malloc( N*N*sizeof(double) );
    double* c = (double*)malloc( N*N*sizeof(double) );
    for( int i = 0; i < N*N; ++i ) {
        a[i] = dist( mt );
        b[i] = dist( mt );
        c[i] = dist( mt );
    }

    std::cout << "initialized." << std::endl;
    const auto start = std::chrono::system_clock::now();

    double* device_a; cudaMalloc( (void**)&device_a, N*N*sizeof(double) );
    double* device_b; cudaMalloc( (void**)&device_b, N*N*sizeof(double) );
    double* device_c; cudaMalloc( (void**)&device_c, N*N*sizeof(double) );

    cudaMemcpy( device_a, a, N*N*sizeof(double), cudaMemcpyHostToDevice );
    cudaMemcpy( device_b, b, N*N*sizeof(double), cudaMemcpyHostToDevice );
    cudaMemcpy( device_c, c, N*N*sizeof(double), cudaMemcpyHostToDevice );

    cudaEventRecord( device_start );

    call_kernel( device_a, device_b, device_c );

    cudaEventRecord( device_finish );
    cudaEventSynchronize( device_finish );

    cudaMemcpy( c, device_c, N*N*sizeof(double), cudaMemcpyDeviceToHost );

    const auto finish = std::chrono::system_clock::now();

    float ms; cudaEventElapsedTime( &ms, device_start, device_finish );
    const double device_time = ms * 1e-3;

    const double host_time = std::chrono::duration_cast<std::chrono::nanoseconds>( finish - start ).count() * 1e-9;

    std::cout << device_time << " seconds (device), " << N*N*N*flop_per_fma/device_time * 1e-9 << " GFLOPS, " << N*N*N*insn_per_fma/device_time << " FPIns/s" << std::endl;
    std::cout << host_time << " second (host), " << N*N*N*flop_per_fma/host_time * 1e-9 << " GFLOPS" << std::endl;
}

kernel.cu↓

static constexpr std::size_t N = 4096;
static constexpr double flop_per_fma = 2.0;
static constexpr double insn_per_fma = 1.0 / 32;

static constexpr std::size_t ThreadsPerBlockJ = 128;
static constexpr std::size_t ThreadsPerBlockI = 1;

static constexpr std::size_t RegisterBlockJ = 4;
static constexpr std::size_t RegisterBlockI = 16;

static constexpr std::size_t SizeJ = RegisterBlockJ * ThreadsPerBlockJ;
static constexpr std::size_t SizeI = RegisterBlockI * ThreadsPerBlockI;

__global__ void kernel( double* a, double* b, double* c ) {
    int i0 = blockIdx.y * SizeI;
    int j0 = blockIdx.x * SizeJ;
    int i2 = threadIdx.y;
    int j2 = threadIdx.x;

    double sum[RegisterBlockJ][RegisterBlockI];
    for( int i1 = 0; i1 < RegisterBlockI; ++i1 )
    for( int j1 = 0; j1 < RegisterBlockJ; ++j1 )
    {
        int i = i0 + (i1 * ThreadsPerBlockI) + i2;
        int j = j0 + (j1 * ThreadsPerBlockJ) + j2;
        sum[j1][i1] = c[i*N+j];
    }

    for( int k = 0; k < N; ++k )
    for( int i1 = 0; i1 < RegisterBlockI; ++i1 )
    for( int j1 = 0; j1 < RegisterBlockJ; ++j1 )
    {
        int i = i0 + (i1 * ThreadsPerBlockI) + i2;
        int j = j0 + (j1 * ThreadsPerBlockJ) + j2;
        sum[j1][i1] = fma( a[i*N+k], b[k*N+j], sum[j1][i1] );
    }

    for( int i1 = 0; i1 < RegisterBlockI; ++i1 )
    for( int j1 = 0; j1 < RegisterBlockJ; ++j1 )
    {
        int i = i0 + (i1 * ThreadsPerBlockI) + i2;
        int j = j0 + (j1 * ThreadsPerBlockJ) + j2;
        c[i*N+j] = sum[j1][i1];
    }
}

void call_kernel( double* device_a, double* device_b, double* device_c ) {
    dim3 grid( N / SizeJ, N / SizeI );
    dim3 block( ThreadsPerBlockJ, ThreadsPerBlockI );
    kernel<<< grid, block >>>( device_a, device_b, device_c );
}

コンパイル：nvcc main.cu -arch compute_70 -code sm_70

なお、ThreadsPerBlockJ、ThreadsPerBlockI、RegisterBlockJ、RegisterBlockI、の四つのパラメータは、おおよそ全部の二冪の組み合わせを調べて、もっともよかったものを採用しました。

実験結果

上記コードを動かした結果、平均的な実行時間は0.0269秒でした。その間に $4096^3$ 回のFMA演算（2FLOP）が行われるので、計算速度は5.1 TFLOPSです。なお、最速の場合には0.02619秒で完了しました（計算速度は5.248 TFLOPS）。

なお、上記結果はGPU内で計算を開始してから計算を終了するまでにかかった時間です。実際にはこれに加えて、CPUからGPUへの転送やGPUからCPUへの転送にも時間がかかります。その時間を含めて考えると、実行時間は約0.069秒で、計算速度に換算すると2.0 TFLOPS程度になります。

GPGPUをフル活用する場合、この転送にかかる時間を他の計算で隠蔽します。しかし、そのような工夫をせずとも、CPUで計算した場合（前回の記事で示したコードの場合、Ryzen 9 7950X上での速度は1.1 TFLOPS程度でした）の二倍程度の速度であるということになります。さすがにGPGPUのような計算に特化したプロセッサの威力といったところです。

cuBLASの場合

以下のコード（CUBLASを使ってみたぞ - 非平衡日常を参考にしました）でcuBLASを動かしたところ、4.9 TFLOPS程度でした。コンパイルは、nvcc main.cu -arch compute_70 -code sm_70 -lcublasで行いました。もしかしたら、性能を引き出せないような使い方をしている可能性もあります。また、そもそも私の書いたコードとは条件が異なる（私が書いたコードは一つの行列サイズに特化したコードになっており、アドレス計算が単純）という点に注意してください。

#include <iostream>
#include <chrono>
#include <random>

#include <cuda_runtime.h>
#include <cublas.h>

static constexpr std::size_t N = 4096;
static constexpr double flop_per_fma = 2.0;
static constexpr double insn_per_fma = 1.0 / 32;

int main( int argc, char* argv[] ) {
    cudaEvent_t device_start; cudaEventCreate( &device_start );
    cudaEvent_t device_finish; cudaEventCreate( &device_finish );

    std::mt19937_64 mt;
    std::uniform_real_distribution dist( -1.0, 1.0 );

    double* a = (double*)malloc( N*N*sizeof(double) );
    double* b = (double*)malloc( N*N*sizeof(double) );
    double* c = (double*)malloc( N*N*sizeof(double) );
    for( int i = 0; i < N*N; ++i ) {
        a[i] = dist( mt );
        b[i] = dist( mt );
        c[i] = dist( mt );
    }

    std::cout << "initialized." << std::endl;
    const auto start = std::chrono::system_clock::now();

    double* device_a; cublasAlloc( N*N, sizeof(double), (void**)&device_a ); cublasSetMatrix( N, N, sizeof(double), a, N, device_a, N );
    double* device_b; cublasAlloc( N*N, sizeof(double), (void**)&device_b ); cublasSetMatrix( N, N, sizeof(double), b, N, device_b, N );
    double* device_c; cublasAlloc( N*N, sizeof(double), (void**)&device_c ); cublasSetMatrix( N, N, sizeof(double), c, N, device_c, N );

    cudaEventRecord( device_start );

    cublasDgemm( 'N', 'N', N, N, N, 1.0, device_a, N, device_b, N, 1.0, device_c, N );

    cudaEventRecord( device_finish );
    cudaEventSynchronize( device_finish );

    cublasGetMatrix( N, N, sizeof(double), device_c, N, c, N );

    const auto finish = std::chrono::system_clock::now();

    float ms; cudaEventElapsedTime( &ms, device_start, device_finish );
    const double device_time = ms * 1e-3;

    const double host_time = std::chrono::duration_cast<std::chrono::nanoseconds>( finish - start ).count() * 1e-9;

    std::cout << device_time << " seconds (device), " << N*N*N*flop_per_fma/device_time * 1e-9 << " GFLOPS, " << N*N*N*insn_per_fma/device_time << " FPIns/s" << std::endl;
    std::cout << host_time << " second (host), " << N*N*N*flop_per_fma/host_time * 1e-9 << " GFLOPS" << std::endl;
}

何をやっているのか

最初に一回ロード、最後に一回ストア

CPUでのコードと同じようにc[i*N+j] = fma( a[i*N+k], b[k*N+j], c[i*N+j] );などと書いてしまうと全然性能が出ません。これはGPUはCPUと違い、ミュータブルなデータをL1キャッシュに保持しないためであると思われます。ロードは常にL2キャッシュに取りに行き、ストアは常にL2キャッシュに書き戻しに行ってしまいます。これをふせぐため、sum変数を作ってレジスタに載せています。

レジスタ ブロッキング

前々回の記事で指摘しましたが、行列積コードはレジスタブロッキングをしないとロード命令やストア命令ばかりになってしまいます。そこで、RegisterBlockI×RegisterBlockJのレジスタブロッキングを行いました。

ここのコードはソフトウェアパイプライン化を行っているように見えるかもしれませんが、ソフトウェアパイプライン化による性能向上はほとんどないと思われます。なぜなら、NVIDIAのGPUは、大量のスレッドを生成することでレイテンシを隠すための命令を供給し、高性能を達成するように設計されているからです。

ワープ単位の考慮

ワープ

NVIDIA用語の「ワープ（Warp）*1」は、スレッド32個*2のまとまりです。この32個組のスレッドは、ハードウェア上特別扱いされており、そのことを意識してプログラムを書くことでGPGPUの性能を引き出すことができます。具体的には、以下のような特別扱いがなされます。

32個のスレッドは足並みをそろえて実行される。分岐命令があると足並みがそろわなくなるので遅くなる（ワープダイバージェンス）
- よほど分岐しまくる場合を除けば、SIMD長32の命令セットでマスク実行しているイメージでよいはず
ワープ内のスレッドが足並みをそろえてロード命令を実行する時、すべてのスレッドが同じキャッシュライン（128 Byte）から読み出すなら速い（コアレスドアクセス）
- キャッシュライン境界にアラインされたfloat[32]を高速に読み出せるようになっている（floatなのは graphics processing unit だった名残）。
- キャッシュラインを読み出す数が多くなると遅くなる、と考えたほうがいいかもしれない（つまり、広い範囲からのGatherは遅い）
- アラインのあった連続アクセス（1キャッシュライン）・アラインがずれた連続アクセス（2キャッシュライン）・ブロードキャスト（1キャッシュライン）、の三種類はそれなりの速度で動くらしい。狭い範囲からのGather（つまり表引き）が速いのかはわからなかった
- アラインされたdouble[32]の読み出しも2キャッシュラインなのでそれなりの速度で動作しているっぽい

よって、「分岐せず」「32個のスレッドで連続するメモリ領域にアクセスする」ようにプログラムを書くべきです。行列乗算の場合、「分岐せず」は簡単に満たせるので、「32個のスレッドで連続するメモリ領域にアクセスする」ように気を付ければ問題ありません。

ワープ内スレッドがコアレスドアクセスするようにする

示したコードは、この点をちゃんと考慮したものになっています。 int j2 = threadIdx.x;の部分がそれで、スレッド番号が近いものが配列の近いところにアクセスするようになっています。レジスタブロッキングで配列の近くのインデクスを使っていないのもこのためです。 32個のスレッドのメモリアクセスパターンが連続になることが最優先なので、近くのインデクスを担当してはいけないのです。

一つのブロックに複数のワープを含める

ブロック

NVIDIA用語の「ブロック」とは、ワープのまとまりです。ブロック内のワープは、全て同じSM (Streaming Multiprocessor) に割り当てられて実行されます。 SMというのは、CPUで言うところの物理コアだと思えばよいでしょう。ブロックとSMは一対一対応というわけではなく、一つのSMには複数のブロックが割り当てられる可能性があります。

複数のワープをブロックにまとめる利点と欠点は以下の通りです。

👍同じSMで実行されるならL1キャッシュを共有できるため、同じデータを使うワープ同士が同じSMで実行されれば通信量を削減できる。こういう時はブロックにまとめたほうが良い
👎一つのSMには一定量（TITAN Vでは256 KiB）のレジスタしかないため、いくらでも多くのワープを一つのSMで実行できるわけではない。つまり、ワープあたりの使用レジスタ数を減らす必要が出てくる
👎同じデータを使うワープをまとめすぎると、そのデータが手に入らないという一つの事象が起こっただけでSM内の多くのワープが止まって、SMでする仕事が何もない状態になりがち

ブロックに含めるワープ

複数のワープをブロックに含めて1つのSMで実行することで、L1キャッシュを共有してデータ通信量を減らす目論見です（キャッシュブロッキング）。今回は、128スレッド、つまり4ワープをまとめて一つのブロックとするのが最適でした。ブロックに含めるべきスレッド数は128～256らしい（参考文献[2]）ので、それとだいたい合っています。

L1キャッシュ

読み取り専用であるとコンパイラが判断したデータのみがキャッシュされるらしいです（参考文献[3]）。行列乗算の入力（aとb）はちゃんと読み取り専用と判断されたようです。昔はこのようなことをするためにはシェアードメモリの読み書きを自分で書かないといけなかったようなので、とても楽になったと言えます。

ブロックのレジスタ使用量

1つのブロックは以下の量のレジスタを消費します。

a：1スレッドあたりdoubleをRegisterBlockI(16)個使用。128スレッドで16 KiB（128スレッド全員が同じ値を持っている）
b：1スレッドあたりdoubleをRegisterBlockJ(4)個使用。128スレッドで4 KiB
sum：1スレッドあたりdoubleをRegisterBlockI * RegisterBlockJ(64)個使用。128スレッドで64 KiB
その他、アドレス計算に必要なレジスタ。1 KiB程度？

合計で85 KiB程度でしょうか。つまり、TITAN VのSMに搭載された256 KiBのレジスタで3つのブロックを動かせるようなサイズ感になっています。

一つのSMでブロックを3つ動かせることは、データ待ちストールの時間を減らすことにかなり役立つようです。ブロックを大きくしすぎて一つのSMで1つのブロックしか動かせないような設定は明らかに性能が低くなりました。 2つのブロックしか動かせないような設定は3つのブロックが動かせる設定と比べてわずかに性能が低くなりました。これ以上ブロックを小さくしても性能はかえって下がってしまったため、3つのブロックを動かすのがベストのようです。

コードの解説まとめ

メモリに毎回アクセスするのではなく、アキュムレータレジスタを用意することで、通信量を大きく削減
4×16のレジスタブロッキング
コアレスアクセスになるようにする
128スレッドをブロックにまとめ、L1キャッシュブロッキング
読み出し専用であればコンパイラが自動的にL1キャッシュを使ってくれる
ブロック3つがSMに入るようにするのがベスト

メモリ配置

MN-Coreにおけるテンソルのメモリ配置表現の記法（参考：MN-Core におけるテンソルのメモリ配置レイアウト表現 - Preferred Networks Research & Development）を借りれば、

aのレイアウト：( (256_Grid, 1_Block, 16_Reg), (4096_Time); B@[128_Block] )
bのレイアウト：( (4096_Time), (8_Grid, 128_Block, 4_Reg) )
cのレイアウト：( (256_Grid, 1_Block, 16_Reg), (8_Grid, 128_Block, 4_Reg) )

となります。

まず、k軸は全部Time軸にしてしまったので、計算順序やスレッド同期の難しい問題に立ち向かわなくてよくなっています。 k軸に対してストリーミング的に処理してほしいなぁという感じのですが、明示的に同期しているわけではないので、そのようになっているかは不明です*3。 j軸はそのほとんどをBlock軸に使ってしまいました。本来はレジスタブロッキングは16×4などという極端な縦横比にすべきではないのですが、Block軸でj軸を使いすぎているので、ここでつじつまを合わせている感じになっています。 i軸は、SMのブロッキングには使わない方がよかったようです。

ストリーミング的に処理されているとすると、L2キャッシュにはkの16周ごとに512 KiBのaと512 KiBのbが読み込まれることになります。 TITAN VのL2キャッシュは4.5 MiBあるので、各スレッドの進捗のずれが64周くらいまでであれば対応できそうです。

各SMのL1キャッシュには、kの1周ごとに128 Byteのaと4 KiBのbが読み込まれることになります。そして、128 Byteのaはブロック内の128個のスレッドのレジスタにブロードキャストされ、合計で16 KiBのレジスタを消費します。 4 KiBのbはブロック内の128個のスレッドに分配され、合計で4 KiBのレジスタを消費します。

まとめ

GPGPUを使って行列乗算を行う、非常に簡単で高速なコードを紹介しました。読み出し専用であれば自動でキャッシュしてくれる機能が搭載されたため、シェアードメモリを使うコードをわざわざ書かなくてよくなって楽でした。性能は5.1 TFLOPSと理論性能の68%程度でしたが、他の実装と比較してみれば、これはそれほど悪い値ではないということがわかりました。

参考文献

[1] prime number, 「プログラムを高速化する話Ⅱ ～GPGPU編～」, p. 151, 2015.
- プログラムを高速化する話Ⅱ ～GPGPU編～ - Speaker Deck
[2] prime number, 「プログラムを高速化する話Ⅱ ～GPGPU編～」, p. 101, 2015.
- プログラムを高速化する話Ⅱ ～GPGPU編～ - Speaker Deck
[3] Fixstars Corporation, 「いまさら聞けない！CUDA高速化入門」, p. 62, 2021.
- いまさら聞けない！CUDA高速化入門 | PPT
- 右下の数字が62のスライド（39と47がないので60 of 90）

*1:発音は「ウォープ」の方が近いけれど、瞬間移動の「ワープ」と同じなので日本語的にはどちらでもよさそう

*2:とは限らない？

*3:そもそもこのカーネルは全ブロックを同時に動かすことは無理です。80個あるSMごとに3つのブロックしか動かせないので、同時には240個のブロックしか動かせません。一方、このカーネルは2048のブロックを持ちます。