TrackedVector.H

Go to the documentation of this file.

/* Copyright 2022-2023 The Regents of the University of California, through Lawrence
*           Berkeley National Laboratory (subject to receipt of any required
 *           approvals from the U.S. Dept. of Energy). All rights reserved.
 *
 * This file is part of ImpactX.
 *
 * FIXME: SHOULD BE REPLACED WITH amrex::Gpu::TrackedVector in v26.05+, see
 *        https://github.com/AMReX-Codes/amrex/pull/5253
 *
 * Authors: Axel Huebl
 * License: BSD-3-Clause-LBNL
 */
#ifndef IMPACTX_TRACKED_VECTOR_H
#define IMPACTX_TRACKED_VECTOR_H
 
#include <AMReX.H>
#include <AMReX_GpuContainers.H>
 
#include <memory>
#include <type_traits>
#include <utility>
#include <vector>
 
 
namespace impactx::elements::mixin
{
    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;
 
    private:
        void register_finalize () {
#ifdef AMREX_USE_GPU
            if (*m_finalize_registered) { return; }
            *m_finalize_registered = true;
 
            std::weak_ptr<std::vector<T>> weak_host = m_host;
            std::weak_ptr<amrex::Gpu::DeviceVector<T>> weak_device = m_device;
            std::weak_ptr<Status> weak_status = m_status;
            std::weak_ptr<bool> weak_registered = m_finalize_registered;
 
            amrex::ExecOnFinalize([weak_host, weak_device, weak_status, weak_registered]() {
                // see: release_gpu()
                auto host   = weak_host.lock();
                auto device = weak_device.lock();
                auto status = weak_status.lock();
                if (host && device && status
                    && *status == Status::device_dirty
                    && !device->empty())
                {
                    // see: to_host()
                    host->resize(device->size());
                    amrex::Gpu::copy(amrex::Gpu::deviceToHost,
                        device->begin(), device->end(), host->begin());
                }
                if (device) {
                    device->clear();
                    device->shrink_to_fit();
                }
                if (status) {
                    *status = Status::host_dirty;
                }
                if (auto reg = weak_registered.lock()) {
                    *reg = false;
                }
            });
#endif
        }
    public:
 
        TrackedVector () = default;
 
        explicit TrackedVector (size_type a_size)
            : m_host(std::make_shared<std::vector<T>>(a_size))
        {
#ifdef AMREX_USE_GPU
            *m_status = Status::host_dirty;
#endif
        }
 
        TrackedVector (size_type a_size, value_type const & a_value)
            : m_host(std::make_shared<std::vector<T>>(a_size, a_value))
        {
#ifdef AMREX_USE_GPU
            *m_status = Status::host_dirty;
#endif
        }
 
        TrackedVector (std::initializer_list<T> a_initializer_list)
            : m_host(std::make_shared<std::vector<T>>(a_initializer_list))
        {
#ifdef AMREX_USE_GPU
            *m_status = Status::host_dirty;
#endif
        }
 
        TrackedVector (std::vector<T> a_vector)
            : m_host(std::make_shared<std::vector<T>>(std::move(a_vector)))
        {
#ifdef AMREX_USE_GPU
            *m_status = Status::host_dirty;
#endif
        }
 
        TrackedVector (TrackedVector const & a_vector)
            : m_host(std::make_shared<std::vector<T>>(*a_vector.m_host))
        {
            *m_status = *a_vector.m_status;
#ifdef AMREX_USE_GPU
            *m_device = *a_vector.m_device;
#endif
        }

        TrackedVector (TrackedVector && a_vector) noexcept
        {
            std::swap(m_status, a_vector.m_status);
            std::swap(m_host, a_vector.m_host);
#ifdef AMREX_USE_GPU
            std::swap(m_device, a_vector.m_device);
            std::swap(m_finalize_registered, a_vector.m_finalize_registered);
#endif
        }
 
        TrackedVector& operator= (TrackedVector const & a_vector) {
            if (this != &a_vector) {
                *m_status = *a_vector.m_status;
                *m_host = *a_vector.m_host;
#ifdef AMREX_USE_GPU
                *m_device = *a_vector.m_device;
#endif
            }
            return *this;
        }

        TrackedVector& operator= (TrackedVector && a_vector) noexcept {
            if (this != &a_vector) {
                std::swap(m_host, a_vector.m_host);
                std::swap(m_status, a_vector.m_status);
#ifdef AMREX_USE_GPU
                std::swap(m_device, a_vector.m_device);
                std::swap(m_finalize_registered, a_vector.m_finalize_registered);
#endif
            }
            return *this;
        }
 
        ~TrackedVector () = default;
 
        enum class Status {
            up_to_date,    
            device_dirty,  
            host_dirty     
        };
 
        [[nodiscard]] Status status () const { return *m_status; }

        [[nodiscard]] std::vector<T> &
        host () {
#ifdef AMREX_USE_GPU
            if (*m_status == Status::device_dirty) {
                to_host();
            }
            *m_status = Status::host_dirty;
#endif
            return *m_host;
        }

        [[nodiscard]] std::vector<T> const &
        host_const () const {
#ifdef AMREX_USE_GPU
            if (*m_status == Status::device_dirty) {
                to_host();
            }
#endif
            return *m_host;
        }
 
#ifdef AMREX_USE_GPU
        [[nodiscard]] amrex::Gpu::DeviceVector<T> &
        device () {
            if (!amrex::Initialized()) {
                throw std::runtime_error("TrackedVector::device() called before AMReX initialize/after AMReX finalize");
            }
            register_finalize();
            if (*m_status == Status::host_dirty) {
                to_device();
            }
            *m_status = Status::device_dirty;
            return *m_device;
        }

        [[nodiscard]] amrex::Gpu::DeviceVector<T> const &
        device_const () const {
            if (!amrex::Initialized()) {
                throw std::runtime_error("TrackedVector::device_const() called before AMReX initialize/after AMReX finalize");
            }
            const_cast<TrackedVector*>(this)->register_finalize();
            if (*m_status == Status::host_dirty) {
                to_device();
            }
            return *m_device;
        }
#else
        [[nodiscard]] std::vector<T> &
        device () { return *m_host; }

        [[nodiscard]] std::vector<T> const &
        device_const () const { return *m_host; }
#endif

        void release_gpu ()
        {
#ifdef AMREX_USE_GPU
            if (*m_status == Status::device_dirty) {
                to_host();
            }
            m_device->clear();
            m_device->shrink_to_fit();
            *m_status = Status::host_dirty;
#endif
        }
 
    private:

        void to_device () const {
#ifdef AMREX_USE_GPU
            if (!amrex::Initialized()) {
                throw std::runtime_error("TrackedVector::to_device() called outside of AMReX initialize/finalize");
            }
            auto const size = m_host->size();
            if (size > 0U) {
                m_device->resize(size);
                amrex::Gpu::copy(amrex::Gpu::hostToDevice,
                    m_host->begin(), m_host->end(), m_device->begin());
            } else {
                m_device->clear();
            }
            *m_status = Status::up_to_date;
#endif
        }

        void to_host () const {
#ifdef AMREX_USE_GPU
            if (!amrex::Initialized()) {
                throw std::runtime_error("TrackedVector::to_host() called outside of AMReX initialize/finalize");
            }
            m_host->resize(m_device->size());
            amrex::Gpu::copy(amrex::Gpu::deviceToHost,
                m_device->begin(), m_device->end(), m_host->begin());
            *m_status = Status::up_to_date;
#endif
        }
 
        mutable std::shared_ptr<Status> m_status = std::make_shared<Status>();
        mutable std::shared_ptr<std::vector<T>> m_host = std::make_shared<std::vector<T>>();
#ifdef AMREX_USE_GPU
        mutable std::shared_ptr<amrex::Gpu::DeviceVector<T>> m_device = std::make_shared<amrex::Gpu::DeviceVector<T>>();
        mutable std::shared_ptr<bool> m_finalize_registered = std::make_shared<bool>(false);
#endif
    };
 
}
 
#endif