radix sort, single threaded & multithreaded (spmd style)

radix_sort.c

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

//
// radix sort
//

static void radix_sort(uint64_t* array, size_t count, uint64_t* scratch)
{
    uint64_t* src = array;
    uint64_t* dst = scratch;

    // loop going from LSB to MSB
    // each iteration will guarantee array will be sorted on all lower "d" bits
    for (size_t d=0; d<64; d+=8)
    {
        size_t counts[256] = { 0 };

        // histogram for digits in all input values
        for (size_t i=0; i<count; i++)
        {
            uint8_t digit = (uint8_t)(src[i] >> d);
            counts[digit]++;
        }

        // prefix sum, histogram -> offsets
        size_t offset = 0;
        for (size_t i=0; i<256; i++)
        {
            size_t c = counts[i];
            counts[i] = offset;
            offset += c;
        }

        // move input values to sorted location
        for (size_t i=0; i<count; i++)
        {
            uint8_t digit = (uint8_t)(src[i] >> d);
            dst[counts[digit]++] = src[i];
        }

        uint64_t* swap = src;
        src = dst;
        dst = swap;
    }
}

//
// test program
//

static uint64_t random64()
{
    static uint64_t x = 0, w = 0;
    x = x*x + (w += 0xb5ad4eceda1ce2a9);
    return x = (x>>32) | (x<<32);
}

int main()
{
    const size_t count = 100*1000*1000;

    uint64_t* array = malloc(count * sizeof(*array));
    uint64_t* scratch = malloc(count * sizeof(*array));

    for (size_t i=0; i<count; i++)
    {
        array[i] = random64();
    }

    // do the sort
    radix_sort(array, count, scratch);

    // check if array is sorted
    for (size_t i=1; i<count; i++)
    {
        if (array[i-1] > array[i])
        {
            printf("ERROR! array is not sorted!\n");
            return 1;
        }
    }

    printf("OK!\n");
}

radix_sort_spmd.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

//
// Threading API
//

// size_t thread_index;               // thread index = [0, thread_count)
// size_t thread_count;               // thread count
// thread_sync();                     // synchronizes all threads
// thread_broacast(&var);             // broadcasts contents of variable (via memcpy) from thread index 0 to other threads
// thread_range(count, &start, &end); // split [0, count) range into "thread_count" sub-ranges - returns [start, end) for current "thread_index"

#if defined(_WIN32)
#  include <windows.h>
#  define THREAD_TLS                   __declspec(thread)
#  define THREAD_BARRIER_TYPE          SYNCHRONIZATION_BARRIER
#  define THREAD_BARRIER_WAIT(barrier) EnterSynchronizationBarrier(barrier, 0)
#else
#  include <unistd.h>
#  include <pthread.h>
#  define THREAD_TLS                   __thread
#  define THREAD_BARRIER_TYPE          pthread_barrier_t
#  define THREAD_BARRIER_WAIT(barrier) pthread_barrier_wait(barrier)
#endif

static THREAD_BARRIER_TYPE thread_barrier;

static THREAD_TLS size_t thread_index;
static            size_t thread_count;

static void thread_sync()
{
    THREAD_BARRIER_WAIT(&thread_barrier);
}

static void thread_broadcast_impl(void* ptr, size_t size)
{
    static void* broadcast_ptr;

    if (thread_index == 0)
    {
        broadcast_ptr = ptr;
    }
    thread_sync();

    if (thread_index != 0)
    {
        memcpy(ptr, broadcast_ptr, size);
    }
    thread_sync();
}
#define thread_broadcast(var) thread_broadcast_impl(&var, sizeof(var))

static void thread_range(size_t count, size_t* start, size_t* end)
{
    size_t index0 = thread_index + 0;
    size_t index1 = thread_index + 1;

    size_t chunk = count / thread_count;
    size_t extra = count % thread_count;

    *start = index0 * chunk + (index0 < extra ? index0 : extra);
    *end   = index1 * chunk + (index1 < extra ? index1 : extra);

    // alternative that can be used when count is uint32_t
    // with this need to be careful to not overflow multiplication
    // *start = index0 * count / thread_count;
    // *end   = index1 * count / thread_count;
}

//
// radix sort
//

static void radix_sort(uint64_t* array, size_t count, uint64_t* scratch)
{
    uint64_t* src = array;
    uint64_t* dst = scratch;

    size_t* thread_counts;
    if (thread_index == 0)
    {
        // this could be static allocation of [max_thread_count][256] array       
        thread_counts = malloc(thread_count * 256 * sizeof(*thread_counts));
    }
    thread_broadcast(thread_counts);

    size_t* counts = thread_counts + 256 * thread_index;

    size_t range_start, range_end;
    thread_range(count, &range_start, &range_end);

    // loop going from LSB to MSB
    // each iteration will guarantee array will be sorted on all lower "d" bits
    for (size_t d=0; d<64; d+=8)
    {
        // histogram for digits in values from input range
        memset(counts, 0, 256 * sizeof(*counts));
        for (size_t i=range_start; i<range_end; i++)
        {
            uint8_t digit = (uint8_t)(src[i] >> d);
            counts[digit]++;
        }
        thread_sync();

        // per-thread histogram -> per-thread offsets
        if (thread_index == 0)
        {
            size_t offset = 0;
            for (size_t i=0; i<256; i++)
            {
                for (size_t t=0; t<thread_count; t++)
                {
                    size_t c = thread_counts[256*t + i];
                    thread_counts[256*t + i] = offset;
                    offset += c;
                }
            }
        }
        thread_sync();

        // move values from input range to sorted location
        for (size_t i=range_start; i<range_end; i++)
        {
            uint8_t digit = (uint8_t)(src[i] >> d);
            dst[counts[digit]++] = src[i];
        }
        thread_sync();

        uint64_t* swap = src;
        src = dst;
        dst = swap;
    }

    if (thread_index == 0)
    {
        free(thread_counts);
    }
}

//
// test program
//

static uint64_t random64()
{
    static uint64_t x = 0, w = 0;
    x = x*x + (w += 0xb5ad4eceda1ce2a9);
    return x = (x>>32) | (x<<32);
}

static void thread_main()
{
    const size_t count = 100*1000*1000;

    uint64_t* array;
    uint64_t* scratch;

    if (thread_index == 0)
    {
        array = malloc(count * sizeof(*array));
        scratch = malloc(count * sizeof(*array));

        for (size_t i=0; i<count; i++)
        {
            array[i] = random64();
        }
    }
    thread_broadcast(array);
    thread_broadcast(scratch);

    // do the sort
    radix_sort(array, count, scratch);

    // check if array is sorted
    if (thread_index == 0)
    {
        for (size_t i=1; i<count; i++)
        {
            if (array[i-1] > array[i])
            {
                printf("ERROR! array is not sorted!\n");
                exit(1);
            }
        }

        printf("OK!\n");
    }
}

#if defined(_WIN32)

static DWORD WINAPI thread_entry(LPVOID arg)
{
    thread_index = (size_t)arg;
    thread_main();
    return 0;
}

int main()
{
    enum { max_thread_count = 64 };

    SYSTEM_INFO sysinfo;
    GetSystemInfo(&sysinfo);

    thread_count = sysinfo.dwNumberOfProcessors;
    if (thread_count > max_thread_count)
    {
        thread_count = max_thread_count;
    }
    InitializeSynchronizationBarrier(&thread_barrier, (LONG)thread_count, -1);

    HANDLE threads[max_thread_count];
    for (size_t i=0; i<thread_count; i++)
    {
        threads[i] = CreateThread(NULL, 0, &thread_entry, (LPVOID)i, 0, NULL);
    }
    WaitForMultipleObjects((DWORD)thread_count, threads, TRUE, INFINITE);
}

#else

static void* thread_entry(void* arg)
{
    thread_index = (size_t)arg;
    thread_main();
    return 0;
}

int main()
{
    enum { max_thread_count = 64 };

    thread_count = sysconf(_SC_NPROCESSORS_ONLN);
    if (thread_count > max_thread_count)
    {
        thread_count = max_thread_count;
    }
    pthread_barrier_init(&thread_barrier, NULL, thread_count);

    pthread_t threads[max_thread_count];
    for (size_t i=0; i<thread_count; i++)
    {
        pthread_create(&threads[i], NULL, &thread_entry, (void*)i);
    }
    for (size_t i=0; i<thread_count; i++)
    {
        pthread_join(threads[i], NULL);
    }
}

#endif