environment.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2026 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#ifdef NF_WITH_MPI_SUPPORT
#include <mpi.h>
#endif
 
#include <cstdlib>
 
#include "NeoN/core/error.hpp"
#include "NeoN/core/info.hpp"
 
 
namespace NeoN
{
 
#ifdef NF_WITH_MPI_SUPPORT
 
namespace mpi
{
 
struct Init
{
    Init(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
    }
 
    ~Init() { MPI_Finalize(); }
};
 
 
class Environment
{
public:
 
    Environment(MPI_Comm commGroup = MPI_COMM_WORLD) : communicator(commGroup)
    {
        MPI_Initialized(&mpiInitialized);
        updateRankData();
        if (std::getenv("NEON_FORCE_HOST_BUFFER") != nullptr) gpuAwareMpi_ = false;
    }
 
    ~Environment() = default;
 
    bool isInitialized() const { return mpiInitialized == 1; }
 
    size_t sizeRank() const { return static_cast<size_t>(mpiSize); }
 
    size_t rank() const { return static_cast<size_t>(mpiRank); }
 
    MPI_Comm comm() const { return communicator; }
 
    bool gpuAwareMpi() const { return gpuAwareMpi_; }
 
    bool& gpuAwareMpi() { return gpuAwareMpi_; }
 
private:
 
    MPI_Comm communicator {MPI_COMM_NULL}; // MPI communicator
    int mpiInitialized {0};
    int mpiRank {-1}; // Index of this rank
    int mpiSize {-1}; // Number of ranks in this communicator group.
    bool gpuAwareMpi_ {true};
 
    void updateRankData()
    {
        if (mpiInitialized)
        {
            MPI_Comm_rank(communicator, &mpiRank);
            MPI_Comm_size(communicator, &mpiSize);
        }
    }
};
 
} // namespace mpi
 
#endif
 
} // namespace NeoN