AMReX_TrackedVector.H

Go to the documentation of this file.

#ifndef AMREX_TRACKED_VECTOR_H
#define AMREX_TRACKED_VECTOR_H
 
#include <AMReX_Config.H>
 
#include <AMReX.H>
#include <AMReX_GpuContainers.H>
 
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
 
namespace amrex::Gpu
{
    template <class T>
    struct TrackedVector
    {
        static_assert(std::is_trivially_copyable<T>(), "TrackedVector can only hold trivially copyable types");
        using value_type      = T;
        using size_type       = std::size_t;
 
#ifdef AMREX_USE_GPU
        using device_vector_type = amrex::Gpu::NonManagedDeviceVector<T>;
#else
        using device_vector_type = std::vector<T>;
#endif
 
        enum class Status {
            up_to_date,    
            device_dirty,  
            host_dirty     
        };
 
        TrackedVector () = default;
 
        explicit TrackedVector (size_type a_size)
            : m_data(std::make_shared<Data>(std::vector<T>(a_size))) {}
 
        TrackedVector (size_type a_size, value_type const & a_value)
            : m_data(std::make_shared<Data>(std::vector<T>(a_size, a_value))) {}
 
        TrackedVector (std::initializer_list<T> a_initializer_list)
            : m_data(std::make_shared<Data>(std::vector<T>(a_initializer_list))) {}
 
        TrackedVector (std::vector<T> a_vector)
            : m_data(std::make_shared<Data>(std::move(a_vector))) {}
 
        TrackedVector (TrackedVector const & a_vector)
        {
            copy_from(a_vector);
        }
 
        TrackedVector (TrackedVector && a_vector) noexcept
        {
            std::swap(m_data, a_vector.m_data);
        }
 
        TrackedVector& operator= (TrackedVector const & a_vector) {
            if (this != &a_vector) {
                copy_from(a_vector);
            }
            return *this;
        }
 
        TrackedVector& operator= (TrackedVector && a_vector) noexcept {
            if (this != &a_vector) {
                std::swap(m_data, a_vector.m_data);
            }
            return *this;
        }
 
        ~TrackedVector () = default;
 
        [[nodiscard]] Status status () const { return m_data->status; }
 
        [[nodiscard]] std::vector<T> &
        host () {
#ifdef AMREX_USE_GPU
            if (m_data->status == Status::device_dirty) {
                to_host();
            }
            m_data->status = Status::host_dirty;
#endif
            return m_data->host;
        }
 
        [[nodiscard]] std::vector<T> const &
        host_const () const {
#ifdef AMREX_USE_GPU
            if (m_data->status == Status::device_dirty) {
                to_host();
            }
#endif
            return m_data->host;
        }
 
#ifdef AMREX_USE_GPU
        [[nodiscard]] device_vector_type &
        device () {
            require_amrex("device");
            register_finalize();
            if (m_data->status == Status::host_dirty) {
                to_device();
            }
            m_data->status = Status::device_dirty;
            return m_data->device;
        }
 
        [[nodiscard]] device_vector_type const &
        device_const () const {
            require_amrex("device_const");
            register_finalize();
            if (m_data->status == Status::host_dirty) {
                to_device();
            }
            return m_data->device;
        }
#else
        [[nodiscard]] device_vector_type &
        device () { return m_data->host; }
 
        [[nodiscard]] device_vector_type const &
        device_const () const { return m_data->host; }
#endif
 
        void release_gpu ()
        {
#ifdef AMREX_USE_GPU
            Data::drain_device(*m_data);
#endif
        }
 
    private:
 
        struct Data {
            Status status = Status::up_to_date;
            std::vector<T> host;
#ifdef AMREX_USE_GPU
            device_vector_type device;
            bool finalize_registered = false;
#endif
            Data () = default;
 
            explicit Data (std::vector<T> h)
                : host(std::move(h))
            {
#ifdef AMREX_USE_GPU
                status = Status::host_dirty;
#endif
            }
 
#ifdef AMREX_USE_GPU
            static void drain_device (Data& d) {
                if (d.status == Status::device_dirty) {
                    amrex::Gpu::streamSynchronize();
                    d.host.resize(d.device.size());
                    if (!d.device.empty()) {
                        amrex::Gpu::copy(amrex::Gpu::deviceToHost,
                            d.device.begin(), d.device.end(), d.host.begin());
                    }
                }
                d.device.clear();
                d.device.shrink_to_fit();
                d.status = Status::host_dirty;
            }
#endif
        };
 
        static void require_amrex (char const * func) {
            if (!amrex::Initialized()) {
                throw std::runtime_error(
                    std::string("TrackedVector::") + func +
                    " called outside of AMReX initialize/finalize");
            }
        }
 
#ifdef AMREX_USE_GPU
        static void require_copyable_without_amrex (Data const& src, Data const& dst)
        {
            if (amrex::Initialized()) {
                return;
            }
            if (src.status == Status::device_dirty) {
                throw std::runtime_error(
                    "TrackedVector::copy_from: source is device_dirty "
                    "outside an AMReX session (host data is stale)");
            }
            if (!src.device.empty()) {
                throw std::runtime_error(
                    "TrackedVector::copy_from: source retains device storage "
                    "outside an AMReX session");
            }
            if (!dst.device.empty()) {
                throw std::runtime_error(
                    "TrackedVector::copy_from: destination retains device storage "
                    "outside an AMReX session");
            }
            if (src.finalize_registered) {
                throw std::runtime_error(
                    "TrackedVector::copy_from: source finalize callback still registered "
                    "outside an AMReX session");
            }
            if (dst.finalize_registered) {
                throw std::runtime_error(
                    "TrackedVector::copy_from: destination finalize callback still registered "
                    "outside an AMReX session");
            }
        }
#endif
 
        void copy_from (TrackedVector const & a_vector) {
#ifdef AMREX_USE_GPU
            Data const& src = *a_vector.m_data;
            require_copyable_without_amrex(src, *m_data);
 
            // Copy status & host data
            m_data->status = src.status;
            m_data->host = src.host;
 
            // Device data & finalize registration:
            // Device copy only makes sense inside a live AMReX session where
            // the device arena is available.  Outside a session both sides
            // must already be empty (drained by Finalize / never allocated).
            if (amrex::Initialized()) {
                if (src.status == Status::host_dirty) {
                    // host_dirty means the device mirror is stale —
                    // avoid a wasteful D->D copy; it will resync on demand.
                    m_data->device.clear();
                    m_data->device.shrink_to_fit();
                } else {
                    // Ensure pending kernels that dirtied src.device
                    // have completed before the D->D copy reads it.
                    if (src.status == Status::device_dirty) {
                        amrex::Gpu::streamSynchronize();
                    }
                    m_data->device = src.device;
                }
 
                if (!m_data->finalize_registered) {
                    register_finalize();
                }
            }
#else
            *m_data = *a_vector.m_data;
#endif
        }
 
        void to_device () const {
#ifdef AMREX_USE_GPU
            require_amrex("to_device");
            auto const size = m_data->host.size();
            if (size > 0U) {
                m_data->device.resize(size);
                amrex::Gpu::copy(amrex::Gpu::hostToDevice,
                    m_data->host.begin(), m_data->host.end(), m_data->device.begin());
            } else {
                m_data->device.clear();
                m_data->device.shrink_to_fit();
            }
            m_data->status = Status::up_to_date;
#endif
        }
 
        void to_host () const {
#ifdef AMREX_USE_GPU
            require_amrex("to_host");
            amrex::Gpu::streamSynchronize();
            m_data->host.resize(m_data->device.size());
            if (!m_data->device.empty()) {
                amrex::Gpu::copy(amrex::Gpu::deviceToHost,
                    m_data->device.begin(), m_data->device.end(), m_data->host.begin());
            }
            m_data->status = Status::up_to_date;
#endif
        }
 
        void register_finalize () const {
#ifdef AMREX_USE_GPU
            require_amrex("register_finalize");
            if (m_data->finalize_registered) { return; }
            m_data->finalize_registered = true;
 
            std::weak_ptr<Data> weak_data = m_data;
 
            amrex::ExecOnFinalize([weak_data]() {
                auto data = weak_data.lock();
                if (!data) { return; }
                Data::drain_device(*data);
                data->finalize_registered = false;
            });
#endif
        }
 
        // mutable: lazy-sync accessors (host_const, device_const) need to
        // update status and trigger copies through the shared_ptr even on
        // const TrackedVector references.
        mutable std::shared_ptr<Data> m_data = std::make_shared<Data>();
    };
 
}
 
#endif