コンシューマー向けGPUは、倍精度演算性能こそ高くないものの、単精度でよければ驚異的な演算性能を持っています。 例えば、NVIDIA社のGeForce RTX 4090は、カタログスペックで82 TFLOPSもの性能があります。 カタログスペックに近い性能を出すことは困難ですが、cuBLASで行列積を行った場合、59 TFLOPSくらいは出すことができます。 初代地球シミュレータは（倍精度なので正しくない比較ですが）理論性能40.96 TFLOPS（LINPACK性能 35.86 TFLOPS＠2002年6月）なので、20年たったとはいえ、簡単に手に入る計算機でそれを超えるような演算性能が出せるのは本当に驚きです。
今回は、GeForce RTX 4090*1上で高速に動作する単精度行列積プログラムを作ってみます。 非常に単純なコードながら58 TFLOPSと、cuBLASに匹敵する性能のプログラムを作ることができました。
最初のコード
以前（行列乗算の最適化入門（GPGPU編） - よーる）に紹介した倍精度用のコードを単精度用に作り変えました。 いくつか定数を変更する必要がありましたが、だいたい30 TFLOPSくらいの性能が出ました。
static constexpr std::size_t N = 4096;
static constexpr float flop_per_fma = 2.0;
static constexpr float insn_per_fma = 1.0 / 32;
static constexpr std::size_t ThreadsPerBlockJ = 32;
static constexpr std::size_t ThreadsPerBlockI = 4;
static constexpr std::size_t RegisterBlockJ = 8;
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( float* a, float* b, float* c ) {
int i0 = blockIdx.y * SizeI;
int j0 = blockIdx.x * SizeJ;
int i2 = threadIdx.y;
int j2 = threadIdx.x;
float 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( float* device_a, float* device_b, float* device_c ) {
dim3 grid( N / SizeJ, N / SizeI );
dim3 block( ThreadsPerBlockJ, ThreadsPerBlockI );
kernel<<< grid, block >>>( device_a, device_b, device_c );
}
shared memory の利用
やはりL1キャッシュだけに頼るのではなく、shared memoryを使わないと高い性能を出すことは難しいようです。 shared memoryとL1キャッシュは同じハードウェアのはずですが、アドレス計算、アドレス変換、タグチェック、などにオーバーヘッドがあるのでしょうか。 次のようにshared memoryを使うように書き換えてみました。
static constexpr std::size_t N = 4096;
static constexpr float flop_per_fma = 2.0;
static constexpr float insn_per_fma = 1.0 / 32;
static constexpr std::size_t ThreadsPerBlockJ = 32;
static constexpr std::size_t ThreadsPerBlockI = 4;
static constexpr std::size_t RegisterBlockJ = 8;
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( float* a, float* b, float* c ) {
int i0 = blockIdx.y * SizeI;
int j0 = blockIdx.x * SizeJ;
int i2 = threadIdx.y;
int j2 = threadIdx.x;
float 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];
}
__shared__ float local_a[SizeI][32];
__shared__ float local_b[32][SizeJ];
for( int k0 = 0; k0 < N; k0 += 32 )
{
__syncthreads();
for( int i1 = 0; i1 < RegisterBlockI; ++i1 ) {
int il = i1 * ThreadsPerBlockI + i2;
int ig = i0 + il;
int kl = j2;
int kg = k0 + kl;
local_a[il][kl] = a[ig*N+kg];
}
for( int j1 = 0; j1 < RegisterBlockJ; ++j1 ) {
int jl = j1 * ThreadsPerBlockJ + j2;
int jg = j0 + jl;
int kl = i2;
int kg = k0 + kl;
local_b[kl+ 0][jl] = b[(kg+ 0)*N+jg];
local_b[kl+ 4][jl] = b[(kg+ 4)*N+jg];
local_b[kl+ 8][jl] = b[(kg+ 8)*N+jg];
local_b[kl+12][jl] = b[(kg+12)*N+jg];
local_b[kl+16][jl] = b[(kg+16)*N+jg];
local_b[kl+20][jl] = b[(kg+20)*N+jg];
local_b[kl+24][jl] = b[(kg+24)*N+jg];
local_b[kl+28][jl] = b[(kg+28)*N+jg];
}
__syncthreads();
for( int k1 = 0; k1 < 32; ++k1 )
for( int i1 = 0; i1 < RegisterBlockI; ++i1 )
for( int j1 = 0; j1 < RegisterBlockJ; ++j1 )
{
sum[j1][i1] = fma( local_a[(i1 * ThreadsPerBlockI) + i2][k1], local_b[k1][(j1 * ThreadsPerBlockJ) + j2], 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( float* device_a, float* device_b, float* device_c ) {
dim3 grid( N / SizeJ, N / SizeI );
dim3 block( ThreadsPerBlockJ, ThreadsPerBlockI );
kernel<<< grid, block >>>( device_a, device_b, device_c );
}
とりあえずshared memoryを利用したコードのお手並み拝見ということでThreadsPerBlockJが32でThreadsPerBlockIが8であることが前提のハードコーディングになっていますが、これで37 TFLOPSとかなり性能を上げることができました。
他のプログラマの知見の利用
私の知識だけではこの辺で策が尽きてきたので、他のプログラマのコードを参考にしてみたいと思います。 sgemm cuda でググって出てきた、GitHub - siboehm/SGEMM_CUDA: Fast CUDA matrix multiplication from scratch を参考にやっていきます。
最後に足す
今までのコードは、行列Cを最初に読み出して、内積演算を行って、最後に行列Cに戻す、というコードでした。 しかし、最初は0で初期化して、内積演算を行って、最後に行列Cに足しこむ、としたほうが速いようです。 これだけで39 TFLOPSまで高速化しました。 最初に読み出す方式は、カーネルが立ち上がった瞬間に大量のメモリ要求が出るのが良くない（最後に読み出す場合、ワープごとに進捗が異なるのでメモリ要求が出るタイミングがばらけるので良い）のでしょうか？
@@ -17,14 +17,7 @@ __global__ void kernel( float* a, float* b, float* c ) { int i2 = threadIdx.y; int j2 = threadIdx.x; - float 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]; - } + float sum[RegisterBlockJ][RegisterBlockI] = {}; __shared__ float local_a[SizeI][32]; __shared__ float local_b[32][SizeJ]; @@ -67,7 +60,7 @@ __global__ void kernel( float* a, float* b, float* c ) { { int i = i0 + (i1 * ThreadsPerBlockI) + i2; int j = j0 + (j1 * ThreadsPerBlockJ) + j2; - c[i*N+j] = sum[j1][i1]; + c[i*N+j] += sum[j1][i1]; } }
SharedMemBlockK の導入
キャッシュラインが128 Byteなこともあり、グローバルメモリへの要求を128 Byte単位にしたほうが効率的だと思ってk軸を32ごとに分割していたのですが、これを小さくするほうが良いようです。
SharedMemBlockKは8や32は悪く、16が良さそうでした。
この結果は、参考にしたウェブサイト（How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog）においてタイリングまで含めて最適化した場合のBK=16と一致しています。
static constexpr int N = 4096;
static constexpr double flop_per_fma = 2.0;
static constexpr double insn_per_fma = 1.0 / 32;
static constexpr int ThreadsPerBlockJ = 32;
static constexpr int ThreadsPerBlockI = 4;
static constexpr int SharedMemBlockK = 16;
static constexpr int RegisterBlockJ = 8;
static constexpr int RegisterBlockI = 16;
static constexpr int SizeJ = RegisterBlockJ * ThreadsPerBlockJ; // 256
static constexpr int SizeI = RegisterBlockI * ThreadsPerBlockI; // 64
static constexpr int LoopX = SizeI * SharedMemBlockK / ThreadsPerBlockJ / ThreadsPerBlockI; // 8
static constexpr int LoopY = SharedMemBlockK / ThreadsPerBlockI; // 4
static constexpr int LoopZ = SizeJ / ThreadsPerBlockJ; // 8
static constexpr int ThreadsPerBlock = ThreadsPerBlockI * ThreadsPerBlockJ; // 128
__global__ void kernel( float* a, float* b, float* c ) {
int i0 = blockIdx.y * SizeI;
int j0 = blockIdx.x * SizeJ;
int i2 = threadIdx.y;
int j2 = threadIdx.x;
int tid = threadIdx.y * ThreadsPerBlockJ + threadIdx.x;
float sum[RegisterBlockJ][RegisterBlockI] = {};
__shared__ float local_a[SizeI][SharedMemBlockK]; // 4 KB
__shared__ float local_b[SharedMemBlockK][SizeJ]; // 16 KB
for( int k0 = 0; k0 < N; k0 += SharedMemBlockK ) {
__syncthreads();
for( int x = 0; x < LoopX; ++x ) {
int il = ( x * ThreadsPerBlock + tid ) / SharedMemBlockK;
int kl = ( x * ThreadsPerBlock + tid ) % SharedMemBlockK;
int ig = i0 + il;
int kg = k0 + kl;
local_a[il][kl] = a[ig*N+kg];
}
for( int z = 0; z < LoopZ; ++z )
for( int y = 0; y < LoopY; ++y )
{
int kl = y * ThreadsPerBlockI + i2;
int jl = z * ThreadsPerBlockJ + j2;
int kg = k0 + kl;
int jg = j0 + jl;
local_b[kl][jl] = b[kg*N+jg];
}
__syncthreads();
for( int k1 = 0; k1 < SharedMemBlockK; ++k1 )
for( int i1 = 0; i1 < RegisterBlockI; ++i1 )
for( int j1 = 0; j1 < RegisterBlockJ; ++j1 )
{
int il = i1 * ThreadsPerBlockI + i2;
int jl = j1 * ThreadsPerBlockJ + j2;
sum[j1][i1] = fma( local_a[il][k1], local_b[k1][jl], 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( float* device_a, float* device_b, float* device_c ) {
dim3 grid( N / SizeJ, N / SizeI );
dim3 block( ThreadsPerBlockJ, ThreadsPerBlockI );
kernel<<< grid, block >>>( device_a, device_b, device_c );
}
これで44 TFLOPSになりました。
global memoryからshared memoryへの転送にfloat4を使う
CUDA言語にはfloatが4つ入ったfloat4という型があります（RGBAやクォータニオンを扱うためでしょうか？）。
各スレッドで行う演算命令にはSIMD命令がありませんが、メモリアクセス命令にはSIMD命令があります。
例えば、global memoryからfloat4をロードするLDG.E.128命令*2や、shared memoryにfloat4をストアするSTS.128命令などがあります。
これを使ってみます。
@@ -31,21 +31,21 @@ __global__ void kernel( float* a, float* b, float* c ) { for( int k0 = 0; k0 < N; k0 += SharedMemBlockK ) { __syncthreads(); - for( int x = 0; x < LoopX; ++x ) { - int il = ( x * ThreadsPerBlock + tid ) / SharedMemBlockK; - int kl = ( x * ThreadsPerBlock + tid ) % SharedMemBlockK; + for( int x = 0; x < LoopX/4; ++x ) { + int il = ( x * ThreadsPerBlock + tid ) / (SharedMemBlockK/4); + int kl = ( x * ThreadsPerBlock + tid ) % (SharedMemBlockK/4) * 4; int ig = i0 + il; int kg = k0 + kl; - local_a[il][kl] = a[ig*N+kg]; + *reinterpret_cast<float4*>(&local_a[il][kl]) = *reinterpret_cast<float4*>(&a[ig*N+kg]); } - for( int z = 0; z < LoopZ; ++z ) + for( int z = 0; z < LoopZ/4; ++z ) for( int y = 0; y < LoopY; ++y ) { int kl = y * ThreadsPerBlockI + i2; - int jl = z * ThreadsPerBlockJ + j2; + int jl = ( z * ThreadsPerBlockJ + j2 ) * 4; int kg = k0 + kl; int jg = j0 + jl; - local_b[kl][jl] = b[kg*N+jg]; + *reinterpret_cast<float4*>(&local_b[kl][jl]) = *reinterpret_cast<float4*>(&b[kg*N+jg]); } __syncthreads();
なんと52 TFLOPSまで速くなりました。 グローバルメモリへの要求を一度にたくさん送ることができたのが高速化の要因でしょうか。
shared memoryからレジスタへの転送にfloat4を使う
__align__(16)とすることで、コンパイラが自動的にLDS.128命令を使ってくれます。
@@ -26,8 +26,8 @@ __global__ void kernel( float* a, float* b, float* c ) { float sum[RegisterBlockJ][RegisterBlockI] = {}; - __shared__ float local_a[SizeI][SharedMemBlockK]; // 4 KiB - __shared__ float local_b[SharedMemBlockK][SizeJ]; // 16 KiB + __align__(16) __shared__ float local_a[SizeI][SharedMemBlockK]; // 4 KiB + __align__(16) __shared__ float local_b[SharedMemBlockK][SizeJ]; // 16 KiB for( int k0 = 0; k0 < N; k0 += SharedMemBlockK ) { __syncthreads(); @@ -50,22 +50,26 @@ __global__ void kernel( float* a, float* b, float* c ) { __syncthreads(); for( int k1 = 0; k1 < SharedMemBlockK; ++k1 ) - for( int i1 = 0; i1 < RegisterBlockI; ++i1 ) - for( int j1 = 0; j1 < RegisterBlockJ; ++j1 ) + for( int i1 = 0; i1 < RegisterBlockI/4; ++i1 ) + for( int j1 = 0; j1 < RegisterBlockJ/4; ++j1 ) + for( int i3 = 0; i3 < 4; ++i3 ) + for( int j3 = 0; j3 < 4; ++j3 ) { - int il = i1 * ThreadsPerBlockI + i2; - int jl = j1 * ThreadsPerBlockJ + j2; + int il = ( i1 * ThreadsPerBlockI + i2 ) * 4 + i3; + int jl = ( j1 * ThreadsPerBlockJ + j2 ) * 4 + j3; - sum[j1][i1] = fma( local_a[il][k1], local_b[k1][jl], sum[j1][i1] ); + sum[j1*4+j3][i1*4+i3] = fma( local_a[il][k1], local_b[k1][jl], sum[j1*4+j3][i1*4+i3] ); } } - for( int i1 = 0; i1 < RegisterBlockI; ++i1 ) - for( int j1 = 0; j1 < RegisterBlockJ; ++j1 ) + for( int i1 = 0; i1 < RegisterBlockI/4; ++i1 ) + for( int j1 = 0; j1 < RegisterBlockJ/4; ++j1 ) + for( int i3 = 0; i3 < 4; ++i3 ) + for( int j3 = 0; j3 < 4; ++j3 ) { - int i = i0 + i1 * ThreadsPerBlockI + i2; - int j = j0 + j1 * ThreadsPerBlockJ + j2; - c[i*N+j] += sum[j1][i1]; + int i = i0 + ( i1 * ThreadsPerBlockI + i2 ) * 4 + i3; + int j = j0 + ( j1 * ThreadsPerBlockJ + j2 ) * 4 + j3; + c[i*N+j] += sum[j1*4+j3][i1*4+i3]; } }
これにより、55 TFLOPSまで性能向上しました。
調べてみるとこれはLDS.128を使ったことによる効果だけではなく、j3の分割を導入したことも少なくない寄与があるようです。
__align__(16)を外したコードの性能は、53 TFLOPSくらいでした。
shared memoryに入れるときに転置
いかにも良さそうな改良ですが、実際には性能が47 TFLOPSに下がってしまいました。
ダブルバッファリング
GitHub - NVIDIA/cutlass: CUDA Templates for Linear Algebra Subroutinesでも採用されているようですが、実際には性能が39 TFLOPSに下がってしまいました。 性能低下の主要因は二倍アンローリングしたことのようですが（もともとレジスタブロッキングでものすごくアンローリングされているので、これ以上アンローリングすると命令キャッシュに収まらなくなる？）、二倍アンローリングにならないように気を付けて書いても性能向上することはありませんでした。
ワープ内の32スレッドが担当する領域を正方形に近くする
現状ではワープ内の32スレッドは横長な（j軸方向に長い）長方形を担当していますが、正方形に近い領域を担当するようにするという改良です。
横長の長方形を担当する場合、32スレッドがfloat4を読み出すと4回アクセスが必要なところ、担当領域の横の長さを1/4にすることで1回のアクセスで済むようにできる効果があると思われます。
しかし実際にやってみてもあまり効果がありませんでした。
レジスタブロッキングを行っているため、shared memoryからの読み出しはボトルネックになっていないことが原因と思われます。
その他
ダブルバッファリングの副産物
ダブルバッファリング自体は性能向上に寄与しなかったものの、shared memoryをたくさん取ることは性能を向上させることがわかりました。 つまり、以下のように、無意味にたくさんのshared memoryを確保することが有効です。
@@ -26,8 +26,8 @@ __global__ void kernel( float* a, float* b, float* c ) { float sum[RegisterBlockJ][RegisterBlockI] = {}; - __align__(16) __shared__ float local_a[SizeI][SharedMemBlockK]; // 4 KiB - __align__(16) __shared__ float local_b[SharedMemBlockK][SizeJ]; // 16 KiB + __align__(16) __shared__ float local_a[SizeI*2][SharedMemBlockK]; // 8 KiB + __align__(16) __shared__ float local_b[SharedMemBlockK*2][SizeJ]; // 32 KiB for( int k0 = 0; k0 < N; k0 += SharedMemBlockK ) { __syncthreads();
よくわかりませんが、性能が57.3 TFLOPSに向上しました。 パディングとかではなく本当に無意味にshared memoryを確保しただけなので、性能が上がるのは不思議です。
shared memoryはSMあたりに100 KiBしかないので、40 KiBも使ってしまうとSMあたりに3つのスレッドブロックを入れることができなくなってしまいます。
ncu*3というプロファイラを使って確認したところ、そもそも元のカーネルでもレジスタを255個も使っていてSMあたりに2つのブロックしか入れられないようです。
そんなにレジスタを使うような部分はないと思うのですが、どういうことなのでしょう。
性能が向上した原因はもしかすると、shared memoryを無駄遣いすることでコンパイラに積極的な最適化の機会を提供できた（どうせ一つのSMにたくさんのスレッドブロックが入らないのでレジスタは贅沢に使ってよい、と指示することにつながった）から、ということかもしれません。
なお、プロファイルを行った副産物として、cuBLASはvoid cutlass::Kernel2<cutlass_80_simt_sgemm_256x128_8x4_nn_align1>(T1::Params) (256, 2, 3)x(256, 1, 1)を呼び出していることがわかりました。
cuBLASの中身は手書きアセンブリとかではなくてCUTLASSなんですね。
あと、一つのSMではスレッドブロックを1つしか実行していないということもわかりました。
一つのSMで複数のスレッドブロックを実行したほうが全体が止まるようなイベントが発生しづらくてよさそうなのですが、そうでもないということなのでしょうか。
レジスタ割り付け
以前書いたコードがそうなってただけなのですが、冷静に考えるとfloat sum[RegisterBlockJ][RegisterBlockI];は転置されていて奇妙です。
しかしながら、この転置をやめると、コードによっては性能が大きく性能が低下します。
ptxを見てもレジスタ割り付けが異なる以外には同じコードになっているようで、不思議です。
Performance Upper Bound Analysis and Optimization of SGEMM on Fermi and Kepler GPUsという論文*4（https://inria.hal.science/hal-00789958/document/）によれば、レジスタもバンクコンフリクトが発生しうるようなので、それが原因でしょうか。
最近のNVIDIA社のGPUでは、レジスタのバンクコンフリクトの影響を緩和させるような機構がついていることがうかがえます。
sassを確認するとレジスタ番号に.reuseのようなアノテーション（？）がついていることがあります。
これは使い方から察するに、次の命令でも使う（ので捨てないでとっておいたほうが良い）ことを意味するものと思われます。
ハードウェア的には、ある種のレジスタキャッシュのようなものになっていると思われます。
このレジスタキャッシュは、バンクコンフリクトの頻度と読み出しポート数のトレードオフの改善に貢献しているはずです。
sm_86向けにコンパイルすると、sm_80向けにコンパイルした場合と異なり、次の命令以外で使う場合にも.reuseのアノテーションがつくことがあります。
命令列から察するに、2エントリ分のレジスタキャッシュがあるようです。
sm_86（GeForce 3000シリーズ）からSMあたりのCUDAコア（FP32演算器）数が増加しましたが、その増えた演算器へのオペランド供給能力を上げる効果もあるのでしょう。
レジスタキャッシュが2エントリあれば、ある種のダブルバッファ的に使うことができます。 そうすると、FFMA命令の一つのオペランドを常にレジスタキャッシュから読み出すような命令列にすることができて、レジスタのバンクコンフリクト頻度を下げられそうです。 しかし、CUDA言語上でそうなるような順番のループを書いても、コンパイラが命令を並べ替えてしまって、そのような命令列を出すことはできませんでした。 ptxを手書きすればよいのかもしれませんが、面倒でやっていません。
また、sm_89（GeForce 4000シリーズ）向けにコンパイルした場合、sm_80と同様、本当に次の命令で使う場合にしか.reuseのアノテーションがつきませんでした。
GeForce 4000シリーズでは、そもそもレジスタキャッシュは1エントリ分しかないのかもしれません。
i3ループの削除
i3での分割は、shared memory内に行列Aが転置された形で格納されているならば、LDS.128が使えるようになるという意味があります。 実際には転置は遅いのでやっておらず、この部分は役に立っていません。 これを削除したところ、57.5 TFLOPSとわずかながら性能向上が得られました。
cuobjdumpでsassを見てみたところ、k1軸方向にアンローリングされてLDS.128命令が使われていました。 なので転置した形でshared memoryに格納するというのはそもそも効果がない最適化だったのです。
レジスタ使用量の明示的指定
上で見たように、レジスタをたくさん使えるとコンパイラに伝えたほうが性能が高くなるような雰囲気だったので、--maxrregcount=255をつけてコンパイルしました。
これにより、58.3 TFLOPSに性能が向上しました。
なお、shared memoryを無駄に確保するのをやめると性能が下がったため、shared memoryを無駄に確保することは、レジスタをたくさん使えるとコンパイラに伝えるのとは異なる効果を生んでいるようです。 かなり意味不明な挙動ですが、何が原因なのかは不明です。
評価に使ったコード
kernel.cu↓
static constexpr int N = 4096;
static constexpr double flop_per_fma = 2.0;
static constexpr double insn_per_fma = 1.0 / 32;
static constexpr int ThreadsPerBlockJ = 32;
static constexpr int ThreadsPerBlockI = 4;
static constexpr int SharedMemBlockK = 16;
static constexpr int RegisterBlockJ = 8;
static constexpr int RegisterBlockI = 16;
static constexpr int SizeJ = RegisterBlockJ * ThreadsPerBlockJ; // 256
static constexpr int SizeI = RegisterBlockI * ThreadsPerBlockI; // 64
static constexpr int LoopX = SizeI * SharedMemBlockK / ThreadsPerBlockJ / ThreadsPerBlockI; // 8
static constexpr int LoopY = SharedMemBlockK / ThreadsPerBlockI; // 4
static constexpr int LoopZ = SizeJ / ThreadsPerBlockJ; // 8
static constexpr int ThreadsPerBlock = ThreadsPerBlockI * ThreadsPerBlockJ; // 128
__global__ void kernel( float* a, float* b, float* c ) {
int i0 = blockIdx.y * SizeI;
int j0 = blockIdx.x * SizeJ;
int i2 = threadIdx.y;
int j2 = threadIdx.x;
int tid = threadIdx.y * ThreadsPerBlockJ + threadIdx.x;
float sum[RegisterBlockJ][RegisterBlockI] = {};
__align__(16) __shared__ float local_a[SizeI*2][SharedMemBlockK]; // 8 KiB
__align__(16) __shared__ float local_b[SharedMemBlockK*2][SizeJ]; // 32 KiB
for( int k0 = 0; k0 < N; k0 += SharedMemBlockK ) {
__syncthreads();
for( int x = 0; x < LoopX/4; ++x ) {
int il = ( x * ThreadsPerBlock + tid ) / (SharedMemBlockK/4);
int kl = ( x * ThreadsPerBlock + tid ) % (SharedMemBlockK/4) * 4;
int ig = i0 + il;
int kg = k0 + kl;
*reinterpret_cast<float4*>(&local_a[il][kl]) = *reinterpret_cast<float4*>(&a[ig*N+kg]);
}
for( int z = 0; z < LoopZ/4; ++z )
for( int y = 0; y < LoopY; ++y )
{
int kl = y * ThreadsPerBlockI + i2;
int jl = ( z * ThreadsPerBlockJ + j2 ) * 4;
int kg = k0 + kl;
int jg = j0 + jl;
*reinterpret_cast<float4*>(&local_b[kl][jl]) = *reinterpret_cast<float4*>(&b[kg*N+jg]);
}
__syncthreads();
for( int k1 = 0; k1 < SharedMemBlockK; ++k1 )
for( int i1 = 0; i1 < RegisterBlockI; ++i1 )
for( int j1 = 0; j1 < RegisterBlockJ/4; ++j1 )
for( int j3 = 0; j3 < 4; ++j3 )
{
int il = i1 * ThreadsPerBlockI + i2;
int jl = ( j1 * ThreadsPerBlockJ + j2 ) * 4 + j3;
sum[j1*4+j3][i1] = fma( local_a[il][k1], local_b[k1][jl], sum[j1*4+j3][i1] );
}
}
for( int i1 = 0; i1 < RegisterBlockI; ++i1 )
for( int j1 = 0; j1 < RegisterBlockJ/4; ++j1 )
for( int j3 = 0; j3 < 4; ++j3 )
{
int i = i0 + i1 * ThreadsPerBlockI + i2;
int j = j0 + ( j1 * ThreadsPerBlockJ + j2 ) * 4 + j3;
c[i*N+j] += sum[j1*4+j3][i1];
}
}
void call_kernel( float* device_a, float* device_b, float* device_c ) {
dim3 grid( N / SizeJ, N / SizeI );
dim3 block( ThreadsPerBlockJ, ThreadsPerBlockI );
kernel<<< grid, block >>>( device_a, device_b, device_c );
}
main.cu↓
#include <iostream>
#include <chrono>
#include <random>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include "kernel.cu"
void test( float*, float*, float* );
void cublas( float* device_a, float* device_b, float* device_c ) {
cublasHandle_t handle; cublasCreate(&handle);
float alpha = 1.0, beta = 1.0;
cublasSgemm( handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N, &alpha, device_b, N, device_a, N, &beta, device_c, N );
}
int main( int argc, char* argv[] ) {
std::mt19937_64 mt;
std::uniform_real_distribution<float> dist( 0.0, 1.0 );
float* a = (float*)malloc( N*N*sizeof(float) );
float* b = (float*)malloc( N*N*sizeof(float) );
float* c = (float*)malloc( N*N*sizeof(float) );
float* d = (float*)malloc( N*N*sizeof(float) );
for( int i = 0; i < N*N; ++i ) {
a[i] = dist( mt );
b[i] = dist( mt );
c[i] = dist( mt );
d[i] = c[i];
}
std::cout << "initialized." << std::endl;
float* device_a; cudaMalloc( (void**)&device_a, N*N*sizeof(float) );
float* device_b; cudaMalloc( (void**)&device_b, N*N*sizeof(float) );
float* device_c; cudaMalloc( (void**)&device_c, N*N*sizeof(float) );
float* device_d; cudaMalloc( (void**)&device_d, N*N*sizeof(float) );
cudaMemcpy( device_a, a, N*N*sizeof(float), cudaMemcpyHostToDevice );
cudaMemcpy( device_b, b, N*N*sizeof(float), cudaMemcpyHostToDevice );
cudaMemcpy( device_c, c, N*N*sizeof(float), cudaMemcpyHostToDevice );
cudaMemcpy( device_d, d, N*N*sizeof(float), cudaMemcpyHostToDevice );
test( device_a, device_b, device_c );
cublas( device_a, device_b, device_d );
cudaMemcpy( c, device_c, N*N*sizeof(float), cudaMemcpyDeviceToHost );
cudaMemcpy( d, device_d, N*N*sizeof(float), cudaMemcpyDeviceToHost );
for( int i = 0; i < N*N; ++i ) {
float err = c[i] - d[i];
if( std::abs(err / d[i] ) > 1e-5 ) {
std::cout << c[i] << " " << d[i] << " " << err << " " << err / d[i] << " @ " << i << std::endl;
exit(1);
}
}
for( int i = 0; i < 50; ++i ) {
test( device_a, device_b, device_c );
}
}
void test( float* device_a, float* device_b, float* device_c ) {
cudaEvent_t device_start; cudaEventCreate( &device_start );
cudaEvent_t device_finish; cudaEventCreate( &device_finish );
cudaEventRecord( device_start );
call_kernel( device_a, device_b, device_c );
cudaEventRecord( device_finish );
cudaEventSynchronize( device_finish );
float ms; cudaEventElapsedTime( &ms, device_start, device_finish );
const float device_time = ms * 1e-3;
std::printf( "%9f second (device), %11.3f TFLOPS, %-6.3e FPIns/s\n", device_time, flop_per_fma*N*N*N/device_time * 1e-12, insn_per_fma*N*N*N/device_time);
}
cublas.cu↓
#include <iostream>
#include <chrono>
#include <random>
#include <cuda_runtime.h>
#include <cublas_v2.h>
static constexpr int N = 4096;
static constexpr double flop_per_fma = 2.0;
static constexpr double insn_per_fma = 1.0 / 32;
void test( float* device_a, float* device_b, float* device_c );
int main( int argc, char* argv[] ) {
std::mt19937_64 mt;
std::uniform_real_distribution<float> dist( 0.0, 1.0 );
float* a = (float*)malloc( N*N*sizeof(float) );
float* b = (float*)malloc( N*N*sizeof(float) );
float* c = (float*)malloc( N*N*sizeof(float) );
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;
float* device_a; cudaMalloc( (void**)&device_a, N*N*sizeof(float) );
float* device_b; cudaMalloc( (void**)&device_b, N*N*sizeof(float) );
float* device_c; cudaMalloc( (void**)&device_c, N*N*sizeof(float) );
cudaMemcpy( device_a, a, N*N*sizeof(float), cudaMemcpyHostToDevice );
cudaMemcpy( device_b, b, N*N*sizeof(float), cudaMemcpyHostToDevice );
cudaMemcpy( device_c, c, N*N*sizeof(float), cudaMemcpyHostToDevice );
for( int i = 0; i < 50; ++i ) {
test( device_a, device_b, device_c );
}
}
void test( float* device_a, float* device_b, float* device_c ) {
cudaEvent_t device_start; cudaEventCreate( &device_start );
cudaEvent_t device_finish; cudaEventCreate( &device_finish );
cublasHandle_t handle; cublasCreate(&handle);
float alpha = 1.0, beta = 1.0;
cudaEventRecord( device_start );
cublasSgemm( handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N, &alpha, device_b, N, device_a, N, &beta, device_c, N );
cudaEventRecord( device_finish );
cudaEventSynchronize( device_finish );
float ms; cudaEventElapsedTime( &ms, device_start, device_finish );
const float device_time = ms * 1e-3;
std::printf( "%9f second (device), %11.3f TFLOPS, %-6.3e FPIns/s\n", device_time, flop_per_fma*N*N*N/device_time * 1e-12, insn_per_fma*N*N*N/device_time);
}
実験環境
- nvcc V12.5.82 (Thu_Jun__6_02:18:23_PDT_2024, compiler.34385749_0)
nvcc main.cu -arch compute_86 -code sm_86 -lcublas --maxrregcount=255
- NVIDIA GeForce RTX 4090
- Compute Capability: 8.9
- Register/SM: 256 KiB
- Shared Memory/SM: 100 KiB
- SM: 128
注意
sm_86向けにコンパイルしています。 sm_89向けにコンパイルするとわずかに性能が下がります。 sm_80向けにコンパイルすると大きく性能が下がります。
cuBLASは一回目の実行はものすごく遅いので、それを除外して計測する必要があります。 おそらく、ライブラリをロードするのに時間がかかっているのだと思います。
まとめ
コンシューマー向けのGPUを使って、単精度行列積を高速化しました。 非常に単純なコードながら、cuBLAS（59.2 TFLOPS）に匹敵する58.3 TFLOPS（理論性能の70.6%、cuBLASの98.4%）を達成することができました。
性能向上の履歴
| 改良 | 性能 |
|---|---|
| （ベース）L1キャッシュに頼るコード | 30 TFLOPS |
| shared memoryを使う | 37 TFLOPS |
| C行列を最後に読む | 39 TFLOPS |
| SharedMemBlockKを32から16に変更 | 44 TFLOPS |
| global memoryからの転送にfloat4を使う | 52 TFLOPS |
| shared memoryからの転送にfloat4を使う | 55 TFLOPS |
| shared memoryを余計に確保する（謎） | 57.3 TFLOPS |
| i3ループの削除 | 57.5 TFLOPS |
| レジスタをたくさん使ってコンパイル | 58.3 TFLOPS |
*1:私はほとんどゲームをしないので、あまり役立っていなくてかわいそう
*2:Eはアドレスが64ビットであることを示すようです（extended address）。GPUメモリが少なかった時代はアドレスが32ビットだったのでしょうか？
*3:WSL2上で動かすためには、Windows上のドライバの設定で「GPUパフォーマンスカウンターへのアクセスを全てのユーザーに許可する」を選ぶ必要がありました。参考：https://peterchng.com/blog/2024/03/02/profiling-cuda-programs-on-wsl-2/
*4:TAGEを考案したAndré Seznec氏が最終著者で、計算機関連のことは何でもやっているんだなぁという感じがあります。第一著者のJunjie Lai氏は今はNVIDIAにいるようです。
