Cooperative groups в CUDA

В CUDA 9 появился принципиально новый функционал по организации взаимодействия нитей между собой - cooperative groups. Этот функционал предоставляет масштабируемое взаимодействие нитей сразу на нескольких уровнях.

До появления cooperative groups основным способом взаимодействия нитей друг с другом была барьерная синхронизация через вызов __syncthreads. Однако этот способ взаимодействия работает только на уровне отдельных блоков. С появлением cooperative groups стало возможным организовывать взаимодействие сразу на всех перечисленных ниже уровнях:

• нитей внутри блока, идущих по одной ветви условия
• варпа или результата деления варпа на части с размером, равным степени двух
• блока
• всей сетки (grid)
• нескольких сеток, выполняющихся на различных GPU (multi-grid)

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

#include    <cooperative_groups.h>

using namespace cooperative_groups;

Фундаментальным типом является thread_group и ряд производных от него классов. Все эти классы представляют собой группу взаимодействующих между собой нитей. Каждый из этих классов содержит следующие методы:

void        sync ();        // выполнить синхронизацию всех нитей в группе
unsigned    size ();        // число нитей в группе
unsigned    thread_rank (); // порядковый номер нити в группе, начиная с 0
bool        is_valid ();    // нарушает ли группу какие-то ограничения API

Кроме того, класс thread_block, соответствующий всем нитям блока, поддерживает два следующих метода:

dim3 group_index  ();   // 3-мерный индекс блока внутри сетки
dim3 thread_index ();   // 3-мерный индекс нити внутри блока

Рассмотрим следующий пример. В нем мы создадим три группы нитей - одну, соответствующую всему блоку (block), одну соответствующую одному варпу (tile32) и еще одну, соответствующую четырем подряд идущим нитям (tile4).

thread_block block  = this_thread_block ();
thread_group tile32 = tiled_partition   ( block, 32 );
thread_group tile4  = tiled_partition   ( tile32, 4 );

В этом фрагменте кода функция this_thread_block возвращает объект, соответствующий группе нитей, образующих текущий блок. Назначение функции tiled_partition - разбить переданную группу нитей на подгруппы с заданным размером. Пока есть ограничения на размер частей - он должен быть степенью двух и не превосходить 32 (т.е. он может быть только 1, 2, 4, 8, 16 или 32 ).

Приводимая ниже функция reduce работает не на уровне блока, а уровне отдельной группы нитей. Вся синхронизация нитей внутри группы осуществляется именно для заданной группы нитей через метод sync.

__device__ int reduce ( thread_group g, int * x, int val )
{
    int lane = g.thread_rank ();    // номер нити внутри группы

    for ( int i = g.size () / 2; i > 0; i /= 2 )
    {
        x [lane] = val;
        g.sync ();
        val += x [lane + i];
        g.sync ();
    }

    return val;
}

Ранее мы рассматривали вариант функции tiled_partition, который на вход получает два параметра - саму группу, которую надо разделить, и размер получаемых подгрупп. Есть и другой вариант этой функции - он является шаблонным и размер подгруппы получает как параметр шаблона. Он возвращает специальный шаблонный класс, параметризованный размером группы. В результате размер возвращаемой подгруппы является константой времени компиляции, что позволяет компилятору проводить различные оптимизации.

thread_block_tile<32> tile32 = tiled_partition<32> ( this_thread_block () );
thread_block_tile<4>  tile4  = tiled_partition<4>  ( this_thread_block () );

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

template <typename group_t>
__device__ int reduce_sum ( group_t g, int * temp, int val )
{
    int lane = g.thread_rank ();    // номер нити в группе

    #pragma unroll
    for ( int i = g.size () / 2; i > 0; i /= 2 )
    {
        x [lane] = val;
        g.sync ();
        val += x [lane + i];
        g.sync ();
    }

    return val; 
}

Шаблонный класс thread_block_tile также поддерживает ряд дополнительный методов, соответствующих функциям CUDA по взаимодействию нитей на уровне варпа - shfl(), shfl_down (), shfl_up() shfl_xor (), any (), all (), ballot (), match_any () и match_all ().

template <unsigned size>
__device__ int tiled_reduce ( thread_block_tile<size> g, int val )
{
    #pragma unroll
    for ( int i = g.size () / 2; i > 0; i /= 2 )
        val += g.shfl_down ( val, i );

    return val;
}

Есть еще один вид группы нитей - coalesced. Поскольку CUDA выполняет нити группами по 32 подряд идущих нити (варп), то при выполнении условных конструкций иногда возникает ситуация, когда часть нитей варпа выполняет определенный фрагмент кода, а другая - нет. Так вот - coalesced group - это группа нитей, возникающих в этом случае.

Тем самым если у нас есть оператор if-else, то в каждой из его ветвей мы получаем свою coalesced-группу, соответствующую всем нитям, для которых условие выполнено. И в пределах этой группы можно проводить синхронизацию, через метод sync, а также вызывать другие методы для варпа (такие как shfl). Такой группе нитей соответствует класс coalesced_group.

auto block = this_thread_block ();
if ( block.thread_rank () % 2 )
{
    coalesced_group acative = coalesced_threads ();
    . . .
    val += active.shfl ( val, 0 );
    active.sync ();
}

В качестве примера использования coalesced-групп мы рассмотрим оптимизация вызова функции atomicInc. Вместо того, чтобы вызывать ее для всех нитей варпа, для которых выполнено заданное условие инкремента, мы посчитаем число таких нитей в варпе и выделим одну нить, которая прибавит сразу это число. В результате мы заметно уменьшаем число вызовов атомарной функции.

__device__ int atimicAggInc ( int * ptr )
{
    coalesced_group g = coalesced_threads ();   // все нити варпа, для которых нужно выполнить inc
    int prev;

    if ( g.thread_rank () == 0 )                // выбираем одну нить
        prev = atomicInc ( ptr, g.size () );    // делаем инкркмент от всех нитей группы

                                                // возвращаем правильное значение
    prev = g.thread_rank () + g.shfl ( prev, 0 );

    return prev;
}

Можно получить группу нитей, соответствующую всей сетке, но только при условии, что сетка была запущена правильным образом, гарантирующим, что все нити сетки будут одновременно выполняться. Для этого нужно для запуска сетки использовать не стандартную конструкцию, а следующую функцию, введенную в CUDA C API:

cudaError_t cudaLaunchCooperativeKernel ( const void * kernel, dim3 gridDim, dim3 blockDim,
        void ** args, size_t sharedMem, cudaStream_t stream );

Здесь параметр args - это указатель на массив указателей на отдельные аргументы вызова ядра. Запускаемое ядро не может использовать динамический параллелизм и сам GPU должен поддерживать свойство cudaDevAttrCooperativeLaunch. Проверить эту поддержку можно при помощи следующего вызова:

cudaDeviceGetAttribute ( &attr, cudaDevAttrCooperativeLaunch, device );

Для получения числа блоков, которые можно одновременно можно запустить на одном мультипроцессоре, следует использовать следующую функцию:

cudaOccupancyMaxActiveBlocksPerMultiprocessor ( &numBlocks, kernel, maxThreads, 0 );

Ниже приводится пример синхронизации всех нитей сетки внутри ядра.

__global__ void kernel ()
{
    grid_group grid = this_grid ();
    . . .
    grid.sync ();
    . . .
}

На данный момент кооперативный запуск поддерживается только для GPU Pascal и Volta и только на Linux. Ниже приводится пример кода, демонстрирующий кооперативный запуск и синхронизацию нитей всей сетки.

__global__ void kernel ( int x )
{
    grid_group grid = this_grid ();

    printf ( "Before sync %d:%d, x = %d\n", grid.thread_rank (), grid.size (), x );

    grid.sync ();

    printf ( "After sync %d:%d, x = %d\n", grid.thread_rank (), grid.size (), x );
}

int main ()
{
    int attr, numBlocks;
    int numThreads = 256;
    int device = 0;

    cudaDeviceProp  devProp;

    cudaGetDeviceProperties ( &devProp, device );

    printf ( "SM count %d\n", devProp.multiProcessorCount );

    cudaDeviceGetAttribute ( &attr, cudaDevAttrCooperativeLaunch, device );

    if ( !attr )
    {
        printf ( "Cooperative launch not supported %d\n", attr );
        exit   ( 1 );
    }

    cudaOccupancyMaxActiveBlocksPerMultiprocessor ( &numBlocks, kernel, numThreads, 0 );

    printf ( "max blocks per SM %d\n", numBlocks );

    int x = 17;
    void * args [] = { &x };

    cudaLaunchCooperativeKernel ( (void *)kernel, dim3 ( numBlocks * devProp.multiProcessorCount ), dim3 ( numThreads ),  args, 0, 0 );

    cudaDeviceSynchronize ();

    return 0;
}

Обратите внимание, что для сборки данного примера необходимо использовать опции "-rdc=true -lcudadevrt".

Также существует и multi_grid_group - группа нитей, соответствующий нескольким сеткам, запущенным на разных GPU в одной и той же системе.

multi_grid_group multi_grid = this_multi_grid ();

. . .

multi_grid.sync ();

Для запуска сеток сразу на нескольких GPU служит функция cudaLaunchCooperativeKernelMultiDevice.

struct cudaLaunchParams params [NUM_DEVICES];

for ( int i = 0; i < NUM_DEVICES; i++ )
{
    params [i].func      = (void *) kernel;
    params [i].gridDim   = grid;
    params [i].blockDim  = block;
    params [i].sharedMem = 0;
    params [i].stream    = streams [i];        // нельзя использовать поток 0
    params [i].args      = args;
}

cudaLaunchCooperativeKernelMultiDevice ( params, NUM_DEVICES );