view release on metacpan or  search on metacpan

xgboost/cub/cub/grid/grid_queue.cuh  view on Meta::CPAN

#include "../util_namespace.cuh"
#include "../util_debug.cuh"

/// Optional outer namespace(s)
CUB_NS_PREFIX

/// CUB namespace
namespace cub {


/**
 * \addtogroup GridModule
 * @{
 */


/**
 * \brief GridQueue is a descriptor utility for dynamic queue management.
 *
 * \par Overview
 * GridQueue descriptors provides abstractions for "filling" or
 * "draining" globally-shared vectors.
 *
 * \par
 * A "filling" GridQueue works by atomically-adding to a zero-initialized counter,
 * returning a unique offset for the calling thread to write its items.
 * The GridQueue maintains the total "fill-size".  The fill counter must be reset
 * using GridQueue::ResetFill by the host or kernel instance prior to the kernel instance that
 * will be filling.
 *
 * \par
 * Similarly, a "draining" GridQueue works by works by atomically-incrementing a
 * zero-initialized counter, returning a unique offset for the calling thread to
 * read its items. Threads can safely drain until the array's logical fill-size is
 * exceeded.  The drain counter must be reset using GridQueue::ResetDrain or
 * GridQueue::FillAndResetDrain by the host or kernel instance prior to the kernel instance that
 * will be filling.  (For dynamic work distribution of existing data, the corresponding fill-size
 * is simply the number of elements in the array.)
 *
 * \par
 * Iterative work management can be implemented simply with a pair of flip-flopping
 * work buffers, each with an associated set of fill and drain GridQueue descriptors.
 *
 * \tparam OffsetT Signed integer type for global offsets
 */
template <typename OffsetT>
class GridQueue
{
private:

    /// Counter indices
    enum
    {
        FILL    = 0,
        DRAIN   = 1,
    };

    /// Pair of counters
    OffsetT *d_counters;

public:

    /// Returns the device allocation size in bytes needed to construct a GridQueue instance
    __host__ __device__ __forceinline__
    static size_t AllocationSize()
    {
        return sizeof(OffsetT) * 2;
    }


    /// Constructs an invalid GridQueue descriptor
    __host__ __device__ __forceinline__ GridQueue()
    :
        d_counters(NULL)
    {}


    /// Constructs a GridQueue descriptor around the device storage allocation
    __host__ __device__ __forceinline__ GridQueue(
        void *d_storage)                    ///< Device allocation to back the GridQueue.  Must be at least as big as <tt>AllocationSize()</tt>.
    :
        d_counters((OffsetT*) d_storage)
    {}


    /// This operation sets the fill-size and resets the drain counter, preparing the GridQueue for draining in the next kernel instance.  To be called by the host or by a kernel prior to that which will be draining.
    __host__ __device__ __forceinline__ cudaError_t FillAndResetDrain(
        OffsetT fill_size,
        cudaStream_t stream = 0)
    {
#if (CUB_PTX_ARCH > 0)
        (void)stream;
        d_counters[FILL] = fill_size;
        d_counters[DRAIN] = 0;
        return cudaSuccess;
#else
        OffsetT counters[2];
        counters[FILL] = fill_size;
        counters[DRAIN] = 0;
        return CubDebug(cudaMemcpyAsync(d_counters, counters, sizeof(OffsetT) * 2, cudaMemcpyHostToDevice, stream));
#endif
    }


    /// This operation resets the drain so that it may advance to meet the existing fill-size.  To be called by the host or by a kernel prior to that which will be draining.
    __host__ __device__ __forceinline__ cudaError_t ResetDrain(cudaStream_t stream = 0)
    {
#if (CUB_PTX_ARCH > 0)
        (void)stream;
        d_counters[DRAIN] = 0;
        return cudaSuccess;
#else
        return CubDebug(cudaMemsetAsync(d_counters + DRAIN, 0, sizeof(OffsetT), stream));
#endif
    }


    /// This operation resets the fill counter.  To be called by the host or by a kernel prior to that which will be filling.
    __host__ __device__ __forceinline__ cudaError_t ResetFill(cudaStream_t stream = 0)
    {
#if (CUB_PTX_ARCH > 0)