allocator.h

/*
 * cugar
 * Copyright (c) 2011-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 the 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 NVIDIA CORPORATION 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/basic/threads.h>
#include <cugar/basic/cuda/arch.h>
#include <cub/util_allocator.cuh>
#include <thrust/device_ptr.h>

namespace cugar {




struct byte_caching_device_allocator
{
  public:
    typedef char              value_type;
    typedef char*             pointer;
    typedef const char*       const_pointer;
    typedef char&             reference;
    typedef const char&       const_reference;
    typedef size_t            size_type;
    typedef int64             difference_type;

    CUGAR_HOST_DEVICE
    char* allocate(size_type num_bytes)
    {
      #if defined(CUGAR_DEVICE_COMPILATION)
        return NULL;
      #else
        if (s_caching_device_allocator == NULL)
            init();

        void* ptr;
        cuda::check_error( s_caching_device_allocator->DeviceAllocate( &ptr, num_bytes ), "cugar::caching_device_allocator::allocate()" );
        return (char*)ptr;
      #endif
    }

    CUGAR_HOST_DEVICE
    void deallocate(char *ptr, size_type n)
    {
      #if !defined(CUGAR_DEVICE_COMPILATION)
        cuda::check_error( s_caching_device_allocator->DeviceFree( ptr ), "cugar::caching_device_allocator::deallocate()" );
      #endif
    }

    static void init(const size_t max_cached_bytes = 128u*1024u*1024u)
    {
        ScopedLock scoped_lock( &s_mutex );
        if (s_caching_device_allocator == NULL)
            s_caching_device_allocator = new cub::CachingDeviceAllocator( 8, 2u, 10u, max_cached_bytes );
    }

private:
    CUGAR_API static cub::CachingDeviceAllocator* volatile s_caching_device_allocator;
    static Mutex                                 s_mutex;
};

template <typename T>
struct caching_device_allocator
{
    typedef T                                       value_type;

    typedef thrust::device_ptr<T>                   pointer;

    typedef thrust::device_ptr<const T>             const_pointer;

    typedef thrust::device_reference<T>             reference;

    typedef thrust::device_reference<const T>       const_reference;

    typedef std::size_t                             size_type;

    typedef typename pointer::difference_type       difference_type;

    __host__ __device__
    inline pointer address(reference r) { return &r; }
    
    __host__ __device__
    inline const_pointer address(const_reference r) { return &r; }

    CUGAR_HOST
    inline pointer allocate(size_type cnt,
                            const_pointer = const_pointer(static_cast<T*>(0)))
    {
        return pointer( (T*)base_allocator.allocate( cnt * sizeof(T) ) );
    } // end allocate()

    CUGAR_HOST
    inline void deallocate(pointer p, size_type cnt)
    {
        base_allocator.deallocate((char*)p.get(), cnt * sizeof(T));
    } // end deallocate()

private:
    byte_caching_device_allocator base_allocator;
};


} // namespace cugar