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()`で初期化することになります。
```C++
//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が提供する関数ではなく、システム依存であること。システムによっては提供されていないかもしれない。
1. `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`を用いて次のように行います。
```C++
//グローバルな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$分割することを同時に伝えます。次のようにします。
```C++
//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](https://www.ibm.com/docs/en/pessl/5.5?topic=blacs-gridmap-routine)