CUDAで行列演算

CUDAの最適化手法についてきちんと理解しておきたいと思ったので,行列計算を題材に最適化を行ってみる
成果物は https://github.com/so298/cuda-gemm-optimization-practice に置いていく
参考資料として
- https://leimao.github.io/article/CUDA-Matrix-Multiplication-Optimization/
-
https://siboehm.com/articles/22/CUDA-MMM
ここらへんの記事がかなり参考になりそうだと感じたので,似たような方針で最適化を行いたい
実行環境は NVIDIA A100-SXM4-40GB
で行う

ナイーブな実装
とりあえず最適化はほぼ何も行わないナイーブな実装を行った (01-naive)
また,行列の各要素の計算をするスレッドの割り付け方法を変えたバージョンも用意した (01-2-non-coalesced)
実行結果は以下の通り
01-naive
$ ./exe/01-naive -n 2048 -m 2048 -k 2048 -i 100
N: 2048
K: 2048
M: 2048
num_iter: 100
Elapsed time: 630 ms
Elapsed time per iteration: 6.3 ms
Throughput: 2726.96 GFLOPs
01-2-non-coalesced
$ ./exe/01-2-non-coalesced -n 2048 -m 2048 -k 2048 -i 100
N: 2048
K: 2048
M: 2048
num_iter: 100
Elapsed time: 5850 ms
Elapsed time per iteration: 58.5 ms
Throughput: 293.673 GFLOPs
01-naive
のほうが10倍くらい速い

01-2
のほうが遅くなってしまった理由の考察
// size_t r = idx / M
size_t r = idx % N; // 01-2
...
sum += A[r * K + k] * B[k * M + c];
01の方では size_t r = idx / M
としていたのに対し, 01-2 の方では size_t r = idx % N
としており,隣接するスレッドにおいて異なる行番号が割り当てられている.
この場合,行列A
の各要素へのアクセスの際に隣接するスレッドでストライドがK
のメモリアクセスが生じている.
グローバルメモリにストライドアクセスする際にはストライドを1にしないとメモリアクセスの帯域が急速に低下するらしい[1].これが原因となって性能差が生じているのだと思う.

GFLOPs の計算について
Throughput (GFLOPs) は

block tilingの実装
行列を小行列(tile)に分解して,計算を行うtileをthread block内で高速にアクセスできるメモリ領域であるshared memoryにコピーしてメモリアクセスにかかる時間を削減する方法を block tiling という.
block tilingの実装
実行結果
$ ./exe/02-block-tiling -n 2048 -m 2048 -k 2048 -i 100
N: 2048
K: 2048
M: 2048
num_iter: 100
Elapsed time: 381 ms
Elapsed time per iteration: 3.81 ms
Throughput: 4509.15 GFLOPs
2726 GFLOPs -> 4509 GFLOPs で大体1.6倍くらい速くなった.

ハードウェアの限界を引き出すとどの程度の速度が出るのかは知っておきたい
まず,A100 40GBのピーク性能(FP32)は 19.5 TFLOPsらしい
また,cuBLASによる最適化された実装も用意して計測したところ次のようになった.
$ ./exe/99-cublas -n 2048 -m 2048 -k 2048 -i 100
N: 2048
K: 2048
M: 2048
num_iter: 100
Elapsed time: 98 ms
Elapsed time per iteration: 0.98 ms
Throughput: 17530.5 GFLOPs
17530.5 GFLOPs = 17.5 TFLOPs
実装は以下
自分の実装は行指向(row-major)だがcuBLASは列指向(column-major)となっている.行指向で行った結果と一致させるために,列指向の計算を行指向の計算を転置して行ったものとみなせば良い.