はじめに
PhysiKyuのAdventCalendar2024の11日目です。
昨今、生成AI技術の急拡大に伴い、GPUの重要性が社会にも認知されてきています。
この記事では、GPUのどのような特徴が膨大な数値計算に向いているのかの仕組みと、簡単なコード例を解説します。
興味を持ってもらうため、今回GPUを使うことによってどのくらいパフォーマンスの向上が得られたのかの結果を先に書きます。
今回は$2048\times 2048$の2つの行列の積をC++で計算してみました。
素朴に3重ループで計算した場合と、GPUを使って計算した場合では、実に
約75倍GPUを使うコードの方が速い
という結果になりました。
もちろん、今回は説明のため、どちらのコードも最適化をほとんど考慮していないので、あくまでも参考程度に。
数値計算の際に、GPUという選択肢があることを頭の片隅に置いておいてもらえれば幸いです。
この記事のコード例を手元のPCで実行したい場合、NVIDIA GPUを取り付けたPCでCUDAの環境設定を行う必要があります。これらの準備が面倒な場合は、Google Colabや、AWSのGPUをつかえるEC2インスタンスを使えばよいです。
GPUとは
GPUとは、Graphics Processing Unitの略で、その名の通り、ディスプレイに描画される画面の描画を行うために使われています。
現代では、FPSゲームに代表されるように、3Dの空間をカメラであっちこっち向きながら自由に探索するということが当たり前に行われているのでなかなか意識しないかもしれませんが、このような作業には、メモリ上に保存されている個別の3Dモデル(頂点データ)を、画面の座標系に変換したうえで、ピクセル単位で全体を埋めて、色の決定やテクスチャのマッピングをするといった処理を行い、かつ、ゲーム内のオブジェクトの位置などを更新するという処理を、最低でも毎秒30回程度行う必要があります。CPUだけでこのような処理をおこなうのは、さすがに無理があることは想像に難くないのではないでしょうか?
そこでGPUという、CPUよりはすこし汎用性や命令の種類を減らす代わりに、同時に大量の計算を行うことのできるチップが登場しました。当初はグラフィックスのために使われていたGPUですが、その膨大な計算を行える能力は、徐々に科学計算にも応用されていくようになり、今まさに、膨大な計算を必要とするAI分野に必須のものとなっています。
並列実行の概念
同時に大量の計算をするとは、どういうことでしょうか。たとえば、行列の計算を考えてみましょう。行列の積を計算するプログラムをC++で書けといわれたら、以下のようになるのではないでしょうか。
const int N = 2048;
// 行列を1次元で宣言. M[i * N + j]でMijを表すと約束する.
std::vector<int> A(N * N, 1); // 簡単のためすべての要素を1で初期化
std::vector<int> B(N * N, 1);
std::vector<int> C(N * N, 0);
// C = A × B
for (int i = 0; i < N; ++i) {
for (int j = 0; j < N; ++j) {
for (int k = 0; k < N; ++k) {
C[i * N + j] += A[i * N + k] * B[k * N + j];
}
}
}
ここではvectorは配列だと思ってもらって大丈夫です1。行列を素朴にネストしたvectorとして表してしまうとパフォーマンス上問題がある2ので、1次元のvectorとして表して、行列$M$の$(i, j)$成分はM[i*N+j]
として表すことにします。ただし i や j は0から始まります。
3重ループがありますから、Nが大きくなるとかなりの計算時間($\Theta(N^3)$)になることが想像できるのではないでしょうか。自分のノートパソコンでは、$N = 2048$のとき、最適化コンパイルを使わない場合150秒ほど時間がかかりました。
さて、今のコードでは、行列$C$の要素を1つずつ計算しています。$c_{12}$を計算した次のループで$c_{13}$を計算するようになっているわけです。しかし、$C$の各要素の計算は完全に独立していますから、ある要素を計算するのに、他の要素の計算が終わるのを待つ必要はありません。各要素を並列して計算したいです。
もちろん、最近のPC用CPUなら複数のコアをもつことが普通なので、並列処理は可能ですが、高々数十のコアしか持っていないので、数十個の計算を並列に行うことしかできません。
一方、GPUでは、数百から数万のコアをもつため、CPUよりも圧倒的に多くの演算を並列に行うことが可能です。GPUをつかって、行列$C$の各要素の計算を並列に行うことができれば、もっと速く計算を終えられそうです。
GPUの構造
それでは、GPUの実際の構造はどのようになっているのでしょうか。ここでは、NVIDIA GPUの用語を使って解説していきます。後でCUDAをつかったコード例を解説しますが、そのコード例を理解するのに重要な用語を重点的に解説します。
ソフトウェアの観点から
スレッド
スレッドは並列処理の基本単位になります。さきほどの行列の例だと、行列$C$の各要素を計算する部分を一つのスレッドで行うことになります。
c_{ij} = \sum_{k=0}^Na_{ik}b_{kj}
スレッドブロック
スレッドブロックとは、スレッドのまとまりです。ソフトウェアの観点からは、スレッドを3次元に並べたようなものだと思えばよいです。スレッドを点だと思うと、スレッドブロックは3次元座標系のようなもので、$(1, 2, 1)$といったインデックスによりスレッドを特定することができます。スレッドの個数には制限があり、私の NVIDIA GeForce RTX 4070 Laptop GPU だと、ブロック内での上限は1024, ブロック内の各次元での上限は順に、1024, 1024, 64でした。
グリッド
グリッドは、スレッドブロックのまとまりです。こちらも、スレッドブロックを3次元に並べたものだと思ってよいです。こちらにも次元ごとに制限があり、私の環境では順に、2147483647, 65535, 65535でした。
私の技術的な限界で、2次元でしか書けませんが、以下のようなイメージになります。
ハードウェアの観点から
CUDAコア
1つ前の節で、GPUは大量のコアを持っていると言いました。コアというのは、簡単に言うと演算器のことで、NVIDIA GPUにおいてはCUDAコアと呼ばれます。CUDAコアは大量にあり、理論的にはCUDAコアの数だけ並列にスレッドを実行できます。
ワープ
ワープとは、CUDAにおいて並列実行されるスレッドのまとまりを指し、通常は32個のスレッドで構成されます。本来であれば、すべてのスレッドが独立した命令列を同時に実行できるのが理想ですが、ハードウェアの制約により、実際には1つの命令列を複数のスレッドで共有して実行します。このスレッドのまとまりがワープです。
ただし、スレッドのまとまりをワープ単位に制限することで、次のような課題が生じます。たとえば、80個のスレッドを並列実行したい場合、32スレッド単位での処理が必要になるため、3つのワープ(96スレッド)が必要です。この場合、96スレッド中の16スレッド(96 - 80)は実際には処理を行わない「無駄なスレッド」となりますが、これは設計上のトレードオフということになります。
SM (Streaming Multiprocessor)
CUDAコアやその他のユニットは、このSMにまとまっています。各スレッドブロックは、実行時にSMに割り当てられ、SM内部でさらにワープ(32スレッド単位)に分割されて実行されます。ワープがどのように分割され、どのような順番で実行されるかについては、少し込み入った話になるため、ここでは詳細を割愛します。
ホストメモリとデバイスメモリ
GPUとCPUの特性の違いから、たいていの場合CPUとGPUではそれぞれ別のメモリを使います。このとき、CPU側のメモリをホストメモリ、GPU側のメモリをデバイスメモリといいます。この構成から、CPUがGPUに処理を依頼するとき、GPUの計算結果をCPUで取得するときに、メモリ間のコピーが必要になります。
CUDAについて
ここまでGPUの説明をしてきましたが、GPUはもともとグラフィックスのためにつくられたチップなので、科学技術計算でつかうには、あまり適したプログラミング環境がありませんでした。そこでNVIDIAによって、C言語に機能追加をするという形で提供されたプログラミング言語が、 CUDA(Couputed Unified Device Architecture) です。
CUDAでプログラミングを行う際、まずCPU側(ホストプログラム)でGPUに処理を依頼します。このとき、GPU側で実行されるコードは カーネルコード と呼ばれます。この形態のため、ホストメモリのデータをデバイスメモリ(GPU側のメモリ)にコピーし、GPUが計算を実行します。計算結果はデバイスメモリに格納され、その後、ホストメモリにコピーして使用します。
CUDAを使ってみる
それではいよいよ、GPUを使って行列の積の計算をやってみましょう。
まずは、GPU側で実行されるカーネル関数を見てみましょう。これが行列$C$の一つの要素を計算するスレッドで実行される関数になります。
__global__ void matMul(const int* A, const int* B, int* C, int N) {
int row = blockIdx.y * blockDim.y + threadIdx.y; // 行番号
int col = blockIdx.x * blockDim.x + threadIdx.x; // 列番号
if (row < N && col < N) {
int sum = 0;
for (int k = 0; k < N; ++k) {
sum += A[row * N + k] * B[k * N + col];
}
C[row * N + col] = sum;
}
}
ここに出てくる、blockIdx
やblockDim
、threadIdx
はCUDAの組み込み変数です。また、__global__
というキーワードはCUDA特有のキーワードです。これらについては後で解説します。
この関数を使うコードが以下です。
const int N = 2048;
// 行列を1次元で宣言. M[i * N + j]でMijを表すと約束する.
// これはホストメモリに置かれる
std::vector<int> A(N * N, 1);
std::vector<int> B(N * N, 1);
std::vector<int> C(N * N, 0);
// デバイスメモリに各行列用のメモリ確保
int* d_A;
int* d_B;
int* d_C;
cudaMalloc(&d_A, N * N * sizeof(int));
cudaMalloc(&d_B, N * N * sizeof(int));
cudaMalloc(&d_C, N * N * sizeof(int));
// ホストメモリからデバイスメモリに行列をコピー
cudaMemcpy(d_A, A.data(), N * N * sizeof(int), cudaMemcpyHostToDevice);
cudaMemcpy(d_B, B.data(), N * N * sizeof(int), cudaMemcpyHostToDevice);
// スレッドブロックあたりのスレッドの配置
dim3 threadsPerBlock(16, 16);
// グリッドあたりのスレッドブロックの配置
dim3 blocksPerGrid((N + threadsPerBlock.x - 1) / threadsPerBlock.x,
(N + threadsPerBlock.y - 1) / threadsPerBlock.y);
// GPUで行列の積の計算を実行
matMul<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);
// デバイスメモリから計算結果をホストメモリの行列Cにコピー
cudaMemcpy(C.data(), d_C, N * N * sizeof(int), cudaMemcpyDeviceToHost);
// デバイスメモリを解放
cudaFree(d_A);
cudaFree(d_B);
cudaFree(d_C);
このコードを実行すると、2秒ほどで計算が終了しました。最初の単純な3重ループによるコードの比べると、実に75倍の速度改善です。
今回は説明のためエラーハンドリングについては考慮していないので、このコードを試す目的以外でそのまま使うのはお勧めしません
解説
スレッドブロックの作成
以下のコードで、行列$2048\times2048$の行列$C$の計算を、それぞれが$16×16$個のスレッドを持つ、$128\times128$個のスレッドブロックで分担できるようにしています。
// スレッドブロックあたりのスレッドの配置
dim3 threadsPerBlock(16, 16);
// グリッドあたりのスレッドブロックの配置
dim3 blocksPerGrid((N + threadsPerBlock.x - 1) / threadsPerBlock.x,
(N + threadsPerBlock.y - 1) / threadsPerBlock.y);
dim3
というのはCUDA特有の型で、指定のない次元のサイズを1にしています。 (今回だと2次元しか指定しておらず、$z$は指定していません)。
CUDAの組み込み変数
CUDAには、各スレッドが、自分自身がグリッド内でどの位置にいるのかを知るための組み込み変数があります。これにより、各スレッドが、行列$C$のどの要素を計算すればいいのかを知ることができます。
blockIdx
: グリッド内でのスレッドブロックのインデックス。blockIdx.x
やblockIdx.y
で各次元のインデックスを取得できます。
blockDim
: スレッドブロックのサイズ。blockDim.x
やblockDim.y
によって各次元のサイズを取得できます。
threadIdx
: スレッドブロック内でのスレッドのインデックス。threadIdx.x
やthreadIdx.y
で各次元のインデックスを取得できます。
したがって、matMul
関数の以下のコードで、スレッドが計算すべき行列$C$の要素に行番号と列番号を計算することができます。
int row = blockIdx.y * blockDim.y + threadIdx.y; // 行番号
int col = blockIdx.x * blockDim.x + threadIdx.x; // 列番号
今回の場合、スレッドブロックのサイズは$(16, 16)$です。スレッドブロックがグリッドの中で、$x$方向のインデックスが4、$y$方向のインデックスが7、そして、スレッドがそのスレッドブロックの中で、$x$方向のインデックスが3、$y$方向のインデックスが10なら、このスレッドが計算する$C$の要素は、
$7\times16+10 = 122$行目、 $4\times16+3 = 67$行目ということになります。
そして、matMul
関数の以下のコードで、実際に$C$の要素を計算しています。なお、スレッドブロックの分割の仕方によっては、行列$C$から飛び出してしまう行と列を得ることもあるので、以下のような条件文が必要になります。
if (row < N && col < N) {
int sum = 0;
for (int k = 0; k < N; ++k) {
sum += A[row * N + k] * B[k * N + col];
}
C[row * N + col] = sum;
}
ホストメモリからデバイスメモリへのコピー
実際にGPUに計算を依頼するには、ホスト側で初期化した行列$A, B, C$を、ホストメモリからデバイスメモリにコピーしないといけません。それが以下の部分です。
// デバイスメモリに各行列用のメモリ確保
int* d_A;
int* d_B;
int* d_C;
cudaMalloc(&d_A, N * N * sizeof(int));
cudaMalloc(&d_B, N * N * sizeof(int));
cudaMalloc(&d_C, N * N * sizeof(int));
// ホストメモリからデバイスメモリに行列をコピー
cudaMemcpy(d_A, A.data(), N * N * sizeof(int), cudaMemcpyHostToDevice);
cudaMemcpy(d_B, B.data(), N * N * sizeof(int), cudaMemcpyHostToDevice);
malloc
とは、メモリ割り当て(memory allocation)の略で、CやC++でメモリの確保に使われますが、cudaMalloc
はそれのCUDA版で、デバイスメモリにメモリを確保します。デバイスメモリに確保されるメモリの先頭アドレスを保持する変数(d_A
やd_B
など)のポインタをcudaMemcpy
に渡すことで、ホスト側が、デバイスメモリのどこにデータが入っているかを知ることができます。N * N * sizeof(int)
がちょうど、各要素が整数の$N\times N$の行列のメモリ容量に等しいです。
デバイスメモリからホストメモリへのコピー
以下のコードで、GPUでの計算結果をホストメモリにコピーしています。
// デバイスメモリから計算結果をホストメモリの行列Cにコピー
cudaMemcpy(C.data(), d_C, N * N * sizeof(int), cudaMemcpyDeviceToHost);
// デバイスメモリを解放
cudaFree(d_A);
cudaFree(d_B);
cudaFree(d_C);
cudaFree
はcudaMalloc
に対応する形で、デバイスメモリに確保したメモリ領域を解放します。
カーネルコード
最後に、もっとも奇妙な構文たちを解説します。
matMul
関数にくっついている、__global__
というキーワードは、CUDA特有のキーワードで、このmatMul
関数がCPUから呼び出されて、GPUで実行される関数であることを意味します。他にも__host__
や__device__
といったキーワードがあります。
__global__
がついた関数は、。<<<グリッドサイズ, スレッドブロックサイズ>>>
という風に、スレッドが自分の位置を知るための情報を渡したうえで、ホスト側で以下のように呼び出されます
// GPUで行列の積の計算を実行
matMul<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);
今回は、コード例を理解するための知識だけを解説しましたが、CUDAには他にもたくさん知っておいてほしい概念がいくつかありますので、興味を持った方はぜひ勉強してみてください。
終わりに
GPUをつかったプログラミングに触れてもらうため、GPUそのものと、CUDAを使ったコードを解説してみました。GPUを使った計算の威力を感じていただけたのではないでしょうか。
今回は最適化については何も考慮せずにコードを書きましたが、GPUの構造を意識した最適化を施せばさらに速くすることができます。(最初の3重ループも然りですが。)
また、今回は理解に重点を置いたため、一からコードを書いてみましたが、行列やベクトルの為のライブラリcuBLASやディープラーニング用のライブラリcuDNNといった、すでに最適化されたライブラリが提供されているので、実際にはこちらを使っていくことになるのかなと思います。