OpenMP OMPTでスレッドごとのtaskの実行数をカウントする

# OpenMP OMPTとは
OMPTはfirst-party toolがOpenMPプログラムのモニタリング等に使用するためのインタフェースであり、OpenMP API仕様[1]で定められています。ここでfirst-party toolとはOpenMPプログラムに直接リンクされるtoolのことを指します。ざっくり言えば(1) 対象のOpenMPプログラムを書き換えずにtoolのコードをリンクすることにより、(2) OpenMPの各種イベント(例えばparallel構文の開始時や終了時)に対してコールバックを登録し、(3) プログラムの内部動作をモニタリングする、というようなことができます。

なお仕様書にはOMPTが何の略か書かれていない(と思う)のですが、恐らくOpenMP Toolの略です。OMPTの詳細については仕様書を参照してください。

[1] OpenMP API Specification 5.2: https://www.openmp.org/wp-content/uploads/OpenMP-API-Specification-5-2.pdf

# 使い方
LLVMのOpenMPライブラリであるlibompではOMPTの機能が概ねサポートされています(※)。なおgccのOpenMPライブラリ(libgomp)では現時点でOMPTはサポートされていないようです。

以下の例ではLLVM 12を使用しています。

※ 最新のサポート状況: https://clang.llvm.org/docs/OpenMPSupport.html
(以下では適宜5.2の仕様書を参照していますが、上のページの通り現時点ですべての仕様がサポートされているわけではありません)

## 簡単な例
まずモニタリング対象のOpenMPプログラムを用意します。スレッド毎にprintfするだけの以下のコードを使います。

```c:hello.c
#include <stdio.h>
#include <omp.h>

int main() {
#pragma omp parallel
	{
		int id = omp_get_thread_num();
		int num_threads = omp_get_num_threads();

		printf("Hello %d/%d\n", id, num_threads);
	}
	return 0;
}
```

```bash:実行結果
$ clang -fopenmp hello.c -o hello
$ OMP_NUM_THREADS=4 ./hello
Hello 0/4
Hello 1/4
Hello 3/4
Hello 2/4
```

次にOMPTでparallel構文の開始時と終了時にコールバック関数を登録してprintfを入れてみます。対象のコードとは別のファイルに以下のOMPTの処理コードを書きます (同じファイルにOMPTの処理を記述しても問題はありませんが、対象のコードは変更しない想定です)。

```c:ompt.c
#include <stdio.h>
#include <omp-tools.h>

// parallel構文の開始に呼ばれるコールバック
// Spec 5.2: 19.5.2.3
static void my_ompt_callback_parallel_begin(
		ompt_data_t *encountering_task_data,
		const ompt_frame_t *encountering_task_frame,
		ompt_data_t *parallel_data,
		uint32_t requested_team_size,
		int flags,
		const void *codeptr_ra)
{
	printf("parallel begin\n");
}

// parallel構文の終了時に呼ばれるコールバック
// Spec 5.2: 19.5.2.4
static void my_ompt_callback_parallel_end(
		ompt_data_t *parallel_data,
		ompt_data_t *encountering_task_data,
		int flags,
		const void *codeptr_ra)
{
	printf("parallel end\n");
}

// OMPT初期化時に呼ばれる処理
// Spec 5.2: 19.5.1.1
int my_ompt_initialize(ompt_function_lookup_t lookup,
					int initial_device_num,
					ompt_data_t *tool_data) {
	printf("ompt_initialize\n");

	// 19.6.1.3 for ompt_set_callback
	// 19.6.3 for lookup
	ompt_set_callback_t ompt_set_callback = 
		(ompt_set_callback_t)lookup("ompt_set_callback");

	// should check return value
	ompt_set_callback(ompt_callback_parallel_begin,
				(ompt_callback_t)my_ompt_callback_parallel_begin);
	ompt_set_callback(ompt_callback_parallel_end,
				(ompt_callback_t)my_ompt_callback_parallel_end);

	// return non-zero for success
	return 1;
}

// OMPT終了時に呼ばれる処理
// Spec 5.2: 19.5.1.2
void my_ompt_finalize(ompt_data_t *tool_data) {
	printf("ompt_finalize\n");
}

static ompt_start_tool_result_t my_ompt_start_tool_result = {
	&my_ompt_initialize,
	&my_ompt_finalize,
	0, // toolで使用する値を保持できる (未使用)
};

// 最初に呼ばれる処理
// Spec 5.2: 19.2.1
ompt_start_tool_result_t* ompt_start_tool(
		unsigned int omp_version,
		const char *runtime_version)
{
	printf("omp version: %u, runtime_version: %s\n", omp_version, runtime_version);
	return &my_ompt_start_tool_result;
}
```

ポイントは以下です。

 - OMPTのコードで使用する型の定義はomp-tools.hにあるためこれをincludeします
    - Fedoraの場合libomp-develパッケージの中に含まれています

```
$ dnf repoquery -l libomp-devel | grep omp-tools.h
/usr/lib64/clang/12.0.1/include/omp-tools.h
```

 - OMPTを有効にするためにはompt_start_tool関数を定義する必要があります。この関数で有効なompt_start_tool_result_tを戻り値として返します。ompt_start_tool_result_tにはomptの初期化時/終了時に呼び出される処理を登録します
 - OMPT初期化時に呼び出される関数(ここではmy_ompt_initialize)においてコールバック関数の登録を行います
     - OMPT APIの関数はヘッダファイルには定義されておらず、使用したいAPI関数は初期化関数に渡されるlookupを利用して関数ポインタを取得する必要があります。まずコールバック関数を設定するためのompt_set_callbackをlookupしています
     - 次にlookupしたompt_set_callbackを利用して、parallel構文の開始時と終了時にコールバックを設定します。それぞれのコールバックでは単純にprintfを呼んでいるだけです
 - OMPT終了時に呼び出される関数(ここではmy_ompt_finalize)で、必要に応じてクリーンアップを行います

各種コールバック関数の種類や引数の意味などはすべて仕様書に定められている通りです。

後はこのファイル(ompt.c)を目的のプログラムにリンクするだけです。つまり単純に一緒にコンパイルしたり、shared libaryを作成してコンパイル時にリンクしたりすればよいです。が、仕様書をよく読むと実は目的のバイナリに直接リンクしなくともOMP_TOOL_LIBRARIES環境変数にshared libraryを渡すことでもOMPTを使用できると書かれています。ここではそちらの方法を試してみます。

```bash:実行結果
$ clang -fPIC -shared ompt.c -o ompt.so
$ clang -fopenmp hello.c -o hello
$ OMP_TOOL_LIBRARIES=./ompt.so OMP_NUM_THREADS=4 ./hello
omp version: 201611, runtime_version: LLVM OMP version: 5.0.20140926
ompt_initialize
parallel begin
Hello 0/4
Hello 1/4
Hello 3/4
Hello 2/4
parallel end
ompt_finalize
```

登録したコールバック関数が呼ばれていることが分かります。

(簡単に実行時の流れを説明します。まず#pragma omp...の構文はコンパイラによりランタイムライブラリ関数の呼び出しに変換されています。最初にライブラリ関数が呼ばれたタイミングなどでライブラリの初期化処理が走りますが、その中でompt_start_toolのシンボルがプロセスのアドレス空間に存在するかやOMP_TOOL_LIBRARIESで指定されたshared libraryが読み込めるかをチェックし、OMPT初期化処理を行います。そしてラインタイム処理の各要点においてコールバック関数が設定されている場合はそれを呼び出します)

## スレッドごとにtaskの実行数をカウントする
もう一つ別の例を見てみます。

OpenMPではtask構文を利用して動的にworker threadが実行する処理単位を生成することができます。ここではtask構文で生成されたtaskを実際に各worker threadがいくつ実行したのかをOMPTを利用してカウントしてみます。

まず対象のプログラムを用意します。taskを100個生成するプログラムです (意味はありません)。

```c:task.c
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>

void f() {
	int time = (rand() % 10000); // 0 ~ 9999us

	usleep(time);
}

int main() {
#pragma omp parallel
#pragma omp single
	{
		for (int i = 0; i < 100; i++) {
#pragma omp task
			f();
		}
	}
	return 0;
}
```

次に仕様書を読むと以下のことが分かります

 - taskが完了した時はtask_scheduleコールバックが呼ばれ、そのとき引数のprior_task_statusがompt_task_completeである
 - スレッドの生成/終了時にはそれぞれthread_begin/thread_endのコールバックが呼ばれる
    - このときそのスレッドに固有のデータ領域として引数として渡される"ompt_data_t *thread_data"(の指す領域)が使用できる
    - 他のOMPTのコールバック関数において、ompt_get_thread_data関数を呼び出すことで実行スレッドのthread_dataを取得することができる

よって

- ompt_data_tはuint64とvoid *のunion型になっているので、ここではuint64とみなしてそのスレッドで実行したtaskの数をカウントするカウンタして利用する (もしより大きな領域が必要な場合はメモリをallocしてそのアドレスをthread_dataに保持する)
- thread_beginコールバックの中でthread_dataのカウントを0にする (ただし仕様として初期化時は0になっているはず)
- task_scheduleコールバックの中でstatusがompt_task_completeならばthread_dataを取得してカウンタをカウントアップする
- thread_endコールバックの中で合計のカウントをprintfする

という処理を行なえばプログラムの終了時に各worker threadで実行されたtaskの数が分かるはずです（プログラムによってはtaskの実行中に別のtaskに切り替わることなどが有り得ますが、そのような複雑な状況はここでは考慮していません）。

コードは以下のようになります。

```c:ompt-task.c
#include <stdio.h>
#include <omp-tools.h>

// ompt_get_thread_dataもlookup経由で取得する
// Spec 5.2: 19.6.1.5
ompt_get_thread_data_t ompt_get_thread_data;

// Spec 5.2: 19.5.2.1
static void my_ompt_callback_thread_begin(
		ompt_thread_t thread_type,
		ompt_data_t *thread_data)
{
	thread_data->value = 0;
}

// Spec 5.2: 19.5.2.2
static void my_ompt_callback_thread_end(
		ompt_data_t *thread_data)
{
	printf("executed tasks: %ld\n", thread_data->value);
}

// Spec 5.2: 19.5.2.10
static void my_ompt_callback_task_schedule(
		ompt_data_t *first_task_data,
		ompt_task_status_t prior_task_status,
		ompt_data_t *second_task_data)
{
	// count executed task upon completion
	if (prior_task_status == ompt_task_complete) {
		ompt_data_t *thread_data = ompt_get_thread_data();
		thread_data->value++;
	}
}

int my_ompt_initialize(ompt_function_lookup_t lookup,
					int initial_device_num,
					ompt_data_t *tool_data) {
	printf("ompt_initialize\n");

	ompt_set_callback_t ompt_set_callback = 
		(ompt_set_callback_t)lookup("ompt_set_callback");
	ompt_get_thread_data =
		(ompt_get_thread_data_t)lookup("ompt_get_thread_data");

	// should check error
	ompt_set_callback(ompt_callback_thread_begin,
				(ompt_callback_t)my_ompt_callback_thread_begin);
	ompt_set_callback(ompt_callback_thread_end,
				(ompt_callback_t)my_ompt_callback_thread_end);
	ompt_set_callback(ompt_callback_task_schedule,
				(ompt_callback_t)my_ompt_callback_task_schedule);

	return 1;
}

void my_ompt_finalize(ompt_data_t *tool_data) {
	printf("ompt_finalize\n");
}

static ompt_start_tool_result_t my_ompt_start_tool_result = {
	&my_ompt_initialize,
	&my_ompt_finalize,
	0,
};

ompt_start_tool_result_t* ompt_start_tool(
		unsigned int omp_version,
		const char *runtime_version)
{
	printf("omp version: %u, runtime_version: %s\n", omp_version, runtime_version);
	return &my_ompt_start_tool_result;
}
```

先程と同様にコンパイル、実行します。

```bash:実行結果
$ clang -fPIC -shared ompt-task.c -o ompt-task.so
$ clang -fopenmp task.c -o task
$ OMP_TOOL_LIBRARIES=./omp-task.so OMP_NUM_THREADS=4 ./task
omp version: 201611, runtime_version: LLVM OMP version: 5.0.20140926
ompt_initialize
executed tasks: 27
executed tasks: 21
executed tasks: 24
executed tasks: 28
ompt_finalize
```
期待通りに動いているようです。

## おわりに
簡単なOMPTの例を見てみました。ある処理の実行回数やあるブロックの実行時間の測定などは比較的簡単に行えそうです。なおOMPTのコールバックからOpenMP runtimeのAPI(例えばomp_get_thread_num())を呼ぶことができないと書かれている(Spec 5.2, 19.5)のでその点は注意が必要です。