はじめに

CUDAプログラミングでは、関数の前に関数識別子__host__ __device__を付けることで、関数をホスト・デバイスのどちらでも実行できるようになります。関数をホスト・デバイスの両方で実行できると、「GPU計算がCPU計算に対して何倍高速なのか?」といったことを簡単に検証できるため非常に便利です。ただしスレッド同期を使用する関数(リダクション処理や行列積計算など)では、ホストとデバイスで実行した際に挙動が変化してしまい異なる計算結果になってしまう、という問題点が発生してしまいます。

そこで今回、スレッド同期を含む場合でもホスト・デバイス両方で正常に動作する関数を実現するために、c++20で実装されたコルーチン機能を使用した実装を行います。
(前置きが長いので、実装手法だけ知りたい方はcoroutineを使って同期関数と同等の動作をホスト側で再現するまで読み飛ばしてください。)

CUDAにおける同期関数

CUDAでは、同一ブロック内のスレッド間で同期を取るデバイス関数 __syncthreads()が定義されています。ここでは下記に示すリダクション処理を例に取り、その挙動を見てみます。
まず、リダクション処理を行う関数の実装例を下記に示します。

#define BLOCK_N 128

// reduction function
__host__ __device__
void calcSumReduction(float *s_array, const float *array,
                      const int index) {
  // first reduction
  int shift = BLOCK_N/2;
  s_array[index] = array[index] + array[index+shift];

  // syncthreads
  __syncthreads();

  // iterative reduction
  while (shift > 1) {
    shift >>= 1;
    if (index < shift) {
      s_array[index] = s_array[index] + s_array[index+shift];
    }
  
    // syncthreads
    __syncthreads();
  }

  // output result
  if (index == 0) printf("result = %3.3e\n", s_array[0]);
}

// device kernel
__global__
void calcSumReductionOnDevice(const float *d_array) {
  // shared array
  extern __shared__ float s_array[];

  // index
  const int index = threadIdx.x;

  // calc sum reduction
  calcSumReduction(s_array, d_array, index);
}

// main function
int main(int argc, char **argv) {
  // allocation
  float *h_array, *d_array;
  h_array = new float[BLOCK_N];
  cudaMalloc(&d_array, sizeof(float)*BLOCK_N);

  // initialization
  for (int i = 0; i < BLOCK_N; i++) {
    h_array[i] = 1.0;
  }
  cudaMemcpy(d_array, h_array, sizeof(float)*BLOCK_N, cudaMemcpyHostToDevice);

  // call device kernel
  calcSumReductionOnDevice<<<1, BLOCK_N, sizeof(float)*BLOCK_N/2>>>(d_array);
  cudaDeviceSynchronize();
}

result = 1.280e+02

上記コード内で実際にリダクション処理を行っているのがcalcSumReduction関数ですが、その中で同期関数__syncthreads()が何度か呼び出されています。ここでの同期関数の役割としては、「シェアードメモリ(スレッド間で共有するメモリ)へのアクセスをスレッドセーフにすること」になります。

同期関数の役割をもう少し明確にするために、下図を使って説明したいと思います(東工大のCUDA講習会の資料¹を参考にしました)。ここでは簡単のため、4スレッドを用いて長さ8の配列に格納された値の総和を求める際の処理フローを示しています。

具体的な処理フローは下記のようになります。

スレッド0-3が配列の0-7番地に格納された値を足し合わせ、シェアードメモリの0-3番地に格納する。
同期関数を呼び出す。
スレッド0-1がシェアードメモリの0-3番地に格納された値を足し合わせシェアードメモリの0-1番地に格納する。
同期関数を呼び出す。
スレッド0がシェアードメモリの0-1番地に格納された値を足し合わせシェアードメモリの0番地に格納する。

処理フロー中、シェアードメモリへの値格納の後に同期関数が呼び出されていますが、これにより、例えばスレッド1がシェアードメモリに値を格納する前にスレッド0が当該シェアードメモリ番地にアクセスしてしまう、といったトラブルを防止することができるようになります。

このことから分かるように、特にシェアードメモリを使用してスレッド間で情報を共有する場合には、同期関数の呼び出しは不可欠と言えます。

スレッド同期を含む関数をホスト実行する際の問題

話を戻します。
今回行いたいのは、前節で説明したようなスレッド同期を含むデバイス関数のホスト対応ですが、残念ながら関数識別子__host__を宣言するだけではホスト実行はおろか、コンパイルすることすらできません。

calling a __device__ function("__syncthreads") from a __host__ __device__ function is not allowed

先述のように__syncthreads()はデバイス専用関数なので、ホスト側から呼び出すことができずコンパイルエラーになります。「じゃあ__CUDA_ARCH__マクロでスイッチングしてホスト側で無効化すれば良いのでは」と思われる方がいるかもしれませんが(コンパイル自体はこれで通ります)、calcSumReduction関数をforループで回しても正しくリダクション操作を行うことができません。

なぜリダクション操作が正しく行われないのか、先ほどの図をもとに説明します。calcSumReduction関数をホスト側で実行する際には、デバイス実行の際に各スレッドが担っていた操作を順々に実施していくことになります。しかしホスト実行の場合には、各スレッドがシェアードメモリに値を格納した直後の同期操作(デバイス関数では__syncthreads()が担っていました)や関数の一時中断が行われないために次の処理に進んでしまい、まだ値が格納されていない要素にアクセスしてしまうという意図しない動作が行われてしまうわけです。

勿論、openMPやMPIまたはstd::threadなどを使ってマルチスレッディングでCUDAのようなリダクション操作をホスト実行することはできますが、これらのテクニックを使用する場合にはスレッド上限の存在という制約が付きまといますし、そもそもホスト関数とデバイス関数で実装が大幅に変わってしまうため本末転倒です。

coroutineを使って同期関数と同等の動作をホスト側で再現する

大分前置きが長かったですが、ここからはタイトルにもある「スレッド同期を含むデバイス関数のホスト実行」の実現手法について説明していきます。

ホスト関数とデバイス関数が共通化できない要因については前節で述べた通りですが、今回はc++20で実装された coroutine² を悪用利用することでスレッド同期を含むデバイス関数を無理やりホスト側でも実行できるようにします。

coroutineについて

coroutineの詳細についてはここでは割愛しますが、端的に言えば「いったん処理を中断した後に続きから処理を再開できる」といった動作をサポートする機能になります。この機能を応用すると、例えば同一の関数を複数起動させたのちに、指定した箇所で全ての関数を一時中断/再開させるといったスレッド同期さながらの動作を実現することができます(ただし、実際にはマルチスレッドではなくシングルスレッドで実行されます)。

サンプルコード

では、coroutineを使用してホスト・デバイス両対応させたリダクション関数の実装を見てみましょう。calcSumReduction関数に関しては、CUDAにおける同期関数で記述したものに少しだけ加筆・修正したものになります。

#define BLOCK_N 128

// reduction function
__host__ __device__
Generator calcSumReduction(float *s_array, const float *array,
                           const int index) {
  // first reduction
  int shift = BLOCK_N/2;
  s_array[index] = array[index] + array[index+shift];

  // syncthreads
#ifdef __CUDA_ARCH__
  __syncthreads();
#else
  co_yield 0;
#endif
  
  // iterative reduction
  while (shift > 1) {
    shift >>= 1;
    if (index < shift) {
      s_array[index] = s_array[index] + s_array[index+shift];
    }
  
    // syncthreads
#ifdef __CUDA_ARCH__
    __syncthreads();
#else
  co_yield 0;
#endif
  }

  // output result
  if (index == 0) printf("result = %3.3e\n", s_array[0]);
}

// host kernel
void calcSumReductionOnHost(const float *h_array) {
  // shared array
  float *s_array = new float[BLOCK_N/2];

  // generate thread
  std::vector<Generator> threads;
  for (int index = 0; index < BLOCK_N/2; index++) {
    threads.emplace_back(calcSumReduction(s_array, h_array, index));
  }

  // run function
  bool adapt_flag = true;
  while ( adapt_flag ) {
    for ( auto &thread : threads ) {
      adapt_flag = thread.resume();
    }
  }
}

// device kernel
__global__
void calcSumReductionOnDevice(const float *d_array) {
  // shared array
  extern __shared__ float s_array[];

  // index
  const int index = threadIdx.x;

  // calc sum reduction
  calcSumReduction(s_array, d_array, index);
}

// main function
int main(int argc, char **argv) {
  // allocation
  float *h_array, *d_array;
  h_array = new float[BLOCK_N];
  cudaMalloc(&d_array, sizeof(float)*BLOCK_N);

  // initialization
  for (int i = 0; i < BLOCK_N; i++) {
    h_array[i] = 1.0;
  }
  cudaMemcpy(d_array, h_array, sizeof(float)*BLOCK_N, cudaMemcpyHostToDevice);

  // call host kernel
  calcSumReductionOnHost(h_array);

  // call device kernel
  calcSumReductionOnDevice<<<1, BLOCK_N, sizeof(float)*BLOCK_N/2>>>(d_array);
  cudaDeviceSynchronize();
}

まず、calcSumReduction関数を見てみると、デバイス関数(__CUDA_ARCH__マクロ有効時)では従来通り__syncthreads()が呼び出される一方で、ホスト関数(__CUDA_ARCH__マクロ無効時)ではco_yieldが定義されていることが分かります。co_yieldは、現在実行している関数を途中で中断し呼び出し元に値を返すためのキーワードであり、ブレークポイントの役割を担っています。また、返り値で使用されているGeneratorはcoroutineの戻り値や関数への参照などを管理するクラスであり、ホスト側実行でのみ使用されます(Generatorクラスの詳細な実装については、最後に示したいと思います)。

次に、今回新しく作成したcalcSumReductionOnHost関数を見てみましょう。この関数では、初めにGeneratorクラスのベクトル型が定義され、これに対してcalcSumReduction関数の戻り値を次々と格納しています。なおこの時点で、各indexに対応したcalcSumReduction関数が最初のブレークポイント(first reduction)まで実行されます。その後、Generatorクラスのメンバ変数であるresume()によって中断されていたcalcSumReduction関数の処理が次々と再開されます。この中断・再開の一連の動作はwhileループを抜けるまで(calcSumReduction関数の全処理が終了するまで)繰り返し実行されます。

この時の動作について、再びリダクション処理の図を参考に見てみましょう。今回の場合には、デバイス関数で__syncthreads()を定義していた箇所に代替としてco_yieldキーワードを定義しています。これにより、2回目以降のリダクション操作(co_yieldの直後)は、すべてのループにおいて直前の動作(シェアードメモリへの値の格納)の終了後に実施されます。すなわち、ホスト関数におけるco_yieldはデバイス関数における__syncthreads()と実質的に同義になり、それぞれにおける計算結果も同じになります。

さて、もう一度calcSumReduction関数を見てみましょう。今回の目的としては「スレッド同期を含むデバイス関数をホスト側でも実行できるようにする」ことでしたが、

同期/中断処理を表す__syncthreads()/co_yield以外はホスト関数とデバイス関数で共通化できている
__syncthreads()/co_yieldは__CUDA_ARCH__マクロでホスト関数/デバイス関数生成時に自動的に切り替えられる
ことを鑑みると、実質的にホスト関数とデバイス関数の共通化を実現することができたのではないでしょうか。



__host__ __device__
Generator calcSumReduction(float *s_array, const float *array,
                           const int index) {
  // first reduction
  int shift = BLOCK_N/2;
  s_array[index] = array[index] + array[index+shift];

  // syncthreads
#ifdef __CUDA_ARCH__
  __syncthreads();
#else
  co_yield 0;
#endif
  
  // iterative reduction
  while (shift > 1) {
    shift >>= 1;
    if (index < shift) {
      s_array[index] = s_array[index] + s_array[index+shift];
    }
  
    // syncthreads
#ifdef __CUDA_ARCH__
    __syncthreads();
#else
  co_yield 0;
#endif
  }

  // output result
  if (index == 0) printf("result = %3.3e\n", s_array[0]);
}

なお、サンプルコードの実行結果は下記のようになります。ホスト実行とデバイス実行でちゃんと同じ結果が出力されていますね。

result = 1.280e+02 (ホスト実行)
result = 1.280e+02 (デバイス実行)

おわりに

というわけで、coroutineを使ってスレッド同期を含むデバイス関数をホスト側でも実行できるようにすることができました。わざわざGPU対応させたプログラムをCPUで動かす奇特な人もあまりいない気もしますが、実装される場合にはご参考まで。
なおc++のcoroutine機能を使用したプログラムは、gcc11.0またはCUDA12.0以降でしかコンパイルできません。特にCUDA12.0についてはつい最近リリースされたばかりなので、あまり気軽に試せるものでもないですね。

おまけ

coroutineの戻り値や関数への参照などを管理するクラスGeneratorの実装を下記に示します。基本的には参照元²のパクリですが、コンストラクタなどの一部関数はデバイス側で呼び出せないため、__CUDA_ARCH__マクロで無効化しています。

struct Generator {

public :
  struct promise_type;
  using handle = std::coroutine_handle<promise_type>;
  handle coro;

  struct promise_type {

   public:
    int current_value; // hold value returned by co_yield

   public:
    static auto get_return_object_on_allocation_failure() { return Generator{ nullptr }; }
    auto get_return_object() { return Generator{ handle::from_promise(*this) }; }
    auto initial_suspend() { return std::suspend_always{}; }
    auto final_suspend() noexcept { return std::suspend_always{}; }
    void unhandled_exception() { std::terminate(); }
    void return_void() {}
    auto yield_value(const int value) {
      current_value = value;
      return std::suspend_always{};
    }

  };

public :
  bool resume() { return (coro) ? (coro.resume(), !coro.done()) : false; }
  int current_value() { return coro.promise().current_value; }
#ifndef __CUDA_ARCH__ // constructor & destructor are available when called by host
  Generator(handle h) : coro(h) {}
  Generator(Generator const&) = delete;
  Generator(Generator && rhs) : coro(rhs.coro) { rhs.coro = nullptr; }
  ~Generator() { if (coro) coro.destroy(); }
#endif

};

参考