[기본] Nividia의 Thrust란?

PyTorch의 C++쪽 코드(caffe2, aten, c10, ..)를 보게 되면 thrust라는 라이브러리를 include하여 사용하는 것을 볼 수 있습니다. 해당 라이브러리는 Nividia에서 개발한 Parallel Algorithms Library입니다. STL과 유사한 high-level interface로 구현되어 있어서 개발자들이 STL을 사용하듯이 병렬 프로그래밍을 할 수 있게 구현되어 있습니다.

 

그리고 위에서 언급했듯이 PyTorch의 Core 부분에서도 아래의 기능들을 include하여 사용하고 있습니다.

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/functional.h>
#include <thrust/tuple.h>
#include <thrust/pair.h>
#include <thrust/iterator/reverse_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/execution_policy.h>
#include <thrust/sort.h>
#include <thrust/unique.h>
#include <thrust/device_ptr.h>
#include <thrust/extrema.h>
#include <thrust/find.h>
#include <thrust/inner_product.h>
#include <thrust/sequence.h>
#include <thrust/adjacent_difference.h>
#include <thrust/scan.h>
#include <thrust/scatter.h>
#include <thrust/for_each.h>
#include <thrust/fill.h>
#include <thrust/gather.h>
#include <thrust/generate.h>
#include <thrust/transform.h>
// ...

---

그러면 아래의 대표적인 사용 예를 보면서 thrust의 쓰임새에 대해 이야기 해보겠습니다.

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/generate.h>
#include <thrust/sort.h>
#include <thrust/copy.h>
#include <thrust/random.h>

int main() {
  // 32M개의 연속적인 임의의 숫자를 생성합니다.
  thrust::default_random_engine rng(1337);
  thrust::uniform_int_distribution<int> dist;
  thrust::host_vector<int> h_vec(32 << 20);
  thrust::generate(h_vec.begin(), h_vec.end(), [&] { return dist(rng); });

  // 생성한 데이터를 device 메모리로 전달합니다.
  thrust::device_vector<int> d_vec = h_vec;

  // 데이터를 정렬합니다. (정렬은 Merge Sort를 사용합니다.)
  thrust::sort(d_vec.begin(), d_vec.end());

  // device 메모리의 정렬된 데이터를 host로 전달합니다.
  thrust::copy(d_vec.begin(), d_vec.end(), h_vec.begin());
}

위와 같이 기본적인 알고리즘(sort, find, transform, ...)들과 자료 구조(tuple, vector, pair, ...)들은 thrust에 구현이 되어 있습니다. 일반적으로 우리가 작성하는 알고리즘에 비해 최적화도 많이 되어있는 상태이기 때문에, thrust에 존재한다고 하면 가져다 사용하는걸 추천합니다. 

 

그렇다면 thrust를 어떻게 사용할 수 있을까요? CMake 빌드 시스템을 사용하는 중이라면 아래와 같이 CMakeLists.txt에 추가함으로써 굉장히 간단하게 추가할 수 있습니다. (단, CMake 최소 버전이 3.15 이상이야 합니다.)

find_package(Thrust REQUIRED CONFIG)
thrust_create_target(Thrust)
target_link_libraries(MyProgram Thrust)

그런데 만약 CMake가 Trust의 위치를 찾지 못할 경우 아래처럼 지정해줌으로써 해결 할 수 있습니다.

$ cmake . -DThrust_DIR=<thrust git repo root>/thrust/cmake/

 

 

 

 

참고

https://github.com/NVIDIA/thrust

GitHub - NVIDIA/thrust: The C++ parallel algorithms library.