Работа с библиотекой thrust в CUDA

Одной из самых красивых и гибких библиотек для CUDA является библиотека thrust. Основным отличием данной библиотеки от других распространенных библиотек для CUDA является то, что thrust - это библиотека, основанная на использовании шаблонов (template) языка С++.

Все классы и функции в этой библиотеки - шаблонные, все, что вам нужно для работы с этой библиотекой - это подключить соответствующие заголовочные файлы - в thrust нет никаких стандартных библиотечных файлов (.lib, .a, .dll, .so).

Ряд понятий и подходов thrust заимствовала из STL. thrust предоставляет в ваше распоряжение набор различных параллельных примитивов, таких как различные преобразования, сортировка, операции reduce и scan. Применяя thrust, многие действия могут быть записаны просто и понятно с использованием минимального объема кода. Все последние версии CUDA включают в себя thrust, так что для работы с thrust никаких дополнительных установок вам не понадобится.

Все, что вводит thrust, помещается в пространстве имен thrust. В частности это удобно тем, что позволяет избежать путаницы с рядом операций, имеющих те же самые имена (и параметры), что и операции, вводимые в STL (например, sort). Обычно, если thrust вводит операцию, уже существующую в STL, то для нее используется то же имя и аналогичный синтаксис.

Вектора

В отличии от STL, предоставляющего в распоряжение программиста большое число различных контейнеров, thrust предоставляет всего два контейнера - thrust::host_vector и thrust::device_vector.

Оба этих класса соответствуют одной и той же структуре данных - вектору. Только один из этих классов хранит содержимое вектора в памяти CPU, а другой - в памяти GPU. Выбор в качестве типа контейнера только вектора - вполне естественный шаг. Ведь с точки зрения быстродействия, вектор - это идеальная структура данных, которая очень хорошо ложится на архитектуру GPU (впрочем и на архитектуру CPU также).

Как и std::vector, вектора в thrust - это контейнеры, параметризованные типом хранящихся в них значений. Соответственно контейнеры в thrust могут хранить значения самых разных типов, сами управляют памятью и могут динамически изменять свой размер. Работа с векторами в thrust очень напоминает работу с векторами в STL. Например, вы также можете использовать оператор [] для доступа к произвольному элементу вектора. При этом совсем не важно где на самом деле хранится содержимое вектора - в памяти CPU или же в памяти GPU. Однако важно иметь в виду, что за подобным обращением к содержимому вектора, хранящегося в памяти GPU, на самом деле стоит скрытый вызов cudaMemcpy, поэтому это довольно дорогая операция.

Ниже приводится простой пример работы с векторами в thrust, обратите внимание на использование оператора = для копирования данных между векторами.

#include <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/reduce.h>

using namespace std;

int sequence ()
{
    static int i = 0;

    return ++i;
}

int main ()
{
    thrust::host_vector<int> hv ( 1000 );
    thrust::device_vector<int> dv ( 1000 );

    thrust::generate ( hv.begin (), hv.end (), sequence );
    dv = hv;

    int sum = thrust::reduce ( dv.begin (), dv.end () );

    cout << "Sum is: " << sum << endl;

    return 0;
}

Обратите внимание на то, что thrust сам управляет выделением памяти для векторов - вы не можете построить вектор, используя уже выделенную для этого память. Однако можно получить указатель на блок памяти, в котором хранятся значения вектора при помощи вызова thrust::raw_pointer_cast.

thrust::device_vector<int> v (1000);
int * rawPtr = thrust::raw_pointer_cast(v.data));
myKernel<<<threads, blocks>>> ( rawPtr, v.size () );

Итераторы

Так же как и в STL, thrust использует итераторы для задания места (или диапазона) в векторе (или массиве). Все контейнеры предоставляют стандартные методы, такие как begin () и end () для получения итераторов. Так же как и в STL вы можете прибавить к итератору целое число, например v.begin()+5 даст итератор на 5-й элемента вектора v.

Однако в thrust итератор - это не просто указатель. Тип итератора содержит в себе информацию о том, на какую память, CPU или GPU, указывает данный итератор. За счет этого уже на этапе компиляции происходит выбор подходящей реализации - выбирается реализация для CPU или же реализация для GPU.

Так же как и в STL, вместо итератора можно передать обычный указатель. Но в этом случае, поскольку на этапе компиляции указатель не содержит в себе информации о том, на какую именно память он ссылается, то будет считаться, что указатель указывает на память CPU и будет выбрана реализация именно для CPU.

Поэтому, если у вас есть указатель на память GPU, например полученный через cudaMalloc, то для того, чтобы можно было передать этот указатель в алгоритмы thrust, его нужно завернуть в специальный тип - thrust::device_ptr.

int * rawPtr;

cudaMalloc ( (void**)&rawPtr, N * sizeof ( int ) );

thrust::device_ptr ptr ( rawPtr );

// some operations on ptr

cudaFree ( rawPtr );

Чтобы получить из thrust::device_ptr оригинальный указатель следует использовать уже знакомую функцию thrust::raw_pointer_cast.

Также для выделения и освобождения памяти GPU вместо cudaMalloc и cudaFree можно использовать и функции из thrust - thrust::device_malloc и thrust::device_free:

thrust::device_ptr<int> p  = thrust::device_malloc<int> ( N );

// perform some operations

thrust::device_free ( p );

Простейшие операции над векторами

Простейшими операциями над векторами (которые, как и все операции в thrust, применимы как thrust::device_vector, так и к thrust::host_vector) являются fill, sequence и copy.

Простейшим из них является fill, который заполняет вектор, заданный двумя итераторами (на самом деле вместе итераторов можно использовать device_ptr) - началом и концом - необходимым значением.

thrust::device_vector<int> d ( N );

thrust::fill ( d.begin (), d.end (), 77 );

Функция thrust::sequence также заполняет вектор значениями, но уже не одним значением, а значениями, полученными из арифметической последовательности (заданной первым элементом и шагом).

thrust::device_vector<int> d ( N );

        // fill with 7, 5, 3, 1, -1, ...
thrust::sequence ( d.begin (), d.end (), 7, -2 );

Для последних двух параметров этой функции есть значения по умолчанию - если их не задавать, то для заполнения будут использованы значения 0, 1, 2, 3, ...

Функция thrust::copy просто копирует участок вектора на заданное место, для задания участка и места используются итераторы или device_ptr.

thrust::copy ( src.begin (), src.end (), dst.begin () );

Обратите внимание, что thrust::copy может копировать данные между памятью CPU и GPU. Также обратите внимание, что функции, рассматриваемые далее, например, thrust::transform, не выделяют памяти - перед их вызовом у вас в соответствующем векторе должно быть достаточно места.

Алгоритмы

Библиотека thrust предоставляет ряд параллельных реализация для многих алгоритмов, заметно облегчающих жизнь программистов на CUDA. Обратите внимание, что за исключением thrust::copy, во всех функциях thrust, в которые передаются итераторы, все передаваемые итераторы должны указывать память CPU, либо все передаваемые итераторы должны указывать в память GPU. Невыполнение этого требования приведет к ошибке компиляции.

Простейшим алгоритмом, поддерживаемым thrust, является thrust::transform, выполняющий некоторую операцию над входными значениями и записывающим результат выполнения операции по заданному итератору. Сама операция задается при помощи функтора, при этом файл thrust/functional.h содержит ряд уже готовых к использованию функторов.

Простейшим случаем является использование унарного функтора - есть один входной вектор, есть функтор, принимающий один аргумент, и есть выходной вектор.

thrust::transform ( a.begin (), a.end(), output.end (), operation );

Рассмотрим как можно ввести свою унарную операцию на примере вычисления модуля числа - пусть мы хотим вычислить значение модуля для всех элементов заданного вектора. Для этого необходимо ввести соответствующий функтор - класс или структуру для которой определен оператор (), вычисляющий модуль числа.

template<typename T>
struct abs_functor
{
    __device__ __host__ T operator ()( const T& x ) const
    {
        return x >= 0 ? x : -x;
    }
};

После того, как мы определили функтор, можно передать его в thrust::transform, например, следующим образом:

thrust::transform ( v.begin (), v.end (), v.begin (), abs_functor () );

Обратите внимание, что изменение вектора происходит in-place - изменяется сам исходный массив. Также обратите внимание на использование круглых скобок после имени функтора - это вызов конструктора, так как в вызов нужно передать созданный экземпляр объекта класса abs_functor.

В файле thrust/functional.h определены следующие унарные и бинарные функторы - plus, minus, multiplies, divides, modulus, maximum, minimum, bit_or, bit_and, bit_xor, bit_and, equal_to, not_equal_to, greater, less, greater_equal, less_equal, logical_and, logical_or и logical_not.

В случае бинарной операции на вход поступают два входных вектора, один выходной и сама операция.

thrust::transform ( a.begin (), a.end (), b.begin (), output.begin (), thrust::multiplies<float> () );

Как и для унарного функтора, можно вводить и свои бинарные функторы и thrust будет с ними работать. В качестве примера рассмотрим часто встречающуюся в пакетах линейной алгебры операцию SAXPY. Ее можно реализовать при помощи использования нескольких вызовов thrust::transform, однако будет гораздо эффективнее сделать это всего за один вызов. Для этого необходимо реализовать функтор, получающий на вход два аргумента и возвращающий результат их обработки как переопределенный оператор ().

struct saxpy
{
    const float a;

    saxpy ( float _a ) : a (_a) {}

    float   __device__ __host__ operator () ( const float& x, const float& y ) const 
    {
        return a*x + y;
    }
};

Одной из распространенных операций над массивами данных является редукция - мы объединяем все элементы массива при помощи заданной бинарной операции. Так для нахождения суммы всех элементов массива можно просто выполнить редукцию, используя в качестве бинарной операции сложение двух чисел.

Операция редукции поддерживается thrust через функцию thrust::reduce. Эта функция получает на вход два итератора, задающих диапазон в массиве, начальное значение и функтор, задающий используемую бинарную операцию. Так для того, чтобы найти сумму всех элементов заданного вектора мы в качестве начального значения используем 0, а в качестве операции - thrust::plus.

thrust::device_vector<int> a ( N );

// fill a with data

// find the sum
int sum = thrust::reduce( a.begin (), a.end (), 0, thrust::plus<int> () );

Частным случаем редукции является нахождения числа элементов, равных заданному (или удовлетворяющих заданному условию). Для есть готовые функции thrust::count и thrust::count_if, определенные в файле count.h.

int c = thrust::count ( a.begin (), a.end (), 7 );

Кроме этого thrust предоставляет еще один вариант редукции - thrust::transform_reduce. Эта функция выполняет операцию редукции над результатом применения преобразования элементов массива. Ее можно заменить двумя отдельными вызовами, но использование такого комбинированного вызова оказывается быстрее. Следующий пример показывает использование этой функции для нахождения суммы квадратов элементов заданного вектора.

template<typename T>
struct square
{
    __device__ __host__ T operator ()( const T& x ) const
    {
        return x*x;
    }
};

После определения такого функтора для нахождения суммы квадратов элементов вектора можно выполнить следующий вызов:

float s = thrust::transform_reduce ( v.begin (), v.end (), square<float>(), 0, thrust::plus<float>() );

Еще одной распространенной операцией, поддерживаемой thrust, является нахождение префиксной суммы, причем поддерживается два варианта нахождения префиксной суммы - включающий (x0, x0+x1, x0+x1+x0,...) и исключающий (0, x0, x0+x1, x0+x1+x2,...).

thrust::inclusive_scan ( x.begin (), x.end (), x.begin () );    // do in-place
thrust::exclusive_scan ( a.begin (), a.end (), b.begin () );

Также этим функциям можно передать на вход начальное значение (по умолчанию - 0) и бинарную операцию (по умолчанию - сложение). По аналогии с редукцией есть операции thrust::transform_inclusive_scan и thrust::transform_exclusive_scan.

Одной из очень важных и часто встречающихся операций, выполняемых над массивами данных, является сортировка. thrust предоставляет несколько различных функций для сортировки. Простейшими из них являются thrust::sort и thrust::stable_sort, определяемые в файле sort.h. Эти функции получают на вход два итератора, задающие диапазон данных для сортировки. Также можно передать функтор - операцию сравнения элементов вектора.

thrust::device_vector<float> v ( N );

// fill with data

thrust::sort ( v.begin (), v.end () );

Отличие обычной сортировки и stable_sort заключается в том, что stable_sort не меняет порядок элементов, которые не нужно менять местами с точки зрения функции сравнения. Обычная сортировка этого не гарантирует.

Иногда возникает ситуация, когда у нас есть отдельный массив данных и отдельный массив ключей, определяющих упорядочение элементов. Для этого случая есть функции thrust::sort_by_key и thrust::stable_sort_by_key.

thrust::sort_by_key ( keys.begin (), keys.end (), data.begin () );

Итераторы

Кроме стандартных итераторов, возвращаемыми такими методами как begin и end (а также объектов типа device_ptr, которые можно использовать вместо итератора), в thrust есть несколько специальных итераторов, которые могут оказаться весьма полезными.

Одним из простейших подобных итераторов является thrust::constant_iterator, всегда возвращающий одно и то же значение. Пусть например, нужно умножить все элементы заданного вектора на 2. Можно для этого построить специальный унарный функтор, умножающий на 2, и вызвать thrust::transform. Но можно использовать стандартный бинарный функтор thrust::multiplies, а в качестве итератора для второго массива передать thrust::constant_iterator, проинициализированный значением 2.

thrust::constant_iterator<int> two ( 2 );

thrust::transform ( a.begin (), a.end (), two, output.begin (), thrust::multiplies<int> () );

Похожим итератором является thrust::counting_iterator. Он возвращает арифметическую последовательность чисел, начиная с заданного числа. Следующий пример демонстрирует как найти индексы всех ненулевых элементов заданного вектора при помощи этого итератора. Для копирования элементов, удовлетворяющих заданному условию, используется функция thrust::copy_if.

thrust::device_vector<int> stencil(N);
    // fill with some values
. . .   

 thrust::device_vector<int> indices(N);
   
// compute indices of nonzero elements
 typedef thrust::device_vector<int>::iterator IndexIterator;
   
// use make_counting_iterator to define the sequence [0, 8)
 IndexIterator indices_end = thrust::copy_if( thrust::make_counting_iterator(0),
                                              thrust::make_counting_iterator(N),
                                              stencil.begin(),
                                              indices.begin(),
                                              thrust::identity<int>() );

Подобно тому, как thrust::transform можно объединить с рядом других методов, его также можно объединить и с итератором, получая при этом итератор, который будет возвращать значение, вычисленное для очередного элемента вектора. Для получения такого итератора из обычного итератора и функтора служит функция thrust::make_transform_iterator.

thrust::reduce ( thrust::make_transform_iterator ( a.begin (), thrust::negate<int> () ),
                 thrust::make_transform_iterator ( a.end    (), thrust::negate<int> () ) );

Одним из самых мощных итераторов в thrust является thrust::zip_iterator, объединяющий вместе несколько итераторов в один, возвращающий значения типа thrust::tuple. Так если у нас есть два вектора (возможно содержащий значения совершенно разных типов), то мы легко можем получить итератор, который будет пробегать сразу по обоим этим массивам следующим образом:

auto first = thrust::make_zip_iterator ( thrust::make_tuple ( a.begin (), b.begin () ) );
auto last =  thrust::make_zip_iterator ( thrust::make_tuple ( a.end   (), b.end   () ) );

Теперь рассмотрим боле сложный пример - у нас есть два массива векторов. Однако для обеспечения "правильного" паттерна доступа к памяти вместо массива трехмерных векторов хранятся три массива из float. Соответственно у нас всего 6 векторов из float и мы хотим построить массив из попарных скалярных произведений этих векторов. Для этого можно воспользоваться следующим кодом, использующим thrust::zip_iterator-ы для того, чтобы получать доступ сразу ко всем компонентам текущего вектора.

typedef thrust::tuple<float,float,float> Float3;

// This functor implements the dot product between 3d vectors
struct DotProduct : public thrust::binary_function<Float3,Float3,float>
{
    __host__ __device__
    float operator()(const Float3& a, const Float3& b) const
    {
        return thrust::get<0>(a) * thrust::get<0>(b) + // x components
                   thrust::get<1>(a) * thrust::get<1>(b) + // y components
                   thrust::get<2>(a) * thrust::get<2>(b); // z components
    }
};

. . .

thrust::transform(  thrust::make_zip_iterator(make_tuple(A0.begin(), A1.begin(), A2.begin())),
                    thrust::make_zip_iterator(make_tuple(A0.end(), A1.end(), A2.end())),
                    thrust::make_zip_iterator(make_tuple(B0.begin(), B1.begin(), B2.begin())),
                    result.begin(),
                    DotProduct() );

Использование лямбд в thrust

К сожалению стандартные ламбды из языка С++ на момент написания этой статьи не работают с thrust. Однако есть другой очень красивый способ вводить ламбды в thrust. Он заключается в использовании специальных переменных с именами _1, _2, _3 и так далее из пространства имен thrust::placeholders. Они позволяют прямо на месте определять необходимые функторы и использовать их в вызовах функций из thrust. Ниже приводится простой пример увеличивающий каждый элемент из массива на единицу при помощи функции thrust::for_each (которую мы рассмотрим позже).

#include <thrust/functional.h>

using namespace thrust::placeholders;

thrust::for_each(x.begin(), x.end(), _1++);

Аналогично можно ввести ранее рассмотренную операцию saxpy:

// saxpy: x = a * x + y
 int a = 42;
 thrust::transform(x.begin(), x.end(), y.begin(), x.begin(), a * _1 + _2);

Update: На самом деле сейчас в CUDA уже поддерживаются лямбды, однако, как и любая функция в CUDA, она выполняется на GPU или CPU. Обычно это определяет то, в выполняемом где кода она создается. Но есть возможность явно это задать, как показано ниже.

thrust::partition ( dv2.begin (), dv2.end (), [] __device__ ( int x ) { return x % 2 == 0; } );

На данный момент возможность явного задания для на должна выполняться является экспериментальной и требует задания для nvcc следующей опции --expt-extended-lambda. При ее задании для любой лямбды можно явно задать область выполнения и любая лямбда с правильной областью выполнения может быть использована в thrust в качестве функтора.

Дополнительные функции

Кроме перечисленных выше в библиотеке thrust есть много других функций. Ниже мы приведем краткое описание некоторых из них.

Для вектора можно проверить выполнение заданного логической функции (предиката) при помощи функция thrust::all_of, thrust::any_of и thrust::none_of, каждый из них получает на вход три аргумента - два итератора, задающий диапазон для проверки и функтор, который будет вычисляться для каждого элемента массива. Так в следующем примере проверяется, что все элементы меньше нуля.

bool flag = thrust::all_of ( a.begin (), a.end (), _1 < 0 );

Функции thrust::copy_if и thrust::remove_if позволяют скопировать все элементы удовлетворяющие заданному функтору или же удалить все элементы, удовлетворяющие заданному функтору. Так следующий пример скопирует только четные элементы.

#include <thrust/copy.h>
  ...
struct is_even
  {
    __host__ __device__
    bool operator()(const int x)
    {
      return (x % 2) == 0;
    }
};
  ...

  thrust::copy_if  (a.begin(), a.end (), result, is_even());

Функция thrust::remove_if фактически переупорядочивает элементы, перемещая элементы, не удовлетворяющие заданному условию в конец. Возвращается итератор на конец набор элементов, для которых условие выполнено.

int *new_end = thrust::remove_if(A, A + N, is_even());

При помощи функции thrust::generate можно заполнить вектор значениями, возвращаемыми заданной функцией или функтором. Так следующий пример заполнят вектор случайными целыми числами.

thrust::generate ( v.begin (), v.end (), rand );

Функция thrust::for_each выполняет заданный функтор для каждого элемента массива или вектора. Так следующий пример увеличивает значение каждого элемента на единицу. На вход он получает два итератора, задающий диапазон, и функтор.

thrust::for_each(x.begin(), x.end(), _1++);

При помощи thrust::partition можно переупорядочить элементы массива так, чтобы элементы, удовлетворяющие заданному функтору шли первыми. Так в следующем примере осуществляется переупорядочение так, чтобы четные элементы шли впереди нечетных. Возвращается итератор, такой что на всем диапазоне от первого переданного итератор до возвращенного итератора заданный функтор возвращает истинное значение.

#include <thrust/partition.h>
  ...
  struct is_even
  {
    __host__ __device__
    bool operator()(const int &x)
    {
      return (x % 2) == 0;
    }
  };
  ...
  int A[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
  const int N = sizeof(A)/sizeof(int);
  thrust::partition(A, A + N, is_even());

С помощью функций из библиотеки thrust можно легко реализовывать многие более сложные операции. Так ниже приводится пример построения гистограммы при помощи библиотеки thrust.

template <typename Vector1, typename Vector2>
void dense_histogram(const Vector1& input, Vector2& histogram)
{
    typedef typename Vector1::value_type ValueType; // input value type
    typedef typename Vector2::value_type IndexType; // histogram index type

    // copy input data (could be skipped if input is allowed to be modified)
    thrust::device_vector<ValueType> data(input);

    // sort data to bring equal elements together
    thrust::sort(data.begin(), data.end());

    // number of histogram bins is equal to the maximum value plus one
    IndexType num_bins = data.back() + 1;

    // resize histogram storage
    histogram.resize(num_bins);

    // find the end of each bin of values
    thrust::counting_iterator<IndexType> search_begin(0);
    thrust::upper_bound(data.begin(), data.end(),
              search_begin, search_begin + num_bins,
             histogram.begin());

    // compute the histogram by taking differences of the cumulative histogram
    thrust::adjacent_difference(histogram.begin(), histogram.end(), histogram.begin());
}