Kokkos1の勉強をしています。本記事は"The Kokkos Lectures - module 2"2を参考にしています。
前記事ではコードに触れていなかったので、見ていきたいと思います。
Tutorial Exercise 01
Exercise 013では、Beginの通常コードをKokkosの対応コードに書き換えるという問題でした。
まず、Kokkosに対応するため、Kokkos_Core.hppをインクルードし、Kokkosランタイムのinitialize()とfinalize()を挿入します。
Solutionのコード
#include <Kokkos_Core.hpp>
:
Kokkos::initialize( argc, argv );
{
:
}
Kokkos::finalize();
次にループの対応箇所を見ていきます。
配列x, yの初期化は1次元のループ、配列Aの初期化は2次元のループです。
Beginのコード
// Initialize y vector.
// EXERCISE: Convert outer loop to Kokkos::parallel_for.
for ( int i = 0; i < N; ++i ) {
y[ i ] = 1;
}
// Initialize x vector.
// EXERCISE: Convert outer loop to Kokkos::parallel_for.
for ( int i = 0; i < M; ++i ) {
x[ i ] = 1;
}
// Initialize A matrix, note 2D indexing computation.
// EXERCISE: Convert outer loop to Kokkos::parallel_for.
for ( int j = 0; j < N; ++j ) {
for ( int i = 0; i < M; ++i ) {
A[ j * M + i ] = 1;
}
}
これをKokkosで書き換えたコードは以下です。
Solutionのコード
// Initialize y vector.
Kokkos::parallel_for( "y_init", N, KOKKOS_LAMBDA ( int i ) {
y[ i ] = 1;
});
// Initialize x vector.
Kokkos::parallel_for( "x_init", M, KOKKOS_LAMBDA ( int i ) {
x[ i ] = 1;
});
// Initialize A matrix, note 2D indexing computation.
Kokkos::parallel_for( "matrix_init", N, KOKKOS_LAMBDA ( int j ) {
for ( int i = 0; i < M; ++i ) {
A[ j * M + i ] = 1;
}
});
KOKKOS_LAMBDAは[=]と書いても大丈夫です。
KOKKOS_LAMBDAの式はそれぞれのiまたはjについてfunctorに割り当てられ、演算されます。
参考）Kokkos::parallel_forの使い方4
Usage
Kokkos::parallel_for(name, policy, functor);
Timerの計測している箇所を見ていきます。
Beginのコード
// EXERCISE: Convert outer loop to Kokkos::parallel_reduce.
for ( int j = 0; j < N; ++j ) {
double temp2 = 0;
for ( int i = 0; i < M; ++i ) {
temp2 += A[ j * M + i ] * x[ i ];
}
result += y[ j ] * temp2;
}
これをKokkosで書き換えたコードは以下です。
Kokkos::parallel_reduce( "yAx", N, KOKKOS_LAMBDA ( int j, double &update ) {
double temp2 = 0;
for ( int i = 0; i < M; ++i ) {
temp2 += A[ j * M + i ] * x[ i ];
}
update += y[ j ] * temp2;
}, result );
リダクション演算なので、parallel_reduceを使い、一時変数としてupdateを用意します。
KOKKOS_LAMBDAの式はそれぞれのjについてfunctorに割り当てられ、演算されます。
最後にresultへ値を格納します。
KOKKOS_LAMBDAは[=]と書いても大丈夫です。
参考）Kokkos::parallel_reduceの使い方5
Usage
Kokkos::parallel_reduce(const std::string& name,
const ExecPolicy& policy,
const FunctorType& functor);
以上が、Tutorial Exercise 01のコード対応箇所でした。