CUDAで複数のGPUを使うプログラムを書く
CPUで動くプログラムはそれぞれのGPU(デバイス)を独立に制御する。CUDAのStreamはデバイスごとの仕事のキューになっている。プログラムの実行はStreamの実行開始と同期で行われている。
そして、Stream間の同期を行うにはeventを使う。eventはカーネルの完了を示すことができる。また、1つのデバイス上の1つのStreamにつながれているカーネルは、他のeventを待ち合わせることができる。
複数のGPUで処理を起動する一番基本的な方法は、次のようなものである。まず、cudaSetDevice()でGPUに準備を行わせ、CudaStreamCreate()でStreamに中身を入れる。これを最初のループで全部のGPUについて行う。2番目のfor文のループではfirstKernel<<< >>>()で、全GPUでfirstKernelが起動される。
3番目のfor文ではcudaStreamSynchronize()で各デバイスでのカーネルの実行終了を待ち合わせる。そして、全GPUでカーネルの実行が終わると次に進み、4番目のfor文のループに入り、secondKernel<<<>>>()で2番目のカーネルの実行を開始させる。
これでも良いのであるが、最初のカーネルの実行が16GPUすべてで終わらないと2番目のカーネルの実行が開始されないというのがムダである。次の例では、firstKernelの実行を開始したのちに、cudaEventRecordでこのカーネルの完了を示すEventの受け取りを指示する。そして、2番目のfor文のループのcudaStreamWaitEvent()で他のGPUからのevent通知を受け取り他のGPUでの実行の終了を待ち合わせる。
次のループではそれぞれのGPUがGPU0のカーネル実行の完了を待ち合わせる。すべてのGPUがGPU0の実行終了を受け取れば、最初のステップは終わりである。
このようにeventを使ってカーネル実行の終了を通知すれば時間のかかる同期(cudaStreamSynchronize()など)は全体の終わり以外には使う必要がなくなり、プログラムの実行時間を短縮することができる。
CUDAのCooperativeKernel機能を使えば、もっと簡単に記述できる。どのStreamのどのカーネルが協調して処理を行っているかをlaunchParamsに記述してcudaLaunchCooperativeKernelMultiDevice()を呼び出し、カーネルの中でthis_multi_grid().syncを呼び出すとGPUコードの中で必要な同期が行なわれる。後は、最後にcudaStreamSynchronize()呼び出すだけで良い。
複数GPUでのメモリ管理
一般に複数のGPUがそれぞれメモリを持っている場合のメモリの管理は面倒であるが、DGX-2ではユニファイドメモリですべてのGPUとCPUのメモリがCPUからもGPUからもアクセスできるので、メモリの管理は容易である。
DGX-2では16個のGPUにそれぞれ32GBのHBM2メモリが付いている。しかし、NVSwitchが全対全の高バンド幅のメモリマッピングを提供しており、この接続は通常のアクセスだけでなくAtomicアクセスなどもサポートしている。
これとユニファイドメモリを組み合わせると全GPUで共用される512GBのメモリが存在するように見える。
本当は、全部のメモリがリニアに連続した領域になって欲しいのであるが、cudaMallocでは区分したグローバルなアドレス空間ができてしまい、データとGPUの関連が見えてしまう。
しかし、これもユニファイドメモリを使うと、(この例では)4つのGPUのメモリが連続になり、ポインタ演算が使えるようになる。
次の例は、N*N*Nの3次元のアレイを作る場合で、まず、cudaMallocManaged()でN*N*Nのメモリを確保する。次にcudaMemAdvise()で全GPUからこのメモリが見えるようにする。そして、最後にGPUごとに存在するメモリが連続するメモリアドレスになるように配置する。
(次回は4月19日に掲載します)