nc_mpi_file.h

Go to the documentation of this file.

#pragma once
 
#include "nc_file.h"
#include "dg/backend/mpi_datatype.h"
 
namespace dg
{
namespace file
{
 
struct MPINcFile
{
    using Hyperslab = MPINcHyperslab;
    // ///////////////////////////// CONSTRUCTORS/DESTRUCTOR /////////
    MPINcFile (MPI_Comm comm = MPI_COMM_WORLD)
    : m_comm(comm)
    {}
    MPINcFile(const std::filesystem::path& filename, enum NcFileMode mode =
        nc_nowrite, MPI_Comm comm = MPI_COMM_WORLD)
    : m_comm(comm)
    {
        open( filename, mode);
    }
    MPINcFile(const MPINcFile& rhs) = delete;
    MPINcFile& operator =(const MPINcFile & rhs) = delete;
 
    MPINcFile(MPINcFile&& rhs) = default;
    MPINcFile& operator =(MPINcFile && rhs)
    {
        if( this!= &rhs)
        {
            // We need to check that the rhs has the same communicator
            // else we need to think about different groups (do all ranks in both
            // groups call this function?)
            int result;
            MPI_Comm_compare( this->m_comm, rhs.m_comm, &result);
            // congruent, similar and ident all should be fine
            assert( result != MPI_UNEQUAL);
            this->m_comm  = rhs.m_comm;
            this->m_rank0 = rhs.m_rank0;
            this->m_readonly = rhs.m_readonly;
            this->m_file    = std::move( rhs.m_file);
            this->m_buffer  = std::move( rhs.m_buffer);
            this->m_receive = std::move( rhs.m_receive);
        }
        return *this;
    }
 
    ~MPINcFile() = default;
    // /////////////////// open/close /////////
 
    void open(const std::filesystem::path& filename,
            enum NcFileMode mode = nc_nowrite)
    {
        // General problem: HDF5 may use file locks to prevent multiple processes
        // from opening the same file for write at the same time
        int rank;
        MPI_Comm_rank( m_comm, &rank);
        m_rank0 = (rank == 0);
 
        m_readonly = ( mode == nc_nowrite);
 
        // Classic file access, one process writes, everyone else reads
        mpi_invoke_void( &SerialNcFile::open, m_file, filename, mode);
        MPI_Barrier( m_comm); // all ranks agree that file exists
    }
    bool is_open() const
    {
        return mpi_invoke( &SerialNcFile::is_open, m_file);
    }
 
    void close()
    {
        mpi_invoke_void( &SerialNcFile::close, m_file);
        MPI_Barrier( m_comm); // all ranks agree that file is closed
        // removes lock from file
    }
 
    void sync()
    {
        mpi_invoke_void( &SerialNcFile::sync, m_file);
    }
    int get_ncid() const noexcept
    {
        return m_file.get_ncid();
    }
 
    int get_format()const
    {
        return mpi_invoke( &SerialNcFile::get_format, m_file);
    }
 
    MPI_Comm communicator() const { return m_comm;}
 
    // ///////////// Groups /////////////////
    void def_grp( std::string name)
    {
        mpi_invoke_void( &SerialNcFile::def_grp, m_file, name);
    }
    void def_grp_p( std::filesystem::path path)
    {
        mpi_invoke_void( &SerialNcFile::def_grp_p, m_file, path);
    }
    bool grp_is_defined( std::filesystem::path path) const
    {
        return mpi_invoke( &SerialNcFile::grp_is_defined, m_file, path);
    }
    void set_grp( std::filesystem::path path = "")
    {
        mpi_invoke_void( &SerialNcFile::set_grp, m_file, path);
    }
    void rename_grp( std::string old_name, std::string new_name)
    {
        mpi_invoke_void( &SerialNcFile::rename_grp, m_file, old_name, new_name);
    }
 
    int get_grpid() const noexcept
    {
        return m_file.get_grpid();
    }
 
    std::filesystem::path get_current_path( ) const
    {
        return mpi_invoke( &SerialNcFile::get_current_path, m_file);
    }
 
    std::list<std::filesystem::path> get_grps( ) const
    {
        return mpi_invoke( &SerialNcFile::get_grps, m_file);
    }
    std::list<std::filesystem::path> get_grps_r( ) const
    {
        return mpi_invoke( &SerialNcFile::get_grps_r, m_file);
    }
 
    // //////////// Dimensions ////////////////////////
    void def_dim( std::string name, size_t size)
    {
        mpi_invoke_void( &SerialNcFile::def_dim, m_file, name, size);
    }
    void rename_dim( std::string old_name, std::string new_name)
    {
        mpi_invoke_void( &SerialNcFile::rename_dim, m_file, old_name, new_name);
    }
    size_t get_dim_size( std::string name) const
    {
        return mpi_invoke( &SerialNcFile::get_dim_size, m_file, name);
    }
 
    std::vector<size_t> get_dims_shape( const std::vector<std::string>& dims) const
    {
        return mpi_invoke( &SerialNcFile::get_dims_shape, m_file, dims);
    }
    std::vector<std::string> get_dims(bool include_parents = true) const
    {
        return mpi_invoke( &SerialNcFile::get_dims, m_file, include_parents);
    }
    std::vector<std::string> get_unlim_dims() const
    {
        return mpi_invoke( &SerialNcFile::get_unlim_dims, m_file);
    }
    bool dim_is_defined( std::string name) const
    {
        return mpi_invoke( &SerialNcFile::dim_is_defined, m_file, name);
    }
    // ///////////// Attributes setters
    void put_att ( const std::pair<std::string, nc_att_t>& att, std::string id = "")
    {
        // Fix unresolved type following
        // https://stackoverflow.com/questions/45505017/how-comes-stdinvoke-does-not-handle-function-overloads
        mpi_invoke_void( [this]( auto&&...xs) { m_file.put_att(
            std::forward<decltype(xs)>(xs)...);}, att, id);
    }
    template<class S, class T>
    void put_att( const std::tuple<S,nc_type, T>& att, std::string id = "")
    {
        mpi_invoke_void( &SerialNcFile::put_att<S,T>, m_file, att, id);
    }
    template<class Attributes = std::map<std::string, nc_att_t> > // *it must be usable in put_att
    void put_atts( const Attributes& atts, std::string id = "")
    {
        mpi_invoke_void( &SerialNcFile::put_atts<Attributes>, m_file, atts, id);
    }
 
    // ///////////////// Attribute getters
 
    template<class T>
    T get_att_as(std::string att_name, std::string id = "") const
    {
        return mpi_invoke( &SerialNcFile::get_att_as<T>, m_file, att_name, id);
    }
    template<class T>
    std::vector<T> get_att_vec_as( std::string att_name, std::string id = "") const
    {
        return mpi_invoke( &SerialNcFile::get_att_vec_as<T>, m_file,
            att_name, id);
    }
 
    template<class T>
    std::map<std::string, T> get_atts_as( std::string id = "") const
    {
        return mpi_invoke( &SerialNcFile::get_atts_as<T>, m_file, id);
    }
    std::map<std::string, nc_att_t> get_atts( std::string id = "") const
    {
        return get_atts_as<nc_att_t>( id);
    }
 
    void del_att( std::string att_name, std::string id = "")
    {
        mpi_invoke_void( &SerialNcFile::del_att, m_file, att_name, id);
    }
    bool att_is_defined( std::string att_name, std::string id = "") const
    {
        return mpi_invoke( &SerialNcFile::att_is_defined, m_file, att_name, id);
    }
    void rename_att( std::string old_att_name, std::string
        new_att_name, std::string id = "")
    {
        mpi_invoke_void( &SerialNcFile::rename_att, m_file, old_att_name,
            new_att_name, id);
    }
 
    // //////////// Variables ////////////////////////
    template<class T, class Attributes = std::map<std::string, nc_att_t>>
    void def_var_as( std::string name,
        const std::vector<std::string>& dim_names,
        const Attributes& atts = {})
    {
        mpi_invoke_void( &SerialNcFile::def_var_as<T, Attributes>, m_file, name, dim_names, atts);
    }
    template<class Attributes = std::map<std::string, nc_att_t>>
    void def_var( std::string name, nc_type xtype,
            const std::vector<std::string>& dim_names,
            const Attributes& atts = {})
    {
        mpi_invoke_void( &SerialNcFile::def_var<Attributes>, m_file, name,
            xtype, dim_names, atts);
    }
 
    template<class ContainerType, typename = std::enable_if_t<
        dg::is_vector_v<ContainerType, dg::SharedVectorTag> or
        dg::is_vector_v<ContainerType, dg::MPIVectorTag>>>
    void put_var( std::string name, const MPINcHyperslab& slab,
            const ContainerType& data)
    {
        int grpid = m_file.get_grpid(), varid = 0;
        int e;
        file::NC_Error_Handle err;
        if( m_readonly or m_rank0)
        {
            e = nc_inq_varid( grpid, name.c_str(), &varid);
        }
        if ( not m_readonly)
        {
            MPI_Bcast( &e, 1, dg::getMPIDataType<int>(), 0, m_comm);
        }
        err = e;
        using value_type = dg::get_value_type<ContainerType>;
        m_receive.template set<value_type>(0);
        auto& receive = m_receive.template get<value_type>( );
        const auto& data_ref = get_ref( data, dg::get_tensor_category<ContainerType>());
        if constexpr ( dg::has_policy_v<ContainerType, dg::CudaTag>)
        {
            m_buffer.template set<value_type>( data.size());
            auto& buffer = m_buffer.template get<value_type>( );
            dg::assign ( data_ref, buffer);
            detail::put_vara_detail( grpid, varid, slab, buffer, receive, m_comm);
        }
        else
            detail::put_vara_detail( grpid, varid, slab, data_ref, receive, m_comm);
    }
 
    template<class ContainerType, class Attributes = std::map<std::string, nc_att_t>,
        typename = std::enable_if_t<
        dg::is_vector_v<ContainerType, dg::SharedVectorTag> or
        dg::is_vector_v<ContainerType, dg::MPIVectorTag>>>
    void defput_var( std::string name, const std::vector<std::string>& dim_names,
            const Attributes& atts, const MPINcHyperslab& slab,
            const ContainerType& data)
    {
        def_var_as<dg::get_value_type<ContainerType>>( name, dim_names, atts);
        put_var( name, slab, data);
    }
 
    template<class T, typename = std::enable_if_t<dg::is_scalar_v<T>>>
    void put_var( std::string name, const std::vector<size_t>& start, T data)
    {
        mpi_invoke_void( &SerialNcFile::put_var<T>, m_file, name, start, data);
    }
 
    template<class T, class Attributes = std::map<std::string, nc_att_t>>
    void def_dimvar_as( std::string name, size_t size, const Attributes& atts)
    {
        mpi_invoke_void( &SerialNcFile::def_dimvar_as<T>, m_file, name, size, atts);
    }
    template<class ContainerType, class Attributes = std::map<std::string, nc_att_t>>
    void defput_dim( std::string name,
            const Attributes& atts,
            const MPI_Vector<ContainerType>& abscissas)  // implicitly assume ordered by rank
    {
        unsigned size = abscissas.size(), global_size = 0;
        MPI_Reduce( &size, &global_size, 1, MPI_UNSIGNED, MPI_SUM, 0,
            abscissas.communicator());
        def_dimvar_as<dg::get_value_type<ContainerType>>( name, global_size,
            atts);
        put_var( name, {abscissas}, abscissas);
    }
 
    template<class ContainerType, typename = std::enable_if_t<
        dg::is_vector_v<ContainerType, dg::SharedVectorTag> or
        dg::is_vector_v<ContainerType, dg::MPIVectorTag>>>
    void get_var( std::string name, const MPINcHyperslab& slab,
            ContainerType& data) const
    {
        int grpid = m_file.get_grpid(), varid = 0; // grpid does not throw
        int e;
        file::NC_Error_Handle err;
        if( m_readonly or m_rank0)
        {
            e = nc_inq_varid( grpid, name.c_str(), &varid);
        }
        if ( not m_readonly)
        {
            MPI_Bcast( &e, 1, dg::getMPIDataType<int>(), 0, m_comm);
        }
        err = e;
 
        using value_type = dg::get_value_type<ContainerType>;
        auto& data_ref = get_ref( data, dg::get_tensor_category<ContainerType>());
        set_comm( data, slab.communicator(), dg::get_tensor_category<ContainerType>());
        unsigned size = 1;
        for( unsigned u=0; u<slab.ndim(); u++)
            size *= slab.count()[u];
 
        if constexpr ( dg::has_policy_v<ContainerType, dg::CudaTag>)
        {
            m_buffer.template set<value_type>( size);
            auto& buffer = m_buffer.template get<value_type>( );
            if( m_readonly)
                err = detail::get_vara_T( grpid, varid,
                    slab.startp(), slab.countp(), buffer.data());
            else
            {
                m_receive.template set<value_type>( 0);
                auto& receive = m_receive.template get<value_type>( );
                detail::get_vara_detail( grpid, varid, slab, buffer, receive, m_comm);
            }
            dg::assign ( buffer, data_ref);
        }
        else
        {
            data_ref.resize( size);
            if( m_readonly)
                err = detail::get_vara_T( grpid, varid,
                    slab.startp(), slab.countp(),
                    thrust::raw_pointer_cast(data_ref.data()));
            else
            {
                m_receive.template set<value_type>( 0);
                auto& receive = m_receive.template get<value_type>( );
                detail::get_vara_detail( grpid, varid, slab, data_ref, receive, m_comm);
            }
        }
    }
 
    template<class T, typename = std::enable_if_t<dg::is_scalar_v<T>> >
    void get_var( std::string name, const std::vector<size_t>& start, T& data) const
    {
        mpi_invoke_void( &SerialNcFile::get_var<T>, m_file, name, start, data);
        if( not m_readonly)
            mpi_bcast( data);
    }
 
    template<class ContainerType, typename = std::enable_if_t<
        dg::is_vector_v<ContainerType, dg::SharedVectorTag> or
        dg::is_vector_v<ContainerType, dg::MPIVectorTag>>>
    ContainerType get_var_as( std::string name, const MPINcHyperslab& slab
            ) const
    {
        ContainerType data;
        get_var( name, slab, data);
        return data;
    }
 
    template<class T, std::enable_if_t< dg::is_scalar_v<T>, bool> = true>
    T get_var_as( std::string name, const std::vector<size_t>& start = {}) const
    {
        T var;
        get_var( name, start, var);
        return var;
    }
 
    bool var_is_defined( std::string name) const
    {
        return mpi_invoke( &SerialNcFile::var_is_defined, m_file, name);
    }
 
    nc_type get_var_type(std::string name) const
    {
        return mpi_invoke( &SerialNcFile::get_var_type, m_file, name);
    }
 
    std::vector<std::string> get_var_dims(std::string name) const
    {
        return mpi_invoke( &SerialNcFile::get_var_dims, m_file, name);
    }
 
    std::list<std::string> get_var_names() const
    {
        return mpi_invoke( &SerialNcFile::get_var_names, m_file);
    }
 
    private:
    template<class ContainerType>
    void set_comm( MPI_Vector<ContainerType>& x, MPI_Comm comm, dg::MPIVectorTag) const
    {
        x.set_communicator(comm);
    }
    template<class ContainerType>
    void set_comm( ContainerType&, MPI_Comm, dg::AnyVectorTag) const
    {
        return;
    }
    template<class ContainerType>
    const ContainerType& get_ref( const MPI_Vector<ContainerType>& x, dg::MPIVectorTag) const
    {
        return x.data();
    }
    template<class ContainerType>
    const ContainerType& get_ref( const ContainerType& x, dg::AnyVectorTag) const
    {
        return x;
    }
    template<class ContainerType>
    ContainerType& get_ref( MPI_Vector<ContainerType>& x, dg::MPIVectorTag) const
    {
        return x.data();
    }
    template<class ContainerType>
    ContainerType& get_ref( ContainerType& x, dg::AnyVectorTag) const
    {
        return x;
    }
 
    template<class T>
    void mpi_bcast( T& data) const
    {
        MPI_Bcast( &data, 1, dg::getMPIDataType<T>(), 0, m_comm);
    }
    void mpi_bcast( std::string& data) const
    {
        size_t len = data.size();
        MPI_Bcast( &len, 1, dg::getMPIDataType<size_t>(), 0, m_comm);
        data.resize( len, 'x');
        MPI_Bcast( &data[0], len, MPI_CHAR, 0, m_comm);
    }
    void mpi_bcast( std::filesystem::path& data) const
    {
        std::string name = data.generic_string();
        mpi_bcast( name);
        data = name;
    }
 
    template<class T>
    void mpi_bcast( std::vector<T>& data) const
    {
        size_t len = data.size();
        MPI_Bcast( &len, 1, dg::getMPIDataType<size_t>(), 0, m_comm);
        data.resize( len);
        for( unsigned u=0; u<len; u++)
        {
            if constexpr ( std::is_same_v<T, bool>)
            {
                bool b = data[u];
                mpi_bcast( b);
                data[u] = b;
            }
            else
                mpi_bcast( data[u]);
        }
    }
    template<class T>
    void mpi_bcast( std::list<T>& data) const
    {
        size_t len = data.size();
        MPI_Bcast( &len, 1, dg::getMPIDataType<size_t>(), 0, m_comm);
        data.resize( len);
        for( auto it = data.begin(); it != data.end(); it++)
            mpi_bcast( *it);
    }
    template<class K, class T>
    void mpi_bcast( std::map<K,T>& data) const
    {
        size_t len = data.size();
        MPI_Bcast( &len, 1, dg::getMPIDataType<size_t>(), 0, m_comm);
        std::vector<K> keys;
        std::vector<T> values;
        if( m_rank0)
        {
            for ( const auto& pair : data)
            {
                keys.push_back( pair.first);
                values.push_back( pair.second);
            }
        }
        mpi_bcast( keys);
        mpi_bcast( values);
        if( not m_rank0)
        {
            for( unsigned u=0; u<len; u++)
                data[keys[u]] = values[u];
        }
    }
    // Adapted from
    // https://stackoverflow.com/questions/60564132/default-constructing-an-stdvariant-from-index
    template <class Variant, std::size_t I = 0>
    Variant variant_from_index(std::size_t index) const
    {
        // I increases so we need to stop
        if constexpr( I >= std::variant_size_v<Variant>)
            return {};
        else
            return index == 0
                ? Variant{std::in_place_index<I>}
                : variant_from_index<Variant, I + 1>(index - 1);
    }
    void mpi_bcast( nc_att_t& data) const
    {
        // We need to make sure that data holds the correct type on non-rank0
        size_t idx = data.index();
        MPI_Bcast( &idx, 1, dg::getMPIDataType<size_t>(), 0, m_comm);
        if( not m_rank0)
            data = variant_from_index<nc_att_t>( idx);
        // ... so that visit calls the correct bcast overload
        std::visit( [this]( auto&& arg) { mpi_bcast(arg);}, data);
    }
 
    // Here we handle errors, such that if they occur on rank0 all ranks will
    // throw
    template<class F, class ... Args>
    std::invoke_result_t<F, Args...> mpi_invoke( F&& f, Args&& ...args) const
    {
        using R = std::invoke_result_t<F, Args...>;
        if ( m_readonly)
        {
            return std::invoke(std::forward<F>(f), std::forward<Args>(args)...);
        }
        int err = NC_NOERR;
        R r;
        if( m_rank0)
        {
            try
            {
                r = std::invoke(std::forward<F>(f), std::forward<Args>(args)...);
            }
            catch( const NC_Error& e) { err = e.error();}
        }
        mpi_bcast( err);
        if( err)
            throw NC_Error( err);
        mpi_bcast( r);
        return r;
    }
    template<class F, class ... Args>
    void mpi_invoke_void( F&& f, Args&& ... args) const
    {
        if ( m_readonly)
        {
            std::invoke(std::forward<F>(f), std::forward<Args>(args)...);
            return;
        }
        int err = NC_NOERR;
        if( m_rank0)
        {
            try
            {
                std::invoke(std::forward<F>(f), std::forward<Args>(args)...);
            }
            catch(const NC_Error& e) { err = e.error(); }
        }
        mpi_bcast( err); // this makes all calls blocking if not readonly!
        if( err)
            throw NC_Error( err);
    }
 
 
 
    bool m_rank0;
    bool m_readonly;
    MPI_Comm m_comm;
    SerialNcFile m_file;
    // Buffer for device to host transfer, and dg::assign
    mutable dg::detail::AnyVector<thrust::host_vector> m_buffer, m_receive;
};
 
using NcFile = MPINcFile;
 
}// namespace file
}// namespace dg