行列積を並列化によって高速化する

導入

行列積は愚直に行うと$\mathcal{O}(N^3)$の計算量がかかります。疎行列の場合はデータ構造を工夫することで$\mathcal{O}(N)$くらいまで落とせると思いますが、今回は疎行列を仮定せずに密行列積を高速に求める方法について説明したいと思います。また今回は正方行列に限定して議論します。

実行環境

東大のOakforest-PACSというスパコンを使ってC言語で密行列積のプログラムを実行しました。並列計算ではOpenMPやMPIを用いました。

OpenMPによる簡単な並列化

愚直にCで行列積$C=AB$を実行するとこうなります。

for(i=0; i<n; ++i){
  for(j=0; j<n; ++j){
    C[i][j]=0.0;
    for(k=0; k<n; ++k){
      C[i][j]+=A[i][k]*B[k][j];
    }
  }
}

並列化をせずに高速化しようと思うと、ループのjとkの順番を入れ替えるとメモリに連続アクセスするようになるので良かったりします。ただコンパイル時に最適化オプションをつけるとこれくらいはやってくれるのであまり意味はありません。
ノートパソコンでも簡単にできる並列化の方法としてOpenMPがあります。これは

#pragma omp parallel for private(j,k)
for(i=0; i<n; ++i){
  for(j=0; j<n; ++j){
    C[i][j]=0.0;
    for(k=0; k<n; ++k){
      C[i][j]+=A[i][k]*B[k][j];
    }
  }
}

と一行おまじないを書くだけでスレッドを並列化してくれます。mpiiccだとコンパイラオプションで-qopenmpと書くだけで動いたのですが、#include<omp.h>と書くこともあるらしく良く分かりません。
これだけでも結構速くなりますがせっかくなのでMPIを使ってさらに定数倍の改善を目指してみましょう。

MPIによる並列化

MPIはInitしたり、Comm_rankしたり、Comm_sizeしたり結構ちゃんと書かないといけなくて、OpenMPよりは大変ですが、異なるプロセッサーで互いに作業を分担できるように通信するための便利なAPIです。詳しくはMPIでググってみたらいいと思います。
行列のサイズが$n$の状況でプロセッサーが$k$個使える状況を考えます。そうすると各プロセッサーでCの$\frac{n}{k}$行分を計算するようにすれば理論上は計算量は$\mathcal{O}(n^3/k)$になります。ただしAとBは全プロセッサーが独立にメモリを確保しています。コードの一部を載せます。

int ib = n/numprocs; //numprocsはprocessorの数
#pragma omp parallel for private(j,k)
for(i=0; i<ib; ++i){
  for(j=0; j<n; ++j){
    C[i][j]=0.0;
    for(k=0; k<n; ++k){
      C[i][j]+=A[i][k]*B[k][j];
    }
  }
}

MPIとOpenMPをハイブリッドに使っていますが、プロセッサーとスレッドの数の組み合わせを変えると結構計算時間に影響が出るので、しっかりパラメーター探索をする必要があります。

SUMMAによる行列積アルゴリズム

最近SUMMAというアルゴリズムを知って実装してみたらめちゃくちゃ速かったので紹介したいと思います。SUUMOみたいな名前してる癖にやるやん…
元論文はこのpdfに載っています。
基本的なアイデアは巨大な行列はキャッシュに乗り切らないのでキャッシュに乗るくらいの小行列に分割して、その小行列積を足し合わせて答えを求めようというものです。
ブロック行列の各要素を$A_{ij}$のように書くとすると、$C_{ij}$を求めるのに必要な小行列積は$A_{ik}B_{kj} (k=1,2,3\cdots)$です。逆に$A_{ij},B_{ij}$を必要とする$C$の要素はそれぞれ$C_{ik},C_{kj} (k=1,2,3\cdots)$です。
各プロセッサーは$A_{ij},B_{ij},C_{ij}$を持っていますが、$A_{ij}B_{ij}$だけでは$C_{ij}$は求まらないので、他の小行列は、よそのプロセッサーからもらってくる必要があります。同様に$A_{ij},B_{ij}$を必要とするプロセッサーにはそれを渡してあげなければいけません。先程の考察から$A_{ij}$は左右(同じ行)に、$B_{ij}$は上下(同じ列)に渡さなければならないのでそれをMPI_Bcastによって実現します。MPI_BcastとはMPIの通信関数の一つで、指定した値を指定した複数のプロセッサー全てに送信する機能を持ちます。
こうすることで小行列積がキャッシュアクセスのみで完結することが期待でき、より高速になります。計算量を解析してみましょう。行列のサイズを$n$、プロセッサーの数を$k^2$とおきます。プロセッサーの数は平方数である方が都合が良いです。また$k$は$n$を割る数でないと割り振る小行列のサイズが均等にならないのでプロセッサーの数を調整したり、行列のサイズを余分に大きくすると良いと思います。
そうすると、小行列のサイズが$\frac{n}{k}$となるので、各小行列の計算量は$\mathcal{O}((\frac{n}{k})^3)$です。さらにこの計算を各プロセッサーで$k$回行なってその足し算によって$C_{ij}$を求めるので最終的なオーダーは$\mathcal{O}(n^3/k^2)$になります。前のセクションではプロセッサーの数を$k$と置いていたので理論的な計算量は変わっていませんが実際には速くなります。それは小行列積がメモリアクセスの回数を大幅に削減できることに起因しています。実際には通信のオーバーヘッドで遅くなったりしますが、良い塩梅のところを見つけるとうまい具合にバランスがとれます。

結果

以下の表に実行時間の結果を示します。コンパイルオプションは全て

mpiicc -O3 -xMIC-AVX512 -qopt-report -qopenmp

を用いています。実行オプションには3種類くらいありましたが、詳細は省略してその中から最速のものだけを最終結果として表示します。1分以内に実行が終わらなかったものはTLEと表示しています。

	並列化なし	OpenMPのみ	MPIによる行ブロッキング	SUMMA
$N=1088$	0.28(s)	1.00(s)	0.20(s)	0.02(s)
$N=5440$	TLE	TLE	0.53(s)	0.188(s)
FLOPS	9G	2.5G	599G	1.7T

最後のSUMMAでは1024コアを使っているので、理論性能は約45TFLOPSになります。よって理論性能の3.7%くらいしか出ていないことがわかりますね。まだまだ改善の余地はありそうです。

実装

具体的な実装はGitHubに載せておきます。(もしかしたら間違っているかもしれないのでご指摘お願いします。)実装にあたってこのコードをめちゃくちゃ参考にしました。