xgc_io.hpp

Go to the documentation of this file.

#ifndef XGC_IO_HPP
#define XGC_IO_HPP
#include "space_settings.hpp"

#ifdef ADIOS2
#include "adios2.hpp"
// If ADIOS2 is being used, XGC_IO_Stream is ADIOS2Stream
typedef ADIOS2Stream XGC_IO_Stream;
typedef adios2::Mode XGC_IO_Mode;
#else
// If ADIOS2 is not being used, spoof IO since there is currently no alternative
enum class XGC_IO_Mode{
Write,
Append,
Read,
ReadRandomAccess
};

class XGC_IO_Stream{
public:
    XGC_IO_Stream(){}
    XGC_IO_Stream(const std::string &IOName){}

    void Init(const std::string &IOName){}
    void Open(const std::string &StreamName, const XGC_IO_Mode mode){}
#ifdef USE_MPI
    void Open(const std::string &StreamName, const XGC_IO_Mode mode, const MPI_Comm &comm){}
#endif
    void Close(){}
};
#endif

// Returns an unmanaged view wrapping the input scalar
template<typename T>
inline View<T*, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> scalar_view_wrap(T& scalar){
    return View<T*, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>>(&scalar, 1);
}

struct IOData{
#ifdef ADIOS2
    // These virtual functions must be defined in derived classes
    virtual void adios_def_and_shape(const ADIOS2Stream& stream) const = 0;
    virtual void adios_put(const ADIOS2Stream& stream) const = 0;
    virtual void adios_get(const ADIOS2Stream& stream) = 0;
#endif
    virtual ~IOData(){}
};

// IOScalar is a helper struct to handle an scalar being written/read with ADIOS2
template<class T>
struct IOScalar : public IOData {
    using value_type = T;
    std::string name;
    value_type* value_addr;
    value_type value;
    bool var_is_const;

    // Input value is not const; assign address
    IOScalar(const std::string& name_in, T& value_in)
      : name(name_in),
        var_is_const(false),
        value_addr(&value_in),
        value(value_in){}

    // Input value is const; assign value
    IOScalar(const std::string& name_in, const T& value_in)
      : name(name_in),
        var_is_const(true),
        value(value_in){}

#ifdef ADIOS2
    void adios_def_and_shape(const ADIOS2Stream& stream) const {
        if(!stream.io->InquireVariable<T>(name)){
            stream.io->DefineVariable<T>(name);
        }
    }

    void adios_put(const ADIOS2Stream& stream) const {
        stream.engine->Put(name, value);
    }

    void adios_get(const ADIOS2Stream& stream) {
        if(var_is_const) printf("\nError: Tried to read into const variable\n"); // Should be compiletime check?
        else stream.engine->Get(name, *value_addr);
    }
#endif
};

// The purpose of this struct is to be able to pass types other than long long int/size_t into the adios2::Dims initialization.
struct UpcastVector {
    std::vector<size_t> vec;

    UpcastVector() = default;

    template <typename... Args>
    UpcastVector(Args... args) {
        addValues(args...);
    }

    template <typename First, typename... Rest>
    void addValues(First first, Rest... rest) {
        vec.push_back(static_cast<size_t>(first));
        if constexpr (sizeof...(rest) > 0) {
            addValues(rest...);
        }
    }
};


// IOArray is a helper struct to represent an array being written/read with ADIOS2
template<class T>
struct IOArray : public IOData {
    using view_type = T;
    using dim_type = std::vector<size_t>;
    std::string name;
    view_type view;

    // These are std::vector<size_t>
    dim_type full_size;
    dim_type offsets;
    dim_type size;

    // Dims arguments are implicitly converted to size_t
    // If the view has fewer dimensions than full_size and offsets, pad with outer dimensions of size 1, i.e. assume that
    // the intention is for the data to be written/read contiguously.
    IOArray(const std::string& name_in, const T& view_in, const UpcastVector& full_size_in, const UpcastVector& offsets_in)
      : name(name_in),
        view(view_in), // shallow copy of view
        full_size(full_size_in.vec),
        offsets(offsets_in.vec),
        size(full_size.size()) // size needs to have the same length (i.e. refer to the same number of dimensions) as full_size
    {
        // If the input view has fewer dimensions than the global view, treat the outermost dimensions as size 1
        int dim_offset = size.size() - T::rank;
        // Write the dimensions of the view into size
        for(int i=0; i<size.size(); i++){
            if(i - dim_offset >= 0){
                size[i] = view.extent(i - dim_offset);
            }else{
                size[i] = 1; // Outer dimension doesn't exist, set it to 1
            }
        }
    }

    // Constructor without full size or offset --> full_size = size, offset = 0
    IOArray(const std::string& name_in, const T& view_in)
      : name(name_in),
        view(view_in), // shallow copy of view
        full_size(T::rank),
        offsets(T::rank),
        size(T::rank)
    {
        // Write the dimensions of the view into size
        for(int i=0; i<size.size(); i++){
            size[i] = view.extent(i);

            // full_size = size, offsets are 0
            full_size[i] = size[i];
            offsets[i] = 0;
        }
    }
#ifdef ADIOS2
    void adios_def_and_shape(const ADIOS2Stream& stream) const {
        auto TAG = stream.io->InquireVariable<typename T::value_type>(name);
        if(!TAG){ 
            TAG = stream.io->DefineVariable<typename T::value_type>(name, full_size, offsets, size);
        }else{
            TAG.SetShape(full_size);
        }
    }   
    
    void adios_put(const ADIOS2Stream& stream) const {
        auto TAG = stream.io->InquireVariable<typename T::value_type>(name);
        TAG.SetSelection({offsets, size});
        stream.engine->Put<typename T::value_type>(TAG, view.data(), adios2::Mode::Sync);
    }

    void adios_get(const ADIOS2Stream& stream) {
        auto TAG = stream.io->InquireVariable<typename T::value_type>(name);
        TAG.SetSelection({offsets, size});
        stream.engine->Get<typename T::value_type>(TAG, view.data());
    }
#endif
};

/* This class contains a vector of data ready to be written/read with ADIOS2 */
class XGC_IO{
    std::vector<std::unique_ptr<IOData>> io_data;

    public:

    // Adds a single new entry to the data vector
    template<class T>
    void add(const T& new_data){
        io_data.push_back(std::make_unique<T>(new_data));
    }

    // Adds an unspecified number of entries to the data vector
    template<class T, class... Ts>
    void add(const T& new_data, const Ts&... new_data_rest){
        add(new_data);
        add(new_data_rest...);
    }

    // Writes the data to the stream
    void write(const XGC_IO_Stream& stream) const{
#ifdef ADIOS2
        // Define and shape all data
        for(int i=0; i<io_data.size(); i++)
            io_data[i]->adios_def_and_shape(stream);

        // Stage all data for writing
        for(int i=0; i<io_data.size(); i++)
            io_data[i]->adios_put(stream);

        // Perform puts
        stream.engine->PerformPuts();
#endif
    }

    // Reads the data from the stream
    void read(const XGC_IO_Stream& stream) {
#ifdef ADIOS2
        // Stage all data for reading
        for(int i=0; i<io_data.size(); i++)
            io_data[i]->adios_get(stream);

        // Perform gets
        stream.engine->PerformGets();
#endif
    }
};

#endif