Open6

CUDAで行列演算

so298so298

CUDAの最適化手法についてきちんと理解しておきたいと思ったので,行列計算を題材に最適化を行ってみる
成果物は https://github.com/so298/cuda-gemm-optimization-practice に置いていく

参考資料として

実行環境は NVIDIA A100-SXM4-40GB で行う

so298so298

ナイーブな実装

とりあえず最適化はほぼ何も行わないナイーブな実装を行った (01-naive)
https://github.com/so298/cuda-gemm-optimization-practice/blob/main/src/01-naive.cu

また,行列の各要素の計算をするスレッドの割り付け方法を変えたバージョンも用意した (01-2-non-coalesced)
https://github.com/so298/cuda-gemm-optimization-practice/blob/main/src/01-2-non-coalesced.cu

実行結果は以下の通り

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倍くらい速い

so298so298

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].これが原因となって性能差が生じているのだと思う.

脚注
  1. 9.2.1.4. Strided Accesses | CUDA C++ Best Practice Guide ↩︎

so298so298

GFLOPs の計算について

Throughput (GFLOPs) は 2 NMK / \text{elapsed\_time} で計算している.出力の行列1要素を計算するのにK回の積算とK-1回の加算が必要で,出力行列の要素数はNM個だからこうなる.

so298so298

block tilingの実装

行列を小行列(tile)に分解して,計算を行うtileをthread block内で高速にアクセスできるメモリ領域であるshared memoryにコピーしてメモリアクセスにかかる時間を削減する方法を block tiling という.

block tilingの実装
https://github.com/so298/cuda-gemm-optimization-practice/blob/main/src/02-block-tiling.cu

実行結果

$ ./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倍くらい速くなった.

so298so298

ハードウェアの限界を引き出すとどの程度の速度が出るのかは知っておきたい

まず,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

実装は以下
https://github.com/so298/cuda-gemm-optimization-practice/blob/main/src/99-cublas.cu

自分の実装は行指向(row-major)だがcuBLASは列指向(column-major)となっている.行指向で行った結果と一致させるために,列指向の計算を行指向の計算を転置して行ったものとみなせば良い.