Хэш-таблицы на GPU (CUDA и GLSL)

Одной из широко распространенных структур для поиска данных являются хэш-таблицы. И можно довольно легко написать ее реализацию для однонитевого приложения. Но в случае кода, выполняемого на GPU, когда у вас многие тысячи нитей, это становится гораздо сложнее. При этом нити могут одновременно модифицировать эту таблицу.

Здесь я хочу рассмотреть реализацию подобной таблицы, с которой можно работать на GPU, т.е. из многих тысяч и десятков тысяч нитей. При этом мы будем поддерживать и добавление и удаление элементов. И, что также очень важно, это будет lock-free таблица.

Естественно, чтобы это можно было достаточно просто реализовать, необходимо ввести некоторые ограничения. Мы будем строить хэш-таблицу, переводящую int в int. При таком выборе гарантируется оптимальное выравнивание данных в памяти GPU и то, что запись ключей и значений будет происходить атомарно, т.е. ни одна нить не прочтет не до конца записанное значение.

Второе ограничение - мы изначально выделить таблицу фиксированного размера и не будет проводить ее переаллокацию и/или перестройку (вообще выделять память на стороне GPU довольно сложно даже в CUDA). Кроме того, мы будем использовать таблицы с размером равным степени двух, что позволит вместо отнюдь не самой дешевой функции mod использовать побитовое AND.

Мы рассмотрим реализацию хэш-таблицы и на CUDA и в шейдере на GLSL. Каждая запись в таблице будет представлена в виде следующей структуры:

struct KeyValue
{
    uint32_t key;
    uint32_t value;
};

Кроме того, мы введем следующие константы:

const uint32_t HashTableCapacity = 128 * 1024 * 1024;
const uint32_t Empty = 0xffffffff;

Здесь HashTableCapacity обозначает размер хэш-таблицы (и является степень двух). Второе значение - Empty - обозначает отсутствующее значнение, причем оно может использоваться как в качестве ключа, так и в качестве значения. Вначале мы проинициализируем всю таблицу (и key и value) значением Empty.

В качестве хэш-функции в этом примере мы будем использовать достаточно простую и быструю Murmur3. Для разрешения коллизий в нашей таблице мы будем использовать linear probing, т.е. в поисках значения будем перебирать записи, идущие после той, на которую указывает хэш.

// 32 bit Murmur3 hash
__device__ uint32_t hash(uint32_t k)
{
    k ^= k >> 16;
    k *= 0x85ebca6b;
    k ^= k >> 13;
    k *= 0xc2b2ae35;
    k ^= k >> 16;
    return k & (HashTableCapacity-1);
}

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

__device__ bool hashtableFind ( KeyValue * hashtable, uint32_t key )
{
    uint32_t slot = hash ( key );
 
    while ( true )
    {
        if ( hashtable[slot].key == key )    // found key
            return true;
 
        if (hashtable[slot].key == Empty)
            return false;
        
        slot = (slot + 1) & (HashTableCapacity - 1);
    }
}

Вставка значений в таблицу является гораздо более сложной операцией, поскольку не исключено, что сразу несколько нитей попытаются обратиться к одной и той же записи в таблице (будет хэш-коллизия при записи). Поэтому для чтения ключа из таблицы мы воспользуемся atomic-compare-and--swap(code>atomicCAS). Данная функция читает значение по переданному адресу и сравнивает его со своим вторым аргументом. В случае совпадения происходит запись на этом место значения третьего аргумента.

Эта функция всегда возвращает прочитанное значение (не зависимо от того, была ли запись или нет). Таким образом atomicCAS(&hashtable[slot].key, Empty, key) в случае если в hastable[slot].key было значение Empty (т.е. этот слот свободен), производит запись туда нашего ключа. В противном случае возможны два варианта - на этом месте находился другой ключ или же на этом месте находится наш ключ key. Что именно там было мы можем узнать по возвращенному значению - если Empty или же наш ключ, то мы производим запись значения в этом место. Иначе мы просто переходим к следующему элементу таблицы.

__device__ void hashtableInsert ( KeyValue * hashtable, uint32_t key, uint32_t value )
{
    uint32_t slot = hash(key);
 
    while ( true )
    {
        uint32_t prev = atomicCAS (&hashtable [slot].key, Empty, key );
        
        if ( prev == Empty || prev == key )
        {
            hashtable [slot].value = value;
            return;
        }
 
        slot = (slot + 1) & (kHashTableCapacity-1);
    }
}

Последняя операция, которую нам осталось рассмотреть, это удаления значения (по ключу) из таблицы. Здесь мы не будем использовать атомарные операции, но при удалении найденной записи мы просто в ее поле value запишем значение Empty. Так как ключ на самом деле остается в таблице, то при длительном применении возможно накопление мусора и таблицу лучше почистить. Мы будем считать, что это будет делать отдельное ядро CUDA.

__device__ void hashtableDelete ( KeyValue * hashtable, uint32_t key )
{
    uint32_t slot = hash(key);
 
    while ( true )
    {
        if (hashtable [slot].key == key )    // found key
        {
            hashtable[slot].value = Empty;
            return;
        }
        if ( hashtable[slot].key == Empty )    // key not found, return
            return;
        
        slot = (slot + 1) & (kHashTableCapacity - 1);
    }
}

Этот код можно легко переписать на GLSL. При этом в качестве таблицы мы будем использовать SSBO, а вместо функции atomicCAS мы используем функцию atomicCompSwap.

#version 450

//
// Simple GLSL hashtable implementation
// uint -> uint
// No GPU rehashing supported

struct KeyValue     // entry in table
{
    uint key;
    uint value;
};

layout ( std430, binding = 0 ) buffer Hashtable
{
    KeyValue    hashtable []; 
};

const uint HashTableCapacity = 128 * 1024 * 1024;
const uint Empty = 0xffffffff;

    // 32 bit Murmur3 hash
uint hash ( uint k )
{
    k ^= k >> 16;
    k *= 0x85ebca6b;
    k ^= k >> 13;
    k *= 0xc2b2ae35;
    k ^= k >> 16;

    return k & ( HashTableCapacity - 1 );
}
 
bool hashtableFind ( uint key )
{
    uint slot = hash ( key );
 
    while ( true )
    {
        if ( hashtable [slot].key == key )    // found key
            return true;
 
        if ( hashtable [slot].key == Empty )
            return false;
        
        slot = (slot + 1) & (HashTableCapacity - 1);
    }
}
 
void hashtableInsert ( uint key, uint value )
{
    uint slot = hash ( key );
 
    while ( true )
    {
        uint prev = atomicCompSwap ( hashtable [slot].key, Empty, key );
        
        if ( prev == Empty || prev == key )
        {
            hashtable [slot].value = value;
            return;
        }
 
        slot = (slot + 1) & (HashTableCapacity-1);
    }
}
 
void hashtableDelete ( uint key )
{
    uint slot = hash ( key );
 
    while ( true )
    {
        if (hashtable [slot].key == key )    // found key
        {
            hashtable [slot].value = Empty;
            return;
        }
        if ( hashtable [slot].key == Empty )    // key not found, return
            return;
        
        slot = (slot + 1) & (HashTableCapacity - 1);
    }
}

Ссылки по теме:

A Simple GPU Hash Table

WarpCore: A Library for fast Hash Tables on GPUs

SimpleComputeShaderHashTable