GPGPU : CUDA C 用の Mutex, Stack の実装について

CUDA に対応した GPU から使える Mutex や Stack の実装のコードが面白かったので, コードを以下に載せます.


以下の Mutex のコードは CUDA の Compute capability 1.2 の GPU で動作します.

struct Mutex
{
    int* value;

    void InitializeFromHost( void )
    {
        int state = 0;

        HANDLE_ERROR( cudaMalloc( (void**)& value,
            sizeof( int ) ) );
        HANDLE_ERROR( cudaMemcpy( value,
                                  &state,
                                  sizeof( int ),
                                  cudaMemcpyHostToDevice ) );
    }

    void FinalizeFromHost( void )
    {
        cudaFree( value );
    }

    __device__ void
    Lock( void )
    {
        while ( atomicCAS( value, 0, 1 ) != 0 );
    }

    __device__ void
    Unlock( void )
    {
        atomicExch( value, 1 );
    }
}


以下のスタックのコードは CUDA の Compute capability 2.0 の GPU で動作します.

// 各スレッドごとに使用可能なスタックです.

// Stack in local memory. Managed independently for each thread.
template<class T, int N>
class local_stack
{
private:
    T buf[N];
    int tos;

public:
    __device__ local_stack() :
    tos(-1)
    {
    }

    __device__ T const & top() const
    {
        return buf[tos];
    }

    __device__ T & top()
    {
        return buf[tos];
    }

    __device__ void push(T const & v)
    {
        buf[++tos] = v;
    }

    __device__ T pop()
    {
        return buf[tos--];
    }

    __device__ bool full()
    {
        return tos == (N - 1);
    }

    __device__ bool empty()
    {
        return tos == -1;
    }
};

// こちらは複数のスレッドで利用でき,
// グローバルなメモリ上に置かれるものです.

// Stacks in global memory.
// Same function as local_stack, but accessible from the host.
// Interleaved between threads by blocks of THREADS elements.
// Independent stack for each thread,
// no sharing of data between threads.
template<class T, int N, int THREADS>
class global_stack
{
private:
    T * buf;
    int free_index;

public:
    // buf should point to an allocated global buffer of
    // size N * THREADS * sizeof(T)
    __device__ global_stack(T * buf, int thread_id) :
    buf(buf), free_index(thread_id)
    {
    }

    __device__ void push(T const & v)
    {
        buf[free_index] = v;
        free_index += THREADS;
    }
    
    __device__ T pop()
    {
        free_index -= THREADS;
        return buf[free_index];
    }

    __device__ bool full()
    {
        return free_index >= N * THREADS;
    }

    __device__ bool empty()
    {
        return free_index < THREADS;
    }

    __device__ int size()
    {
        return free_index / THREADS;
    }
};