environment.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2025 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#ifdef NF_WITH_MPI_SUPPORT
#include <mpi.h>
#endif
 
#include "NeoN/core/error.hpp"
#include "NeoN/core/info.hpp"
 
 
namespace NeoN
{
 
#ifdef NF_WITH_MPI_SUPPORT
 
namespace mpi
{
 
struct MPIInit
{
    MPIInit(int argc, char** argv)
    {
#ifdef NF_REQUIRE_MPI_THREAD_SUPPORT
        int provided;
        MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided);
        NF_ASSERT(
            provided == MPI_THREAD_MULTIPLE, "The MPI library does not have full thread support"
        );
#else
        MPI_Init(&argc, &argv);
#endif
    }
 
    ~MPIInit() { MPI_Finalize(); }
};
 
 
class MPIEnvironment
{
public:
 
    MPIEnvironment(MPI_Comm commGroup = MPI_COMM_WORLD) : communicator(commGroup)
    {
        updateRankData();
    }
 
    ~MPIEnvironment() = default;
 
    size_t sizeRank() const { return static_cast<size_t>(mpi_size); }
 
    size_t rank() const { return static_cast<size_t>(mpi_rank); }
 
    MPI_Comm comm() const { return communicator; }
 
private:
 
    MPI_Comm communicator {MPI_COMM_NULL}; // MPI communicator
    int mpi_rank {-1};                     // Index of this rank
    int mpi_size {-1};                     // Number of ranks in this communicator group.
 
    void updateRankData()
    {
        NF_ASSERT(communicator != MPI_COMM_NULL, "Invalid communicator, is null.");
        MPI_Comm_rank(communicator, &mpi_rank);
        MPI_Comm_size(communicator, &mpi_size);
    }
};
 
} // namespace mpi
 
#endif
 
} // namespace NeoN