warp_atomics.h

Go to the documentation of this file.

/*
 * CUGAR : Cuda Graphics Accelerator
 *
 * Copyright (c) 2010-2018, NVIDIA Corporation
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *   * Neither the name of NVIDIA Corporation nor the
 *     names of its contributors may be used to endorse or promote products
 *     derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#pragma once

#include <cugar/basic/types.h>
#include <cugar/bits/popcount.h>
#include <cub/cub.cuh>

namespace cugar {
namespace cuda {




__device__ __forceinline__
unsigned int warp_increment(unsigned int *ptr)
{
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 700
    const unsigned int lane_id = threadIdx.x & 31;

    int pred;
    int mask = __match_all_sync(__activemask(), (unsigned long long)ptr, &pred);
    int leader = __ffs(mask) - 1;    // select a leader

    int res = 0;
    if (lane_id == leader)                  // leader does the update
        res = atomicAdd(ptr, __popc(mask));

    res = __shfl_sync(mask, res, leader);    // get leader’s old value
    return res + __popc(mask & ((1 << lane_id) - 1)); //compute old value
#else
    const uint32 warp_tid   = threadIdx.x & 31;
    const uint32 warp_mask  = __ballot_sync( __activemask(), true );
    const uint32 warp_count = __popc( warp_mask );
    const uint32 warp_scan  = __popc( warp_mask << (warpSize - warp_tid) );

    const uint32 first_tid = ffs(warp_mask) - 1;

    uint32 broadcast_offset = 0;

    // acquire an offset for this warp
    if (warp_tid == first_tid)
        broadcast_offset = atomicAdd( ptr, warp_count );

    // obtain the offset from the first participating thread
    const uint32 offset = cub::ShuffleIndex<32>(broadcast_offset, first_tid, __activemask());

    // compute the per-thread offset
    return offset + warp_scan;
#endif
}

struct warp_static_atomic
{
    __device__ __forceinline__
    warp_static_atomic(uint32* pool)
        : m_dest(pool) {}

    template <uint32 N>
    __device__ __forceinline__
    void add(bool p)
    {
        warp_static_atomic::static_add<N>( p, m_dest );
    }

    template <uint32 N>
    __device__ __forceinline__
    void sub(bool p)
    {
        warp_static_atomic::static_sub<N>( p, m_dest );
    }

    template <uint32 N>
    __device__ __forceinline__
    void add(bool p, uint32* result)
    {
        warp_static_atomic::static_add<N>( p, m_dest, result );
    }

    template <uint32 N>
    __device__ __forceinline__
    void sub(bool p, uint32* result)
    {
        warp_static_atomic::static_sub<N>( p, m_dest, result );
    }

    // --- stateless methods --- //

    template <uint32 N>
    __device__ __forceinline__
    static void static_add(bool p, uint32* dest)
    {
        const uint32 warp_tid   = threadIdx.x & 31;
        const uint32 warp_mask  = __ballot_sync( __activemask(), p );
        const uint32 warp_count = __popc( warp_mask );
        const uint32 warp_scan  = __popc( warp_mask << (warpSize - warp_tid) );

        // perform the atomic
        if (warp_scan == 0 && p)
            atomicAdd( dest, warp_count * N );
    }

    template <uint32 N>
    __device__ __forceinline__
    static void static_sub(bool p, uint32* dest)
    {
        const uint32 warp_tid   = threadIdx.x & 31;
        const uint32 warp_mask  = __ballot_sync( __activemask(), p );
        const uint32 warp_count = __popc( warp_mask );
        const uint32 warp_scan  = __popc( warp_mask << (warpSize - warp_tid) );

        // perform the atomic
        if (warp_scan == 0 && p)
            atomicSub( dest, warp_count * N );
    }

    template <uint32 N>
    __device__ __forceinline__
    static void static_add(bool p, uint32* dest, uint32* result)
    {
        const uint32 warp_tid   = threadIdx.x & 31;
        const uint32 warp_mask  = __ballot_sync( __activemask(), p );
        const uint32 warp_count = __popc( warp_mask );
        const uint32 warp_scan  = __popc( warp_mask << (warpSize - warp_tid) );

        const uint32 first_tid = ffs(warp_mask) - 1;

        uint32 broadcast_offset;

        // acquire an offset for this warp
        if (warp_tid == first_tid)
            broadcast_offset = atomicAdd( dest, warp_count * N );

        // obtain the offset from the first participating thread
        const uint32 offset = cub::ShuffleIndex<32>(broadcast_offset, first_tid, __activemask());

        // compute the per-thread offset
        *result = offset + warp_scan * N;
    }

    template <uint32 N>
    __device__ __forceinline__
    static void static_sub(bool p, uint32* dest, uint32* result)
    {
        const uint32 warp_tid   = threadIdx.x & 31;
        const uint32 warp_mask  = __ballot_sync( __activemask(), p );
        const uint32 warp_count = __popc( warp_mask );
        const uint32 warp_scan  = __popc( warp_mask << (warpSize - warp_tid) );

        const uint32 first_tid = ffs(warp_mask) - 1;

        uint32 broadcast_offset;

        // acquire an offset for this warp
        if (warp_tid == first_tid)
            broadcast_offset = atomicSub( dest, warp_count * N );

        // obtain the offset from the first participating thread
        const uint32 offset = cub::ShuffleIndex<32>(broadcast_offset, first_tid, __activemask());

        // compute the per-thread offset
        *result = offset - warp_scan * N;
    }

private:
    uint32* m_dest;
};

struct warp_atomic
{
    struct temp_storage_type
    {
        union {
            typename cub::WarpScan<uint32>::TempStorage     scan_storage;
            typename cub::WarpReduce<uint32>::TempStorage   reduce_storage;
        };
    };

    __device__ __forceinline__
    warp_atomic(uint32* dest, temp_storage_type& temp_storage)
        : m_dest(dest), m_temp_storage(temp_storage) {}


    __device__ __forceinline__
    void add(uint32 n)
    {
        warp_atomic::add( n, m_dest, m_temp_storage );
    }

    __device__ __forceinline__
    void sub(uint32 n)
    {
        warp_atomic::sub( n, m_dest, m_temp_storage );
    }

    template <uint32 N>
    __device__ __forceinline__
    void add(bool p)
    {
        return warp_atomic::static_add<N>( p, m_dest );
    }

    template <uint32 N>
    __device__ __forceinline__
    void sub(bool p)
    {
        warp_atomic::static_sub<N>( p, m_dest );
    }

    __device__ __forceinline__
    void add(uint32 n, uint32* result)
    {
        warp_atomic::add( n, m_dest, result, m_temp_storage );
    }

    __device__ __forceinline__
    void sub(uint32 n, uint32* result)
    {
        warp_atomic::sub( n, m_dest, result, m_temp_storage );
    }

    template <uint32 N>
    __device__ __forceinline__
    void add(bool p, uint32* result)
    {
        return warp_atomic::static_add<N>( p, m_dest, result );
    }

    template <uint32 N>
    __device__ __forceinline__
    void sub(bool p, uint32* result)
    {
        return warp_atomic::static_sub<N>( p, m_dest, result );
    }

    // --- stateless methods --- //

    __device__ __forceinline__
    static void add(uint32 n, uint32* dest, temp_storage_type& temp_storage)
    {
        // issue a warp-reduction
        const uint32 warp_count = cub::WarpReduce<uint32>(temp_storage.reduce_storage).Sum(n);

        // issue a per-warp atomic
        const uint32 warp_tid = threadIdx.x & 31;
        if (warp_tid == 0)
            atomicAdd( dest, warp_count );
    }

    __device__ __forceinline__
    static void sub(uint32 n, uint32* dest, temp_storage_type& temp_storage)
    {
        // issue a warp-reduction
        const uint32 warp_count = cub::WarpReduce<uint32>(temp_storage.reduce_storage).Sum(n);

        // issue a per-warp atomic
        const uint32 warp_tid = threadIdx.x & 31;
        if (warp_tid == 0)
            atomicSub( dest, warp_count );
    }

    __device__ __forceinline__
    static void add(uint32 n, uint32* dest, uint32* result, temp_storage_type& temp_storage)
    {
        uint32 warp_scan, warp_count;

        // issue a warp-scan
        cub::WarpScan<uint32>(temp_storage.scan_storage).ExclusiveSum(n, warp_scan, warp_count);

        const uint32 warp_tid = threadIdx.x & 31;

        // issue a per-warp atomic
        uint32 base_index;
        if (warp_tid == 0)
            base_index = atomicAdd( dest, warp_count );

        // compute the per-thread offset
        *result = cub::ShuffleIndex<32>( base_index, 0, __activemask() ) + warp_scan;
    }

    __device__ __forceinline__
    static void sub(uint32 n, uint32* dest, uint32* result, temp_storage_type& temp_storage)
    {
        uint32 warp_scan, warp_count;

        // issue a warp-scan
        cub::WarpScan<uint32>(temp_storage.scan_storage).ExclusiveSum(n, warp_scan, warp_count);

        const uint32 warp_tid = threadIdx.x & 31;

        // issue a per-warp atomic
        uint32 base_index;
        if (warp_tid == 0)
            base_index = atomicSub( dest, warp_count );

        // compute the per-thread offset
        *result = cub::ShuffleIndex<32>( base_index, 0, __activemask() ) - warp_scan;
    }

    template <uint32 N>
    __device__ __forceinline__
    static void static_add(bool p, uint32* dest)
    {
        return warp_static_atomic::static_add<N>( p, dest );
    }

    template <uint32 N>
    __device__ __forceinline__
    static void static_sub(bool p, uint32* dest)
    {
        warp_static_atomic::static_sub<N>( p, dest );
    }

    template <uint32 N>
    __device__ __forceinline__
    static void static_add(bool p, uint32* dest, uint32* result)
    {
        return warp_static_atomic::static_add<N>( p, dest, result );
    }

    template <uint32 N>
    __device__ __forceinline__
    static void static_sub(bool p, uint32* dest, uint32* result)
    {
        warp_static_atomic::static_sub<N>( p, dest, result );
    }

private:
    uint32*             m_dest;
    temp_storage_type&  m_temp_storage;
};


} // namespace cuda
} // namespace cugar