Битоническая сортировка при помощи вычислительных шейдеров

Я уже рассматривал на сайте битоническую сортировку при помощи фрагментных шейдеров. Она работает, но на данный момент это далеко не самый оптимальный вариант выполнения сортировки на GPU. Поэтому здесь будет рассмотрен модифицированный вариант битонической сортировки и его реализация при помощи вычислительных шейдеров. Одним из основных недостатков битонической сортировки является большое число обращений к DRAM GPU, поэтому мы рассмотрим как можно использовать быструю разделяемую (shared) память для увеличения быстродействия.

В классическом варианте битонической сортировки строится так называемая сортировочная сеть, состоящая из полуочистителей (half-cleaner). Каждый полуочиститель получает на вход два элемента и упорядочивает их по возрастанию или по убыванию. В классической схеме направления сортировки все время чередуются, что не очень удобно.

Последовательность a₀,...,a_n-1 называется битонической если она или состоит из двух монотонных частей (одна часть возрастает, другая убывает) или получается циклическим вращением такой последовательности.

Если последовательность битоническая то при помощи битонического слияния (bitonic merge) ее можно легко отсортировать. Битоническое слияние это последовательное применение наборов полуочистителей. На самом деле произвольная последовательность может быть приведена к битонической следующим путем.

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

В результате мы получим набор битонических последовательностей из 8 элементов. Их мы отсортируем опять чередуя направление. Тогда мы получим набор битонических последовательностей из 16 элементов и т.д. до тех пор пока мы не придем к полностью отсортированной последовательности.

Но на самом деле существует вариант битонической сортировки в котором вообще нет чередования направлений и все полуочистители всегда сортируют в одном и том же направлении что приводит к более чистому и легко читаемому коду. На рис. 1 приводится схема подобной битонической сортировки для 16 элементов. Каждый вертикальный отрезок на рисунке соответствует одному полуочистителю, выполняющему операцию compare-and-swap (CAS). Обратите внимание, что на самом деле нитей нужно вдвое меньше, чем у нас входных элементов, т.е. 8.

Рис. 1. Схема битонической сортировки 16 элементов без чередования направления

Как видно из рисунка вся сортировка состоит из ряда шагов, на каждом из которых выполняется операция compare-and-swap над всеми элементами входного массива. Разница между шагами заключается в том, как строятся пары сравниваемых элементов. На рисунке видно 10 таких шагов, причем эти шаги делятся на две группы, обозначенные на рисунке зеленым и желтым. И также видно что каждый шаг характеризуется своей высотой - расстоянием начиная с которого мы начинаем сравнивать элементы.

На зеленом шаге мы работаем с группами элементов и для каждого следующего шага высота (т.е. количество элементов в группе h) удваивается. Внутри каждой группы мы сравниваем элементы, симметрично расположенные относительно середины группы. Таким образом если у нас есть группа из h элементов a_s,a_s+1,...,a_s+h-1, то мы будем сравнивать между собой пары (a_s+j,s_s+h-1-j) для j от 0 до h/2. Данный шаг мы будем называть flip, ниже приводится его реализация на С++.

void    flip ( int n, int h )
{
    for ( int t = 0; t < n/2; t++ )
    {
        int q       = ((2 * t) / h) * h;
        int t2      = t % (h/2);

        compareAndSwap ( q + t2, q + h - t2 - 1 );
    }
}

За каждый зеленым блоком (flip) обязательно следует желтый блок. Каждый такой блок также работает с группами по h элементов. И с каждым новым шагом h уменьшается вдвое. Если группа состоит из элементов a_s,a_s+1,...,a_s+h-1, то мы будем сравнивать между собой пары (a_s+j, a_s+j+h/2). Желтый шаг мы будем называть disperse. Ниже приводится реализация желтого шага на С++.

void    disperse ( int n, int h )
{
    for ( int t = 0; t < n/2; t++ )
    {
        uint    q       = ((2 * t) / h) * h;
        uint    t2      = t % (h/2);

        compareAndSwap ( q + t2, q + t2 + h/2  );
    }
}

Тогда весь процесс сортировки массива из n элементов (считая n степенью двух) можно записать на С++ следующим образом:

for ( int h = 2; h <= n; h *= 2 ) 
{
    flip ( n, h );

    for ( int hh = h / 2; hh > 1; hh /= 2 ) 
        disperse ( n, hh );
}

Теперь давайте перепишем ранее приведенный код в виде вычислительного шейдера на GLSL. Для начала давайте рассмотрим реализацию битонической сортировки при помощи разделяемой памяти. Поскольку объем разделяемой памяти, который может быть выделен одному блоку ограничен, то мы начнем со случая, когда число сортируемый элементов n равно 1024. В этом случае мы гарантированно можем разместить весь массив в разделяемой памяти GPU и нам понадобится всего один блок из 512 нитей (на 1024 сортируемый элементов).

Основным свойством разделяемой памяти является то,что она выделяется на каждый блок и является очень быстрой. Но при этом блок не может обратиться к разделяемой памяти другого блока. На n элементов входного массива мы запустим один блок из n/2 нитей, каждая нить при этом отвечает за два элемента массива. Ниже приводится соответствующий шейдер, обратите внимание на обязательное использование барьерной синхронизации - вызовы функции barrier. Она нужно потому, что нити блока могут находится в разных местах кода, какие-то могут убежать вперед, а какие-то наоборот отстать. Гарантируется, что функцию barrier все нити блока пройдут одновременно.

-- compute
#version 450 core

    // Sort 16 sortable elements
#define NUM_ELEMENTS 16

layout ( local_size_x = NUM_ELEMENTS / 2 ) in;

layout ( binding = 0 ) buffer SortData 
{
    uint value [];
};

    // Use local memory to minimize memory access cost.
shared uint local [NUM_ELEMENTS];

void    localCompareAndSwap ( uvec2 idx )
{
    if ( local [idx.x] > local [idx.y] ) 
    {
        uint tmp = local [idx.x];

        local [idx.x] = local [idx.y];
        local [idx.y] = tmp;
    }
}

void    flip ( uint h )
{
    uint    t       = gl_LocalInvocationID.x;
    uint    q       = ((2 * t) / h) * h;
    uint    t2      = t % (h / 2);
    uvec2   indices = q + uvec2 ( t2, h - t2 - 1 );

    localCompareAndSwap ( indices );
}

void    disperse ( uint h )
{
    uint    t       = gl_LocalInvocationID.x;
    uint    q       = ((2 * t) / h) * h;
    uint    t2      = t % (h / 2);
    uvec2   indices = q + uvec2 ( t2, t2 + (h / 2) );

    localCompareAndSwap ( indices );
}

void    main ()
{
    uint    t = gl_LocalInvocationID.x;
    int     n = NUM_ELEMENTS;

        // each thread loads two elements to shared memory
    local [t*2    ] = value [t*2    ];
    local [t*2 + 1] = value [t*2 + 1];
 
    for ( uint h = 2; h <= n; h *= 2 )
    {
        barrier ();
        flip    ( h );

        for ( uint hh = h / 2; hh > 1; hh /= 2 ) 
        {
            barrier  ();
            disperse ( hh );
        }
    }

    barrier ();

        // write data back
    value [t*2    ] = local [t*2    ];
    value [t*2 + 1] = local [t*2 + 1];
}

Приведенный выше шейдер легко выполнит битоническую сортировку 1024 элементов с использованием разделяемой памяти. Использование быстрой разделяемой памяти важно, поскольку у нас имеет место большое число обращений к памяти. Однако данный шейдер может отсортировать только массив небольшого размера, который целиком поместится в разделяемую память. Можно пойти другим путем - все операции проводить с SSBO вообще не используя разделяемую память. Для этого нужно использовать приведенные ниже версии функций bigFlip и bigDisperse, которые работают напрямую с SSBO.

    // Performs full-height flip (h height) over globally available indices.
void bigFlip ( in uint h )
{
    uint tPrime = gl_GlobalInvocationID.x;
    uint halfH  = h >> 1;
    uint q      = ((2 * tPrime) / h) * h;
    uint t2     = tPrime % halfH;
    uint x      = q + t2;
    uint y      = q + h - t2 - 1; 

    globalCompareAndSwap ( ivec2 ( x, y ) );
}

    // full-height disperse (h height) over globally available indices.
void bigDisperse ( in uint h )
{
    uint tPrime = gl_GlobalInvocationID.x;
    uint halfH  = h >> 1;
    uint q      = ((2 * tPrime) / h) * h;
    uint t2     = tPrime % halfH;
    uint x      = q + t2;
    uint y      = q + t2 + halfH;

    globalCompareAndSwap ( ivec2 ( x, y ) );
}

Однако понятно, что быстродействие у такого подхода будет не очень высоким из-за большого числа обращений к довольно медленной DRAM GPU. Однако можно заметно повысить эффективность за счет грамотного использования разделяемой памяти. Для этого давайте представим что мы по-прежнему сортируем 16 элементов, но теперь у нас есть два блока по 4 нити и у каждого блока свой участок разделяемой памяти не доступный другому блоку.

Рис. 2. Битоническая сортировка 16 элементов при помощи двух блоков по 4 нити в каждом.

Как видно из рисунка у нас для 16 элементов есть только один шаг, где нужно нитям одного блока обращаться к памяти другого блока - это bigFlip. Все остальные шаги могут быть реализованы через разделяемую память. В общем случае у нас кроме bigFlip появится еще и bigDisperse.

Но при этом все что можно сделать используя разделяемую память будет делаться с использованием разделяемой памяти. В результате у нас появится четыре операции - localSort, disperse, bigFlip и bigDisperse. Из них первые две работают с разделяемой памятью, вторые две - с DRAM GPU.

Операция localSort это на самом деле объединенные вместе несколько операций flip disperse. Можно эти четыре операции разбить на четыре шейдера, здесь мы будем использовать один большой шейдер, реализующий сразу все эти четыре операции. Какую из них нужно выполнить в данный момент будет задаваться через uniform-переменную algorithm. Подобный подход позволит уменьшить число дорогостоящих переключений между шейдерами.

Далее приводится соответствующий вычислительный шейдер.

-- compute
#version 450 core

#define LocalBitonicMergeSort   0
#define LocalDisperse           1
#define BigFlip                 2
#define BigDisperse             3

layout ( local_size_x = 512 ) in;

layout ( binding = 0 ) buffer SortData 
{
    uint value [];
};

uniform int step;
uniform int algorithm;

shared uint local_value [gl_WorkGroupSize.x * 2];

bool is_smaller ( in const uint left, in const uint right )
{
    return left > right;
}

#define COMPARE is_smaller

void globalCompareAndSwap ( ivec2 idx )
{
    if ( COMPARE ( value [idx.x], value [idx.y] ) )
    {
        uint tmp = value [idx.x];

        value [idx.x] = value [idx.y];
        value [idx.y] = tmp;
    }
}

void localCompareAndSwap ( ivec2 idx )
{
    if ( COMPARE ( local_value [idx.x], local_value [idx.y] ) )
    {
        uint tmp = local_value [idx.x];

        local_value [idx.x] = local_value [idx.y];
        local_value [idx.y] = tmp;
    }
}

    // Performs full-height flip (h height) over globally available indices.
void bigFlip ( in uint h )
{
    uint tPrime = gl_GlobalInvocationID.x;
    uint halfH  = h >> 1;
    uint q      = ((2 * tPrime) / h) * h;
    uint t2     = tPrime % halfH;
    uint x      = q + t2;
    uint y      = q + h - t2 - 1; 

    globalCompareAndSwap ( ivec2 ( x, y ) );
}

    // full-height disperse (h height) over globally available indices.
void bigDisperse ( in uint h )
{
    uint tPrime = gl_GlobalInvocationID.x;
    uint halfH  = h >> 1;
    uint q      = ((2 * tPrime) / h) * h;
    uint t2     = tPrime % halfH;
    uint x      = q + t2;
    uint y      = q + t2 + halfH;

    globalCompareAndSwap ( ivec2 ( x, y ) );
}

    // Performs full-height flip (h height) over locally available indices.
void localFlip ( in uint h )
{
    uint    t = gl_LocalInvocationID.x;
    uint    halfH  = h >> 1;
    uint    t2     = t % halfH;
    uint    q      = h * ( ( 2 * t ) / h );

    barrier ();

    localCompareAndSwap ( ivec2 ( q + t2, q + h - 1 - t2 ) );
}

// Performs progressively diminishing disperse operations (starting with height h)
// on locally available indices: e.g. h==8 -> 8 : 4 : 2.
// One disperse operation for every time we can divide h by 2.
void localDisperse ( in uint h )
{
    uint    t = gl_LocalInvocationID.x;

    for ( ; h > 1 ; h /= 2 ) 
    {
        barrier ();

        uint    halfH  = h >> 1;
        uint    t2     = t % halfH;
        uint    q      = h * ( ( 2 * t ) / h );

        localCompareAndSwap ( ivec2 ( q + t2, q + halfH + t2 ) );
    }
}

    // binary merge sort for local elements
void localSort ( in uint h )
{
    for ( uint hh = 2; hh <= h; hh <<= 1 ) 
    {
        localFlip     ( hh     );
        localDisperse ( hh / 2 );
    }
}

void    main ()
{
    uint t = gl_LocalInvocationID.x;
 
        // Calculate global offset for local workgroup
    uint offset = gl_WorkGroupSize.x * 2 * gl_WorkGroupID.x;

        // This shader can be called in four different modes:
        // 
        //  1. local flip+disperse (up to n == local_size_x * 2) 
        //  2. big flip
        //  3. big disperse
        //  4. local disperse 

    if ( algorithm <= LocalDisperse )
    {
            // for local store data in shared memory
        local_value [t*2    ] = value [offset + t*2    ];
        local_value [t*2 + 1] = value [offset + t*2 + 1];
    }

    switch ( algorithm )
    {
        case LocalBitonicMergeSort:
            localSort ( step );
            break;

        case LocalDisperse:
            localDisperse ( step );
            break;

        case BigFlip:
            bigFlip ( step );
            break;

        case BigDisperse:
            bigDisperse ( step );
            break;
    }

        // write shared data back to buffer
    if ( algorithm <= LocalDisperse )
    {
        barrier ();

            // push to global memory
        value [offset + t*2    ] = local_value [t*2    ];
        value [offset + t*2 + 1] = local_value [t*2 + 1];
    }
}

Ниже приводится код на С++, реализующий битоническую сортировку при помощи описанного шейдера. Обратите внимание что в используемом классе содержится сразу две метода sortSlow и sort. Первый реализует сортировку без использования разделяемой памяти, а второй - с использованием.

#define NUM_ELEMENTS    (1024*1024)     // 1M elements

#define LocalBitonicMergeSort   0
#define LocalDisperse           1
#define BigFlip                 2
#define BigDisperse             3

class   SortWindow : public GlutWindow
{
    Program                 program;
    VertexArray             vao;
    VertexBuffer        buf;
    std::vector<uint32_t>       array;
    const int               blockSize = 512;        // # of threads in a block

public:
    SortWindow () : GlutWindow ( 100, 100, 512, 512, "Compute Sort" )
    {
            // initialize array with random elements to sort
        array.resize ( NUM_ELEMENTS );

        for ( int i = 0; i < NUM_ELEMENTS; i++ )
            array [i] = rand ();

            // load sorting compute shader
        if ( !program.loadProgram ( "global-sort.glsl" ) )
            exit ( "Error building program: %s\n", program.getLog ().c_str () );    
        
            // prepare buffer with data to be sorted
        program.bind  ();
        vao.create    ();
        vao.bind      ();
        buf.create    ();
        buf.bindBase  ( GL_SHADER_STORAGE_BUFFER, 0 );
        buf.setData   ( NUM_ELEMENTS * sizeof ( uint32_t ), array.data (), GL_DYNAMIC_DRAW );

            // perform the sort
        sort ( NUM_ELEMENTS );

        program.unbind    ();

            // get sorted array
        std::vector<uint32_t>   buffer ( NUM_ELEMENTS );

        glGetNamedBufferSubData ( buf.getId (), 0, NUM_ELEMENTS * sizeof ( uint32_t ), buffer.data () );

            // sort original data for comparison
        std::sort ( array.begin (), array.end () );

            // print the diffs
        printf ( "\nDiffs:\n" );

        for ( int i = 0; i < NUM_ELEMENTS; i++ )
            if ( buffer [i] != array [i] )
                printf ( "Diff aat index %d - %d vs %d\n", i, array [i], buffer [i] );
    }

    void redisplay () override
    {
        glClear ( GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT );
    }

    void    dispatch ( int n, int h, int algo )
    {
        const int blockCount = n / ( blockSize * 2 );

        program.setUniformInt ( "step",      h    );
        program.setUniformInt ( "algorithm", algo );

        glDispatchCompute ( blockCount, 1, 1 );
        glMemoryBarrier   ( GL_SHADER_STORAGE_BARRIER_BIT );
    }

    void    sort ( int n )
    {
        int h = blockSize * 2;

        assert ( h <= n );
        assert ( h % 2 == 0 );

        dispatch ( n, h, LocalBitonicMergeSort );

        for ( h *= 2; h <= n; h *= 2 )
        {
            dispatch ( n, h, BigFlip );

            for ( int hh = h / 2; hh > 1; hh /= 2 )
            {
                if ( hh <= blockSize * 2 ) 
                {
                    dispatch ( n, hh, LocalDisperse );
                    break;
                } 
                else
                    dispatch ( n, hh, BigDisperse );
            }
        }
    }

    void    sortSlow ( int n )
    {
        for ( int h = 2; h <= n; h *= 2 ) 
        {
            dispatch        ( n, h, BigFlip );
            glMemoryBarrier ( GL_SHADER_STORAGE_BARRIER_BIT );

            for ( int hh = h / 2; hh > 1; hh /= 2 ) 
            {
                dispatch        ( n, hh, BigDisperse );
                glMemoryBarrier ( GL_SHADER_STORAGE_BARRIER_BIT );
            }   
        }
    }
};

int main ( int argc, char * argv [] )
{
    GlutWindow::init ( argc, argv );
    
    SortWindow  win;
    
    return 0;
}

Исходный код к этой статье можно скачать по этой ссылке. В основу данной статьи легла следующая статья на английском языке.