loadable_buffer.hpp

Go to the documentation of this file.

#ifndef LOADABLE_BUFFER_HPP
#define LOADABLE_BUFFER_HPP

#include "space_settings.hpp"

/* This class facilitates packing Kokkos Views of arbitrary dimension into a single contiguous allocation,
 * which can be useful for data movement such as MPI operations. The buffer must be unloaded in the reverse order
 * in which it is loaded (first in, last out).
 * Currently requires that the incoming Views are of type T.
 */
template<typename T, class Device>
class LoadableBuffer{
    View<T*, CLayout, Device> buffer; 
    int offset; 

    /* size_sum is used to determine the necessary buffer size when using the constructor that takes a template pack of Views
     * The final template has only one View left, so returns the size of that View.
     */
    template<class V>
    int size_sum(V first) {
        return first.size();
    }

    /* size_sum is used to determine the necessary buffer size when using the constructor that takes a template pack of Views
     * The generic template adds the first View in the template pack, then returns the size_sum of the remaining Views in the pack
     */
    template<class V, class... TRest>
    int size_sum(V first, TRest... args) {
        return first.size() + size_sum(args...);
    }

    public:

    /* Constructor where the user specifies the size manually.
     */
    LoadableBuffer(int size)
     : buffer("buffer",size),
       offset(0)
    {}

    /* Constructor where the user inputs a list of Views, the size of which is added to determine how big the buffer needs to be.
     * This is essentially 'previewing' the loading operations to get the allocation size required.
     */
    template<class... Vs>
    LoadableBuffer(Vs... views)
     : buffer("buffer",size_sum(views...)),
       offset(0)
    {}

    /* Loads the input View into the buffer, and offsets the buffer to prepare for the next View. If do_copy is false,
     * then only the offset is performed. This saves an unnecessary copy if the copied value would have been unused, e.g.
     * for a broadcast, but is necessary so that the unload can be done properly.
     */
    template<class V>
    void load(V view, bool do_copy){
        int size_of_view = view.size();
        View<T*, CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> offset_buffer(buffer.data() + offset, size_of_view);
        View<T*, CLayout, typename V::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> unmanaged_view(view.data(), size_of_view);

        if(do_copy) Kokkos::deep_copy(offset_buffer, unmanaged_view);
        offset += size_of_view;
    }

    /* Unloads the buffer into the input View, and offsets the buffer to prepare for the next View. If do_copy is false,
     * then only the offset is performed. This saves an unnecessary copy if the copied value would have been unused, e.g.
     * for a broadcast.
     */
    template<class V>
    void unload(V view, bool do_copy){
        int size_of_view = view.size();
        offset -= size_of_view;
        View<T*, CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> offset_buffer(buffer.data() + offset, size_of_view);
        View<T*, CLayout, typename V::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> unmanaged_view(view.data(), size_of_view);

        if(do_copy) Kokkos::deep_copy(unmanaged_view, offset_buffer);
    }

    /* Returns the size of the buffer.
     */
    int size(){
        return buffer.size();
    }

    /* Returns the pointer to the buffer
     */
    T* data(){
        return buffer.data();
    }
};

#endif