MPI_Comm_splitとScaLAPACKの併用

概要

MPIで並列化されたコードではLAPACKの代わりにScaLAPACKを使うことが候補となります。ここで、ScaLAPACKについてWeb上での資料は少なく、また、よくある使い方の解説では、全プロセスを使う例が殆どです(MPI_CommがMPI_COMM_WORLDの場合に相当)。一方で、MPIにおいては、プロセスを幾つかのグループに分け、各グループではグループ内のプロセスだけで閉じた通信を行う事が可能です(MPI_Comm_splitなどを利用)。しかし、グループ内のプロセスだけでScaLAPACKの関数を利用するには、少しコツが必要でした。試行錯誤の末に上手く動く方法が見つかったので、ここで共有します。

背景

ScaLAPACKは複数プロセスに並列化されたコードにおいて固有値ソルバーなどを提供してくれるライブラリです。プロセス間通信が前提になっていますが、おそらくMPIが主流になる以前から開発されてきた経緯もあり、プロセス間の通信部分はBLACS（BLASではないので注意）というライブラリが担っています。最近ではMPIが主流になり、BLACSはMPIの実質的なラッパーと言われています。ただし、問題の発端はここにあり、BLACSおよびScaLAPACKにはMPIの通信を司るコミュニケーターMPI_Commを直接渡すことが出来ません。代わりにICONTXTというBLACS独自のコミュニケーター相当の変数を利用します。また、ScaLAPACKおよびBLACSがFORTRANで記述されたライブラリということもあり、特にC/C++から利用する場合には、MPI_CommからBLACSのICONTXTへ変換が必要とされています。

良くある利用方法(C/C++からの場合)

C/C++ではMPI_CommからBLACSのICONTXTへ変換する方法として、Csys2blacs_handle()を使うという説明がされています。得られたICONTXTを使って、sl_init()もしくはblacs_gridinit()で初期化することになります。

//MPI_Comm型のmpi_commからBLACSのicontxtへ変換
int icontxt = Csys2blacs_handle(mpi_comm);

//BLACSの初期化
sl_init_(&icontxt, &np_rows, &np_cols);
//blacs_gridinit_(...)を使う場合もある

//以降でScaLAPACK関数を呼び出し
(略...)

//BLACSの終了//
blacs_gridexit_(&icontxt);

ここで2つの問題があります。

1. Csys2blacs_handle()はBLACSやScaLAPACKが提供する関数ではなく、システム依存であること。システムによっては提供されていないかもしれない。
2. Csys2blacs_handle()が正常に値を返すのはMPI_COMM_WORLDを受け取った場合だけ、という環境がある（少なくとも私のRocky Linux 8.6 + Intel MKLの環境）。MPI_Comm_splitなどで新しく発行されたMPIコミュニケーターを渡すと、無効な値(-1)を返す。

後者の問題において無効な値である$-1$が返された場合、以後のScaLAPACK関数の呼び出しでも当然失敗します。

MPI_Comm_splitとScaLAPACKを併用して動作させることが今回の目的です。

方法

説明の為に、ここではMPIで8プロセス並列をする場合を考えます。そして、プロセス0～3のグループ0、プロセス4～7のグループ1の2つに分け、各グループで独立にScaLAPACKを使う事にします。このような分割は、MPI_Comm_splitを用いて次のように行います。

//グローバルなMPIコミュニケーターはデフォルトでMPI_COMM_WORLDです
//グローバルな全プロセス数とプロセスIDを取得
int proc_id, num_procs;
MPI_Comm_rank(MPI_COMM_WORLD, &proc_id);   //proc_id is 0 to 7//
MPI_Comm_size(MPI_COMM_WORLD, &num_procs);  //num_procs = 8 //

//グループに分割//
int num_in_group = num_procs / 2;   //各グループのプロセス数//
int id_in_group = proc_id % num_in_group;  // 0 to 3 //
int group_id = proc_id / num_in_group;  // 0 or 1//
MPI_Comm group_comm;  //グループごとの新しいコミュニケーター
MPI_Comm_split(MPI_COMM_WORLD, group_id, id_in_group, &group_comm);

これで新しいMPIコミュニケーターとしてgroup_commが得られました。

今後はグループ内のプロセスをScaLAPACKと紐づけます。ScaLAPACKは扱う行列を2次元分割して各プロセスに処理を振り分けますが、ここでは各グループごとに4プロセスずつありますので、ScaLAPACKでは 担当領域を$2\times 2$分割することにします。

一方で、グループごとにどのようなプロセスがいるのかをBLACSに伝える手段として、先のgroup_commは渡せません。代わりに、blacs_gridmap()を利用して、グループごとの所属プロセスと、行列を$2\times 2$分割することを同時に伝えます。次のようにします。

//BLACSのデフォルトのICONTXTを取得//
int ZERO = 0;
int icontxt;
blacs_get_(&ZERO, &ZERO, &icontxt); 

//グループごとの担当プロセスを伝える配列//
std::vector<int> usermap(num_in_group);
for(int id = 0; id < num_in_group; ++id){
    usermap[id] = id + num_in_group * group_id;
}
//ここまでで
//group_id==0 => usermap={0,1,2,3}
//group_id==1 => usermap={4,5,6,7}

//行列をnp_rows * np_cols (2 * 2)で分割//
int np_rows = 2;
int np_cols = 2;

//BLACSの初期化(利用するプロセス(グローバルなID)を伝える)//
blacs_gridmap_(&icontxt, &usermap[0], &np_rows, &np_rows, &np_cols);

//例)グループごとに異なる条件でScaLAPACK関数を呼び出し
int matrixA[N*N] = {...}
int matrixB[N*N] = {...}
int matrixC[N*N] = {...}
if(group_id==0){
  pdsyevx( ... , matrixA, ...);  //group0ではmatrixAについて解く
}else{
  pdsyevx( ... , matrixB, ...); //group0とは異なるmatrixBを解く
  pdsyevx( ... , matrixC, ...); //呼び出し回数がグループごとに異なっても良い
}

//BLACSの終了//
blacs_gridexit_(&icontxt);

初期化にblacs_gridmap()を使うので、sl_init()やblacs_gridinit()は呼び出してはいけません。Csys2blacs_handle()も必要ありません。

ポイントはusermap配列にグループごとの所属プロセス一覧をセットしているところです。実際には上記のコードは全てのプロセスごとに実行されますので、プロセス自身の所属するgroup_idに合わせて{0,1,2,3}もしくは{4,5,6,7}をusermapに格納します。そして、blacs_gridmap()に渡すことで、以降のBLACSおよびScaLAPACKの通信がグループ内で閉じたものになります。

その後のScaLAPACKの呼び出しでは、グループごとに異なる行列を渡すことができます。また、関数を呼び出す回数がグループごとに異なっても動作しました(私の環境で試した限りでは)。

一見して不思議なのは、最初にblacs_get()でデフォルトのICONTEXTを取得して、以降はその変数icontxtを使い続けていることです。MPI_COMM_WORLDのように全プロセスでの通信に固定されるのでは？と思ってしまいます。しかし、私の環境ではblacs_gridmap()を呼ぶことでicontxtの中の数値が書き換わるので、何かしら新しいコミュニケーターになっているのでしょう。（それでも、書き換えられたicontxtの値がグループに依らず同じ値であるので、不思議さはぬぐえないですが）

注意点

コードの実装として、ScaLAPACKの関数のラッパーを作りなくなると思います。ただし、ScaLAPACK関数呼び出しの際に毎回blacs_gridmap()とblacs_gridexit()でICONTXTを生成消滅させていると、上の例の様にグループごとに呼び出し回数が異なる場合にデッドロックしました。ICONTXTを消滅させる前であればScaLAPACKの関数の呼び出し回数はグループごとに異なっても良さそうですが、BLACSのICONTXTの生成消滅の回数は全グループで揃える必要がありそうです。ですので、メインループの外などでのICONTXTの生成消滅を行うようにするのがよいでしょう。

おわりに

blacs_gridmap()を使うことで、複数プロセスをグループに分割した場合であっても、グループ内で閉じた通信としてScaLAPACK関数を利用できました。

最初に2つ挙げた問題のうち、Csys2blacs_handle()の件はC/C++特有のことかもしれませんが、それ以降の内容はFORTRANであっても有効化と思います。お役に立てば幸いです。

Reference

[1] IBM document: BLACS_GRIDMAP routine