概要
SGEMM(Single-precision General Matrix-matrix Multiply)は、BLAS(Basic Linear Algebra Subprograms)ライブラリの中核をなす関数であり、$ C=α⋅op(A)⋅op(B)+β⋅C$ という演算を実行します。この関数は、高性能コンピューティング(HPC)分野において不可欠な存在であり、科学計算、工学シミュレーション、さらにはディープラーニングといった最先端の応用に対する計算能力の基盤を提供します。特にディープラーニングでは、ニューラルネットワークの訓練および推論プロセスにおいて、大量の畳み込み演算や全結合層演算が発生します。これらの演算の核心部分は、効率的にSGEMM問題へと変換することが可能であり、高度に最適化された関数ライブラリを利用することで、ハードウェア上での高速化が実現されます。SGEMMは典型的な計算集約型(compute-intensive)タスクであるため、そのGPU上での最適化効率は、上位アプリケーションの性能上限を直接決定づける要因となり、計算パイプライン全体の最適化における重要ポイントとなります。本稿では、PTX4090を対象に、グローバルメモリのアクセス最適化、共有メモリのブロッキング、ベクトル化、レジスタ最適化、ダブルバッファリングといった技術を駆使して単精度演算子の最適化を試みました。その結果、最適化された演算子はCUBLASの性能の約80%に達し、最高では97.5%を記録しました。
1. 基本的なベンチマーク
まず、Polybenchを使用してネイティブカーネルとCUBLASカーネルを実装し、異なるデータセットでバッチ処理を実行して性能曲線を取得しました。
図1.1に示すように、CUBLASのGFLOPSはベースラインの約10倍であり、最適化の余地が大きいことが直感的に理解できます。
以下では、ネイティブカーネルの性能が低い原因について分析します。
2. 理論的分析
2.1. ボウンドの特定
最適化の基準となる行列サイズを 1024x1024x1024(M x N x K)と設定し、ハードウェア指標および計算に基づいて以下のデータを算出しました。
- Theoretical FLOPs = 2 x M x N x K = 2 x 1024 x 1024 x 1024 = 2.147 GFLOPs
- Theoretical FLOPs Peak = 82.58 TFLOPS
- Theoretical Compute Time = Theoretical FLOPs / Theoretical FLOPs Peak = 2.147 / 82,580 = 25.99 µs
- Theoretical Bytes (min) = ((M x K) + (K x N) + 2 x (M x N)) x sizeof(DATA_TYPE) = 3 x 1024 x 1024 x 4 = 12,582,912 bytes = 12.58 MB
- Theoretical Bandwidth = 1.01 TB/s
- Theoretical Mem Time = Theoretical Bytes / Theoretical Bandwidth = 11.9 µs
上記の結果から、理論的にはこのカーネルは計算バウンド(compute bound)であることがわかります。
2.2. ネイティブカーネルの分析
NCU(NVIDIA Nsight Compute)のサンプリングレポートから、以下の指標が得られました。
- Real Time = 484.67 µs
- DRAM Bytes Read Sum = 12,585,728 bytes
- Real FLOPs = 4.437 TFLOPS
- Real Bandwidth = 25.97 GB/s (対応する計算ロード時間 = 470.5 µs)
実行時間の大部分をメモリキャッシュが占めていることがわかります。したがって、ネイティブカーネルはメモリバウンド(memory bound)です。
2.3. CUBLASカーネルの分析
NCUのサンプリングレポートから、以下の指標が得られました。
- Real Time = 54.98 µs
- DRAM Bytes Read Sum = 12,585,728 bytes
- Real FLOPs = 39.11 TFLOPS
- Real Bandwidth = 228.88 GB/s (対応する計算ロード時間 = 53.7 µs)
ネイティブカーネルとCUBLASカーネルを比較すると、帯域幅の利用効率において大きな差があることがわかります。以降の最適化はこの点に注力します。
3. 段階的な最適化戦略と実験結果
3.1. Coalesced Readによるグローバルメモリアクセス最適化
同じワープ内のスレッドによる連続したメモリアクセスは、グループ化されて単一の操作として実行されることがあり、このアクセス方法をグローバルメモリのCoalesced Accessと呼びます。
ネイティブカーネルでは、以下のインデックス計算方式が使用されており、ワープ内スレッドのデータアクセスがCoalesced Accessになる保証がありませんでした。
const uint x = blockIdx.x * blockDim.x + threadIdx.x;
const uint y = blockIdx.y * blockDim.y + threadIdx.y;
グローバルメモリのCoalesced Accessを実現するため、インデックス計算方式を変更しました。
const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);
const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);
注意: このインデックス分割方式を採用するには、以下の条件を満たす必要があります。
- スレッドブロックは1次元で起動すること。
- スレッドブロックのサイズは32x32であること。
性能結果
ネイティブカーネルのスレッドブロック起動が32x32であり、基準データセットサイズが1024x1024であったため、すでにグローバルメモリのCoalesced Accessが実現されていた可能性があります。そのため、現時点では顕著な最適化結果は得られていません。
関連性能指標
ワープステータスにおいて、'Stall LG Throttle' の割合が最大(20.12サイクル)でした。これは、L1命令キューがグローバルメモリ操作で満杯になるのを待機している状態を示唆しています。
3.2. SMEM(Shared Memory)の活用
GPUのメモリ階層には、グローバルメモリ(GMEM)の他に、キャッシュ(L1/L2)および共有メモリ(SMEM)、レジスタが存在します。階層が下がるほど計算ユニットに近くなり、ストレージ容量は小さくなりますが、読み取り速度は向上します。
これまではGMEMのみを使用していましたが、SMEMの活用を試みます。
まず、GMEMからそれぞれ BLOCKSIZE x BLOCKSIZE サイズの A と B のブロックをロードします。各スレッドは引き続き C の要素1つを計算する役割を担います。これらのデータブロックを A の列方向および B の行方向に移動させながら、C の部分和を計算し、最終的な計算結果を得ます。
概念図を以下に示します。
性能結果
関連性能指標:
- Real Time = 406.05 µs
- DRAM Bytes Read Sum = 12,583,680 bytes
- Real FLOPs = 5.296 TFLOPS
- Real Bandwidth = 30.99 GB/s (対応する計算ロード時間 = 396.4 µs)
依然としてメモリバウンドです。
ワープステータスにおいて、'Stall LG Throttle' の割合は20.12サイクルから0.39サイクルに減少しましたが、'Stall MIO Throttle' の割合が最大(23.95サイクル)となりました。これは、共有メモリへのアクセスに起因するパイプラインの混雑を示唆しています。
GMEMからSMEMへのデータロード段階がボトルネックとなっています。この段階の改善策として、ダブルバッファリングや、ロード後の計算部分の増加によるレイテンシ隠蔽、またはベクトル読み出しによる帯域幅利用率の向上が考えられます。ここでは、計算量の増加による最適化を続けます。
3.3. Block Tilingによる最適化
RTX 4090のSMEMサイズは128KB(2^17バイト)です。
64x64のfloat行列は32KB(2^14バイト)を占めるため、1つのSMEMには2つの128x128行列を格納できます。前回の最適化では2x32x32のSMEMしか使用していなかったため、SMEMにより多くのデータを格納・処理させることが考えられます。
ただし、スレッドブロックの次元が通常32x32であるため、64x64や128x128のサイズではSMの利用率が低下する可能性があります。
64x64や128x128のデータをロードするということは、1スレッドあたり複数の(この場合は2x2や4x4)データを処理する必要があることを意味します。これをスレッドブロックのタイリング手法に移行します。
プロセス概念図:
性能結果
スレッドメモリアクセス回数計算
- SMEM: 各スレッドが1つの結果を計算。GMEM: K/32回の外周ループ * 2回のロード。SMEM: K/32回の外周ループ * BLOCKSIZE(=32) * 2回のロード。結果あたりのメモリアクセス: 16分の1 K回のGMEM、2K回のSMEM。
- Block Tiling (1D): 各スレッドが8つの結果を計算。GMEM: K/8回の外周ループ * 2回のロード。SMEM: K/8回の外周ループ * BK(=8) * (1 + TM(=8))。結果あたりのメモリアクセス: 32分の1 K回のGMEM、8分の9 K回のSMEM。
関連性能指標:
- Real Time = 149.28 µs
- DRAM Bytes Read Sum = 12,583,552 bytes
- Real FLOPs = 14.406 TFLOPS
- Real Bandwidth = 84.29 GB/s (対応する計算ロード時間 = 145.7 µs)
依然としてメモリバウンドですが、命令あたりのメモリ待機サイクル数が大幅に減少しました。
ワープのStall指標が大幅に減少しました(上位3つのStall指標を以下に示します)。
- Stall MIO Throttle: 6.31
- Stall Long Scoreboard: 5.52
- Stall Barrier: 2.74
3.4. Block Tiling (2D)による最適化
この手法はBlock Tiling (1D) と同様の考え方ですが、スレッドが処理するデータをさらに拡大するため、スレッドブロックを2次元で処理します。これにより、1スレッドあたり 8x8 の行列要素を計算します。
概念図:
コードの主要部分:
// 結果レジスタを定義
float threadResults[TM * TN] = {0.0};
// A、B用のレジスタ
float regM[TM] = {0.0};
float regN[TN] = {0.0};
for(int BlockIdx = 0; BlockIdx < K ; BlockIdx += BK) {
// A[innerRowA * K + innerColA] = A[innerRowA * K + innerColA];
// B[innerRowB * N + innerColB] = B[innerRowB * N + innerColB];
// データロード
for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) {
As[(innerRowA + loadOffset) * BK + innerColA] =
A[(innerRowA + loadOffset) * K + innerColA];
}
for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) {
Bs[(innerRowB + loadOffset) * BN + innerColB] =
B[(innerRowB + loadOffset) * N + innerColB];
}
// キャッシュが完全にロードされるまでブロック内のスレッドを同期
__syncthreads();
A += BK;
B += BK * N;
// 内積計算
for (int dotIdx = 0; dotIdx < BK; ++dotIdx) {
for (uint i = 0; i < TM; ++i) {
regM[i] = As[(threadRow * TM + i) * BK + dotIdx];
}
for (uint i = 0; i < TN; ++i) {
regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
}
for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
threadResults[resIdxM * TN + resIdxN] +=
regM[resIdxM] * regN[resIdxN];
}
}
}
__syncthreads();
}
for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN] =
alpha * threadResults[resIdxM * TN + resIdxN] +
beta * C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN];
}
}
性能結果
スレッドメモリアクセス回数計算
- Block Tiling (1D): 各スレッドが8つの結果を計算。GMEM: K/8回の外周ループ * 2回のロード。SMEM: K/8回の外周ループ * BK(=8) * (1 + TM(=8))。結果あたりのメモリアクセス: 32分の1 K回のGMEM、8分の9 K回のSMEM。
- Block Tiling (2D): 各スレッドが8x8個の結果を計算。GMEM: K/8回の外周ループ * 2回のロード * 1024/256。SMEM: K/8回の外周ループ * 8回の内積ループ * 2回のロード * 8回のロード。結果あたりのメモリアクセス: 64分の1 K回のGMEM、4分の1 K回のSMEM。
関連性能指標:
- Real Time = 175.23 µs
- DRAM Bytes Read Sum = 12,583,808 bytes
- Real FLOPs = 12.273 TFLOPS
- Real Bandwidth = 71.81 GB/s (対応する計算ロード時間 = 145.7 µs)
データセットが小さい場合、レイテンシ隠蔽効果が低いため、Block Tiling (1D)よりも実行時間が長くなり、FLOPSも低くなりました。
1024規模以上のデータセットテストでは、2DタイリングのFLOPSが1Dタイリングを大幅に上回りました。
ワープのStall指標も大幅に減少しました。
3.5. Vectorization(ベクトル化)、Transpose Write(転置書き込み)、Register(レジスタ)最適化
ベクトル化操作によりメモリアクセスを最適化し、元の4回のfloat型ロード命令を1回のfloat4型(128ビット)ロード操作に統合します。データ総量を維持したまま、メモリロード命令の呼び出し回数を大幅に削減し、メモリ帯域幅の利用効率を向上させます。
メモリの読み書きを伴うモジュールすべてでfloat4型をベクトル化処理し、単一の128ビットデータ転送で複数回の32ビットデータ転送を代替することで、メモリアクセスのスループットを効果的に向上させ、ハードウェアのベクトル演算能力を最大限に引き出します。
ベクトル読み出し時に行列Aの転置を行い、後続の内積計算を容易にします。
内積ループではレジスタを使用し、データ読み出しと再利用の効率を高めます。
コードの主要部分:
// メイン計算部分
for(int BlockIdx = 0; BlockIdx < K ; BlockIdx += BK) {
float4 tmp = reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
// 内積計算のための転置書き込み。
As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;
reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
__syncthreads();
A += BK;
B += BK * N;
for (int dotIdx = 0; dotIdx < BK; ++dotIdx) {
for (uint i = 0; i < TM; ++i) {
// regM[i] = As[(threadRow * TM + i) * BK + dotIdx];
regM[i] = As[dotIdx * BM + threadRow * TM + i];
}
for (uint i = 0; i < TN; ++i) {
regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
}
for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
threadResults[resIdxM * TN + resIdxN] +=
regM[resIdxM] * regN[resIdxN];
}
}
}
__syncthreads();
}
for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
// Cベクトルをレジスタにロード
float4 tmp = reinterpret_cast<float4 *>(
&C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0];
// レジスタ内でGEMM更新を実行
tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x;
tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y;
tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z;
tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w;
// 書き戻し
reinterpret_cast<float4 *>(
&C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
tmp;
}
}
性能結果
関連性能指標:
- Real Time = 156.03 µs
- DRAM Bytes Read Sum = 12,583,552 bytes
- Real FLOPs = 13.783 TFLOPS
- Real Bandwidth = 80.65 GB/s (対応する計算ロード時間 = 152.3 µs)
3.6. Buffering(ダブルバッファリング) + WarpTile
これまでの分析から、データアクセスコストがカーネル処理時間の大部分を占めていることがわかります。
前バージョンのコードでは、データの一貫性を保つために2回の`__syncthreads()`を使用していました。最初の`__syncthreads()`はRead-After-Writeの順序性を保証するために不可欠です。
しかし、2回目の同期は、データが処理される前に他のスレッドに読み取られないように、Write-After-Readの順序性を保証することを目的としていました。
これらの同期が発生する根本原因は、異なるイテレーションで同じメモリ領域にデータを保存していることにあります。実際には、これらのイテレーション間のデータに真の依存関係はありません。もし別の場所に書き込めば、同期は不要になります。
この手法をデータプリフェッチ、またはダブルバッファリングと呼びます。2倍のストレージ領域を確保し、イテレーションごとに交互に使用することで、最後の同期操作を省略し、レイテンシを隠蔽します。
概念図:
コードの主要部分:
// 2倍の領域を確保
__shared__ float smem_a[2][BM * BK];
__shared__ float smem_b[2][BK * BN];
int write_index = 0;
// 外周ループで行列ブロックを走査
for (uint bk_idx = 0; bk_idx < K; bk_idx += BK)
{
// データを共有メモリ smem_a[write_index] および smem_b[write_index] にロード
...
__syncthreads();
// 計算
// データをレジスタ reg_a[write_index] および reg_b[write_index] にロード
...
A += BK;
B += BK * N;
// 元の同期を削除
...
write_index = 1 - write_index; // 読み書きポインタを切り替え
}
性能結果
ダブルバッファリングを適用した後、FLOPSが低下するという奇妙な結果になりました。
これは、Warpタイリングを使用したことに関連している可能性があります。ロードするデータが多すぎるとSMの利用率が低下し、レイテンシ隠蔽効果が減少したと考えられます。
後ほど、`native-double-buffering` のバージョンを別途テストして、この仮説を検証する予定です。
3.7. Buffering_no_WarpTile(ダブルバッファリングのみ)
Warpタイリング部分を削除し、余分なレジスタスペースを確保してSMの占有率を向上させ、性能テストを行ったところ、一定の性能向上が見られました。
コードの主要部分:
__shared__ float As[2][BM * BK];
__shared__ float Bs[2][BK * BN];
// インデックス計算部分
...
float regM[2][TM] = {0.0};
float regN[2][TN] = {0.0};
int write_idx = 0;
// まずデータを As[0][] に転送
float4 tmp = reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
...
// 対応するオフセット
A += BK;
B += BK * N;
for(int BlockIdx = BK; BlockIdx < K ; BlockIdx += BK) {
// 計算部分
...
// データ区間を切り替え
write_idx ^= 1;
// データ読み出し
float4 tmp = reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
...
A += BK;
B += BK * N;
}
// 最後のデータを処理
for (int dotIdx = 0; dotIdx < BK; ++dotIdx) {
...
}
// データ書き込み
for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
...
}
}
3.8. Vector Register Read(ベクトルレジスタ読み出し) + Bank Conflict解決
ベクトルレジスタを宣言すると、システムは自動的に隣接するレジスタ空間をベクトルレジスタ全体に割り当てます。そのため、ベクトルレジスタを使用して最終的な内積計算部分のデータ読み出しを行い、帯域幅利用率を向上させることができます。
共有メモリは、高並列アクセスを実現するために、物理的に32個の独立したアクセス可能なメモリバンク(bank)で設計されています。複数のスレッドが同時に同じバンクにアクセスすると、バンクコンフリクトが発生し、並列データアクセスが直列アクセスに低下して帯域幅利用率が大幅に低下します。この最適化では、パディング技術を採用し、4バイトの追加スペースを埋め込むことで、ベクトル化アクセスのアライメント要件を満たしつつ、有効データの格納方法を分散させ、バンクコンフリクトを削減します。
4. 最適化のまとめ
今回の最適化実験では、共有メモリ、Block Tiling、ベクトル化操作、ダブルバッファリング、ベクトルレジスタ、およびバンクコンフリクト解決技術を適用してSGEMM演算子を最適化しました。最適化後の演算子は、CUBLASの約80%の全体性能を達成し、最高で97.5%に達しました。具体的な加速率は以下の通りです。
本実験を通じて、NVIDIA GPUアーキテクチャのコア原理と最適化メカニズムについて深く理解することができました。実践においては、共有メモリのブロッキング、ベクトル化読み出しといった主要技術を体系的に習得し、バンクコンフリクト問題も効果的に解決しました。さらに、CUDA演算子最適化の標準化されたプロセス、すなわちパフォーマンス分析、ボトルネック特定、最適化戦略の実装といった重要なステップに関する包括的な認識フレームワークを構築できました。これは、今後のHPC開発における強固な理論的および実践的基盤を築くものです。
実験結果は、一部のテストケースでは実装した演算子の性能がCUBLASライブラリのレベルに近づいているものの、全体的な計算効率には依然として顕著な差があることを示しています。簡単な分析によると、この性能差は主に以下の3つの重要な要因に起因しており、これらは今後の実験で重点的に改善すべき方向性です。
- ブロッキング戦略の最適化不足: 現在の実装では固定のブロックサイズを採用しており、入力データの規模に応じて動的に調整できていません。これにより、メモリアクセスパターンが最適でなくなり、GPUのメモリ階層構造を十分に活用できていない可能性があります。
- Tensor Coreハードウェアの未活用: 計算プロセスにおいて、NVIDIA GPUのTensor Core(Matrix Multiply-Accumulate, MMA)ユニットを効果的に利用できていません。これにより、計算集約型操作の性能がハードウェアの理論上の最高値に達していません。
- 非同期メモリコピーハードウェアリソースの利用率不足: グローバルメモリ(GMEM)から共有メモリ(SMEM)へのデータ転送中に、TMA(Tensor Memory Accelerator)ハードウェアユニットを効果的に利用できていません。この非効率的なデータ転送モードは、レジスタリソースの消費を増加させるだけでなく、CUDAコアに余分な負荷をかけ、全体的な計算効率を低下させます。
今後の実験では、これらの問題に焦点を当て、適応型ブロッキング戦略とTensor Core最適化技術を導入することで、演算子の全体性能をさらに向上させることを目指します。
付録
RTX 4090 関連ハードウェア指標
- L1 Cache: 128 KB (per SM) - 2^15 float
- L2 Cache: 72 MB
| メモリタイプ | 容量 | 帯域幅 | レイテンシ | アクセス粒度 |
|---|---|---|---|---|
| Register | 256KB/SM | 80TB/s | 1サイクル | 32ビット/スレッド |
| Shared Memory | 128KB/SM | 5.3TB/s | 20サイクル | 32バイト/ブロック |
| Global Memory | 24GB(デバイス全体) | 1TB/s | 300サイクル | 128バイト/トランザクション |
関連用語
- FLOPs: FLOPの複数形
- FLOPS: 1秒あたりの浮動小数点演算回数 (FLOPs/s)
非同期コピー
「他の場所でのレジスタ使用量を削減する方法」を調べていた際に、`s_xxx[idx] = d_xxx[idx]` の形式で、グローバルメモリから共有メモリへ「ワンステップ」で書き込むように見える手法について言及している記事を見つけました。実際には、これは中間的なレジスタ(tmp)を経由する `tmp = d_xxx[idx]; s_xxx[idx] = tmp;` という段階的な分解プロセスにコンパイルされ、2回目の書き込み時にレジスタへの依存が生じます。計算能力8.6および8.7のデバイスでは、新しい `cuda::memcpy_async` ロード方法を検討することが推奨されます。
NCUサンプリングでは顕著な違いは見られませんでしたが、GPUデータコピーモードについて調査したところ、現代のプロセッサアーキテクチャは主にロード/ストアアーキテクチャであることがわかりました。このアーキテクチャの基本原則は、計算はレジスタ上でのみ行われ、メモリアクセスは専用命令によって完了するというものです。
GMEMからSMEMへのデータ転送プロセスは、CUDAコアによるLDG(グローバルメモリロード)およびSTS(共有メモリストア)命令の実行を占有します。
`cuda::memcpy_async` を使用すると、データを非同期にコピーできるだけでなく、CUDAコアを解放し、TMAを利用できるため、理論的にはある程度の最適化が期待できます。