Thrustの非同期実行 #C++

初めまして。
趣味で雪の物理シミュレーションをC++で書いています。
アドベントカレンダーが空いていたので軽い気持ちで参加しました。
よろしくお願いします。

Thrustの非同期実行についてですが、日本語の文献が
thrustにasyncサポートが入っていた
くらいしか見当たらなかったので、具体例を交えながら環境構築から解説したいと思います。

環境構築

CMakeを使うやり方をご紹介します。
Windowsでは試していないため、ご了承下さい。

以下のソフトウェアをインストールしてください。

CUDA >= 12.3
CMake >= 3.27.7
gccまたはclang

Gitもインストールをお勧めしますが、なくても問題ありません。

Thrustのダウンロード

Thrustをダウンロードします。
1.9.4以上のバージョンを選んで下さい。
Gitを使う方は再帰的にサブモジュールをクローンしてください。

なぜソースコードをダウンロードするかというと、CMakeがThrustを見つけてくれなかったためです。
見つけ方をご存じの方は教えていただけると幸いです。

CMakeプロジェクトの作成

例として、floatとintの配列を作成し、GPUで各要素の2乗を計算し、結果が正しいか確認するプログラムを作成します。
以下のようにCMakeプロジェクトを作成します。

hoge/-- CMakeLists.txt
     |- double.cu
     |- double.h
     |- main.cpp
     |- thrust/

CMakeLists.txtに以下を書きます。

CMakeLists.txt

cmake_minimum_required(VERSION 3.27.7)

# CUDAアーキテクチャを指定する。
# 最初に設定しないと、プロジェクトが作成できません。
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
  set(CMAKE_CUDA_ARCHITECTURES native)
endif()

project(my_program VERSION 0.1.0 LANGUAGES CXX CUDA)

# ここの設定はお好みで。
set(CMAKE_CXX_EXTENSIONS OFF)
set(CMAKE_OPTIMIZE_DEPENDENCIES ON)

# CUDAを見つける。
find_package(CUDAToolkit 12 REQUIRED)

# Thrustを見つける。
add_subdirectory(thrust)
find_package(Thrust REQUIRED CONFIG)
thrust_create_target(Thrust)

# CUDAライブラリを作る。
add_library(double double.cu)
target_compile_options(
  double
  PRIVATE
    # ラムダ式をデバイスコードで使えるようにする。
    -expt-extended-lambda
    # SIMDを無効にする。
    # SIMDの無効化は必要ありませんが、Eigenなど線形代数ライブラリを使う際に必要になります。
    "$<$<COMPILE_LANG_AND_ID:CXX,GNU>:-fno-tree-vectorize>"
    "$<$<COMPILE_LANG_AND_ID:CXX,Clang>:-fno-vectorize>"
)
# ここはお好みで。
target_compile_features(
  double
  PRIVATE
    cuda_std_20
)
# 誰もインストールしないはずなので、簡略に。
target_include_directories(
  double
  PUBLIC 
    ${CMAKE_CURRENT_SOURCE_DIR}
)
# CUDAランタイムとThrustをリンクする。
# 当たり前ですが、ThrustにはCUDAはないので。
target_link_libraries(
  double
  PUBLIC
    CUDA::cudart
    Thrust
)

# my_programを作る。
add_executable(my_program main.cpp)
target_link_libraries(
  my_program
  PRIVATE
    double
)
# ここはお好みで。
target_compile_features(
  my_program
  PRIVATE
    cxx_std_20
)
# ここはお好みで。
target_compile_options(
  my_program
  PRIVATE
    "$<$<CONFIG:Debug>:-g;-O0;-Wall;-Wextra>"
    "$<$<CONFIG:Release>:-DNDEBUG;-O3;-march=native;-mtune=native>"
)
# CUDA 5.0以前のバージョンでは、デバイスコードで使用する関数の実装までインクルードする必要が
# ありましたが、CUDA 5.0以降では、デバイスコードで使用する関数の宣言のみをインクルードするだけで
# 十分になりました。
# このことをコンパイラとリンカーに伝えるために、CUDA_SEPARABLE_COMPILATIONプロパティをONに
# しています。
# ただし、このプロパティは実行ファイルにのみONに設定して下さい。
# 詳しくは参考文献1を参照してください。
set_property(TARGET my_program PROPERTY CUDA_SEPARABLE_COMPILATION ON)

Thrustの非同期実行の書き方

いよいよ本題です。

eventとfuture

Thrustにはeventとfutureがあります。
futureはreduceから返され、その他のアルゴリズムはeventを返します。
どちらもwait()メンバ関数を持ち、非同期実行が終わるまで待ちます。
futureはさらにget()メンバ関数を持ち、実行結果を取得できます。
詳しく知りたい方は参考文献2のプログラムをご覧ください。
eventはunique_eager_eventクラスに, futureはunique_eager_futureクラスに対応しています。

ヘッダーはそれぞれthrust/event.h, thrust/future.hです。
今回は計算結果をhost_vectorに格納するので、使うのは実行が完了したか確認するだけのeventです。

double.h

#include <thrust/event.h>
#include <thrust/host_vector.h>

thrust::device_event Double(thrust::host_vector<float>& doubles, thrust::host_vector<int>& ints);

ここでdevice_eventを使っていますが、ホストでも使えます。
host_eventはありません。
後の欠点の節でこのことに触れます。
__global__はあってもなくてもいいです。

非同期なアルゴリズムと実行ポリシー

Thrustにはcopyやtransformといった、STLのような関数が用意されています。
これらを非同期に実行したい場合は、async名前空間のものを使えばいいです。
例えば、thrust::copyの非同期版はthrust::async::copyです。

実行ポリシーとは、アルゴリズムにホストで実行するのか、デバイスで実行するのか、それとも逐次的に処理するのかを指定するクラスです。
クラス名はそれぞれhost、device、seqです。
アルゴリズムに明示的に指定することで、無駄なデータのコピーを減らせます。
実装はタグディスパッチで、自分でカスタマイズすることも可能です。
詳しくは参考文献3をご覧ください。

実行ポリシー自体は同期的なアルゴリズムに対しても使えますが、非同期ならhostとdeviceのafterメンバ関数にeventとfutureを指定することで実行順序を指定できます。
eventとfutureの個数は何個でもいいです。
また、when_all関数で複数のeventとfutureを一つのeventにまとめることもできます。

これらを踏まえて、以下のように実装してみます。

double.cu

#include <thrust/async/copy.h>
#include <thrust/async/transform.h>
#include <thrust/device_vector.h>

#include "double.h"

template <class T>
thrust::device_event Double(thrust::host_vector<T>& ts) {
  // デバイス側の配列を用意
  auto device_ts = thrust::device_vector<T>();

  // メモリの確保と初期化
  device_ts.resize(ts.size());

  // ホストからデバイスへの非同期コピー
  auto copy_ts_event = thrust::async::copy(
      thrust::host,
      thrust::device,
      ts.begin(),
      ts.end(),
      device_ts.begin()
  );

  // デバイス側での計算
  auto double_ts_event = thrust::async::transform(
      thrust::device.after(copy_ts_event),
      device_ts.begin(),
      device_ts.end(),
      device_ts.begin(),
      [] __device__(T d) { return d * d; }
  );

  // デバイスからホストへの非同期コピー
  auto copy_back_ts_event = thrust::async::copy(
      thrust::device.after(double_ts_event),
      device_ts.begin(),
      device_ts.end(),
      ts.begin()
  );
  
  return copy_back_ts_event;
}

thrust::device_event Double(thrust::host_vector<float>& floats, thrust::host_vector<int>& ints) {
  // 非同期実行の完了をまとめる
  return thrust::when_all(Double(floats), Double(ints));
}

intとfloatの配列のそれぞれに対して、デバイスへのコピーが完了したら、2乗する計算を行い、終わればホストへコピーする処理を順番に実行ポリシーに指定しています。
最後に、両方の配列がホストへコピーされたことを確認するためのeventを作成しています。

あとは、main関数で各配列を作って、Double関数に渡して、結果が正しいか確認するだけです。

main.cpp

#include "double.h"

#include <cassert>
#include <cmath>
#include <iostream>

int main() {
  auto floats = thrust::host_vector<float>();
  auto ints   = thrust::host_vector<int>();

  std::size_t size = 10000;

  floats.reserve(size);
  ints.reserve(size);

  for (std::size_t i = 0; i < size; ++i) {
    floats.push_back(i);
    ints.push_back(i);
  }

  auto event = Double(floats, ints);

  event.wait();

  for (std::size_t i = 0; i < size; ++i) {
    assert(std::abs(floats[i] - i * i) < 1e-5);
    assert(ints[i] == static_cast<int>(i * i));
  }
  std::cout << "Success!\n";

  return 0;
}

本当は浮動小数点数の値の比較をもう少し丁寧にするべきですが、本題と関係ないので良しとします。

利点

一般に非同期実行というと身構える方が多いと思われますが、上の解説を見ると意外と簡単にできると思われたのではないでしょうか。
また、実行ポリシーで実行順序を指定できるのも魅力的です。

非同期処理の利点というよりThrustを使う利点になりますが、実はrocThrustという、CUDAをHIPとROCmに置き換えたThrustをAMDが開発しています。
現時点では、Thrust 1.17.2まで対応していますので非同期実行もできます。
詳しくは参考文献4をご覧ください。

欠点

残念ながら欠点があります。
1つ目の欠点はCUDAがデバイスの時かつデバイス側でしか使えないことです。
ThrustはoneTBBをデバイスとして使うこともできますが、非同期処理はoneTBBでは実装されていません。
シリアルやOpenMPも同様です。
「eventとfuture」節で触れましたが、これがhost_eventがない理由だと考えられます。
しかし、この欠点はさほど問題にならないでしょう。

2つ目の欠点はasync::copyが転送するデータの型に対してtrivially relocatableであることを要求していることです。
trivially relocatableな型とは、簡単に言うとコピーとムーブ(それぞれ構築と代入を含む)、破棄がdefaultまたはdeleteで、仮想メンバ関数と仮想継承を持たず、全てのメンバ変数と基底クラスもtrivially relocatableな型です。
これはmemcpyでホストとデバイス間でデータ転送できるようにするために導入されました。
Thrustでは、THRUST_PROCLAIM_TRIVIALLY_RELOCATABLEマクロ関数を使って、与えた型がtrivially relocatableであると宣言する必要があります。
制約が強いもののパフォーマンスは向上し、私が行った実験ではそうでないものに比べて実行時間が約30%減っています。
詳しくは参考文献5をご覧ください。

実は...

Thrustのリポジトリをご覧になるとわかるのですが、なんとアーカイブにされてます。
どうやら半年前からCCCLという、CUBとlibcudacxxも一緒にしたリポジトリに移ったようです。
このことに気づいたのが一昨日だったので、昨日CCCLでもビルドできるか試していたのですが、上述のやり方ではエラーが出てしまいました。
どうにか解決したかったのですが、時間もなかったのでとりあえずThrustに対して書きました。
もしご存じの方がいらっしゃったら、教えていただけると幸いです。

まとめ

データがtrivially relocatableな型なら、お手軽に非同期実行できます。

ソースコード

GitHubに置いています。