個人的な勉強メモです。
本記事は"The Kokkos Lectures - module 1"1および"Kokkos Core Wiki"2を参考にしています。
Kokkosとは
複数のベンダーのデバイス・アーキテクチャ(CPU,GPU)に対応したコードを書くことのできるC++のプログラミングモデル
- CUDA, HIP, SYCL, HPX, OpenMP, C++スレッドなどのライブラリとして動作する
- GCC, Clang, NVCC, ROCMなどのコンパイラがサポートしている
- OSSプロジェクトであり、エコシステムとしてプロファイル・デバッグツール、数学ライブラリ、FortranやPythonとの統合ユーティリティも提供している
Kokkosの利点
レクチャー1に載っていた利点
- 一つのコードでパフォーマンスポータブル（異種のハードでの高速化）を実現
- OpenMPに比べシンプル
- CUDAやHIPに比べ簡単に書くことが可能
- データアクセスパターンを操作できることも特徴
- エコシステムが充実
また、Kokkosの目指すパフォーマンスポータビリティは以下とされている。
- 複数のアーキテクチャで動作する
- アーキテクチャの特性に合わせたメモリアクセスパターンの操作を行う
- アーキテクチャ独自のフィーチャーを可能な限り活かす
データ並列のコンセプト
OpenMP, OpenACCでの並列化
例えば、OpenMPによる並列化では以下のようにpragmaを挿入する。
#pragma omp parallel for
for(e = 0; e < ne; ++e) {
total = 0;
for(qp = 0; qp < nqp; ++qp) {
total += dot(left[e][qp], right[e][qp];
}
ev[e] = total;
}
ただし、上記はマルチコアCPUでの並列化の方法であり、GPUアーキテクチャでは以下のようになる。(OpenMP4.5)
GPUオフロードではデータのマッピングなどを細かく指定をする必要があるためやや複雑なコードとなる。
#pragma omp target data map(...)
#pragma omp team num_teams(...) num_threads(...) private(...)
#pragma omp distribute
for(e = 0; e < ne; ++e) {
total = 0;
#pragma omp parallel for
for(qp = 0; qp < nqp; ++qp) {
total += dot(left[e][qp], right[e][qp];
}
ev[e] = total;
}
また、同じコードはOpenACCでは以下のように書ける。
#pragma acc parallel copy(...) num_gangs(...) vector_length(...)
#pragma acc loop gang vector
for(e = 0; e < ne; ++e) {
total = 0;
for(qp = 0; qp < nqp; ++qp) {
total += dot(left[e][qp], right[e][qp];
}
ev[e] = total;
}
このように、OpenMPやOpenACCではポータビリティ（移植可能性）は満たすことができるのだが、実はKokkosの定義するパフォーマンスポータブルを満たすことは容易でない。これは、アーキテクチャにより適するメモリアクセスパターンが異なっており、その違いを吸収するための工夫が必要となるためである。（CPUのデータレイアウトのつもりでメモリアクセス操作を行うと、GPUのデータアクセス速度は大きく低下してしまう）
Kokkosでの並列化
Kokkosではメモリアクセスパターンをアーキテクチャに応じて操作できる仕組みがある。
Kokkosではworkという単位の処理をリソースに割り当てる。workの定義は以下である。
- 各イテレーションのbody(for文内のループbody)をworkの単位とする
- イテレーションのindexはworkの識別子を表す
- イテレーションのrangeはworkの全数を表す
例として、以下のコードでworkの割り当てを考える。
for (atomIndex = 0; atomIndex < numberOfAtoms; ++atomIndex) {
atomForces[atomIndex] = calculateForce(...data...); // work
}
Kokkosではworkを割り当てる仕組みとしてfunctorという構造体を導入している。全てのworkはKokkos::parallel_forを介して一つずつfunctorに割り当てられる。
ParallelFunctor functor;
Kokkos::parallel_for(numberOfIterations, functor);
...
// struct Functor {
// void operator()(const int64_t index)
// const {...}
// }
ここで、Kokkos runtimeはイテレーションを並列実行するが、イテレーションの順序は保証されない。
functorは自身のデータメンバを介して、bodyの任意のデータにアクセスすることができるようになっている。
struct AtomForceFunctor {
ForceType _atomForces;
AtomDataType _atomData; //データメンバ
AtomForceFunctor(ForceType atomForce, AromDataType data):
_atomForces(atomForces), _atomData(data) {} // コンストラクタ
void operator()(const int64_t atomIndex) const {
_atomForces[atomIndex] = culculateForce(_atomData);
}
}
この仕組みにより、ランタイムでデータアクセスパターンを操作することができる。
ここで再度、書き方の比較をしてみる。
- Serial
for (int64_t i = 0; i < N; ++i) {
/* loop body */
}
- OpenMP(CPU)
#pragma omp parallel for
for (int64_t i = 0; i < N; ++i) {
/* loop body */
}
- Kokkos
parallel_for(N, [=] (const int64_t i) {
/* loop body */
});
※[=]はc++11のLambdaによる書き方で、コンパイラによりfunctorに変換される。
このように、Kokkosはfunctorの構造体にloop bodyを入れていると考えると、わかりやすい書き方になっている。
シンプルな書き方でメモリアクセスパターンを制御することができるというのが利点である。