AMReX_GpuBuffer.H

Go to the documentation of this file.

#ifndef AMREX_GPU_DEVICE_BUFFER_H_
#define AMREX_GPU_DEVICE_BUFFER_H_
#include <AMReX_Config.H>
 
#include <AMReX_Arena.H>
#include <AMReX_TypeTraits.H>
#include <AMReX_GpuDevice.H>
#include <cstring>
#include <cstdlib>
#include <initializer_list>
#include <memory>
 
namespace amrex::Gpu {
 
template <typename T, std::enable_if_t<std::is_trivially_copyable_v<T>,int> = 0>
class Buffer
{
public:
 
    Buffer (std::initializer_list<T> init)
        : m_size(init.size())
    {
        if (m_size == 0) { return; }
#ifdef AMREX_USE_GPU
        h_data = static_cast<T*>(The_Pinned_Arena()->alloc(m_size*sizeof(T)));
#else
        h_data = static_cast<T*>(std::malloc(m_size*sizeof(T)));
#endif
        std::memcpy(h_data, init.begin(), m_size*sizeof(T));
#ifdef AMREX_USE_GPU
        if (Gpu::inLaunchRegion())
        {
            d_data = static_cast<T*>(The_Arena()->alloc(m_size*sizeof(T)));
            Gpu::htod_memcpy_async(d_data, h_data, m_size*sizeof(T));
        }
#endif
    }
 
    Buffer (T const* h_p, const std::size_t n)
        : m_size(n)
    {
        if (m_size == 0) { return; }
#ifdef AMREX_USE_GPU
        h_data = static_cast<T*>(The_Pinned_Arena()->alloc(m_size*sizeof(T)));
#else
        h_data = static_cast<T*>(std::malloc(m_size*sizeof(T)));
#endif
        std::memcpy(h_data, h_p, m_size*sizeof(T));
#ifdef AMREX_USE_GPU
        if (Gpu::inLaunchRegion())
        {
            d_data = static_cast<T*>(The_Arena()->alloc(m_size*sizeof(T)));
            Gpu::htod_memcpy_async(d_data, h_data, m_size*sizeof(T));
        }
#endif
    }
 
    ~Buffer () { clear(); }
 
    Buffer (Buffer const&) = delete;
    Buffer (Buffer &&) = delete;
    void operator= (Buffer const&) = delete;
    void operator= (Buffer &&) = delete;
 
    [[nodiscard]] T const* data () const noexcept { return (d_data != nullptr) ? d_data : h_data; }
    [[nodiscard]] T* data () noexcept { return (d_data != nullptr) ? d_data : h_data; }
 
    [[nodiscard]] T const* hostData () const noexcept { return h_data; }
    [[nodiscard]] T* hostData () noexcept { return h_data; }
 
    [[nodiscard]] std::size_t size () const noexcept { return m_size; }
 
    void clear ()
    {
#ifdef AMREX_USE_GPU
        if (d_data) { The_Arena()->free(d_data); }
        if (h_data) { The_Pinned_Arena()->free(h_data); }
#else
        std::free(h_data);
#endif
        d_data = nullptr;
        h_data = nullptr;
    }
 
    T* copyToHost ()
    {
#ifdef AMREX_USE_GPU
        if (d_data)
        {
            Gpu::dtoh_memcpy_async(h_data, d_data, m_size*sizeof(T));
            Gpu::streamSynchronize();
        }
#endif
        return h_data;
    }
 
private:
    std::size_t m_size;
    T* d_data = nullptr;
    T* h_data = nullptr;
};
 
}
 
#endif