communicator.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2026 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <vector>
#include <unordered_map>
#include <memory>
 
 
#include "NeoN/fields/field.hpp"
 
#ifdef NF_WITH_MPI_SUPPORT
#include "NeoN/core/mpi/fullDuplexCommBuffer.hpp"
#include "NeoN/core/mpi/operators.hpp"
#include "NeoN/core/mpi/environment.hpp"
#endif
 
namespace NeoN
{
 
#ifdef NF_WITH_MPI_SUPPORT
struct NodeCommMap
{
    label local_idx; 
};
 
using RankCommMap = std::vector<NodeCommMap>;
 
using CommMap = std::vector<RankCommMap>;
 
class Communicator
{
public:
 
    using bufferType = mpi::FullDuplexCommBuffer;
 
    Communicator() = default;
 
    ~Communicator() = default;
 
    Communicator(mpi::Environment mpiEnviron, CommMap rankSendMap, CommMap rankReceiveMap)
        : mpiEnviron_(mpiEnviron), sendMap_(rankSendMap), receiveMap_(rankReceiveMap)
    {
        NF_DEBUG_ASSERT(
            mpiEnviron_.sizeRank() == rankSendMap.size(),
            "Size of rankSendSize does not match MPI size."
        );
        NF_DEBUG_ASSERT(
            mpiEnviron_.sizeRank() == rankReceiveMap.size(),
            "Size of rankReceiveSize does not match MPI size."
        );
    };
 
    template<typename valueType>
    void startComm(Vector<valueType>& field, const std::string& commName)
    {
        NF_DEBUG_ASSERT(
            !CommBuffer_.contains(commName) || (!CommBuffer_[commName]),
            "There is already an ongoing communication for key " << commName << "."
        );
 
        CommBuffer_[commName] = findDuplexBuffer();
        if (!CommBuffer_[commName])
        {
            CommBuffer_[commName] = createNewDuplexBuffer();
        }
 
        CommBuffer_[commName]->initComm<valueType>(commName);
        for (size_t rank = 0; rank < mpiEnviron_.sizeRank(); ++rank)
        {
            auto rankBuffer = CommBuffer_[commName]->getSend<valueType>(rank);
            for (size_t data = 0; data < sendMap_[rank].size(); ++data)
                rankBuffer[data] =
                    field.view()[static_cast<localIdx>(sendMap_[rank][data].local_idx)];
        }
        CommBuffer_[commName]->startComm();
    }
 
    bool isComplete(std::string commName);
 
    template<typename valueType>
    void finaliseComm(Vector<valueType>& field, std::string commName)
    {
        NF_DEBUG_ASSERT(
            CommBuffer_.contains(commName) && CommBuffer_[commName],
            "No communication associated with key: " << commName
        );
 
        CommBuffer_[commName]->waitComplete();
        for (size_t rank = 0; rank < mpiEnviron_.sizeRank(); ++rank)
        {
            auto rankBuffer = CommBuffer_[commName]->getReceive<valueType>(rank);
            for (size_t data = 0; data < receiveMap_[rank].size(); ++data)
                field.view()[static_cast<localIdx>(receiveMap_[rank][data].local_idx)] =
                    rankBuffer[data];
        }
        CommBuffer_[commName]->finaliseComm();
        CommBuffer_[commName] = nullptr;
    }
 
private:
 
    mpi::Environment mpiEnviron_;    
    CommMap sendMap_;                
    CommMap receiveMap_;             
    std::vector<bufferType> buffers; 
    std::unordered_map<std::string, bufferType*>
        CommBuffer_; 
    bufferType* findDuplexBuffer();
 
    bufferType* createNewDuplexBuffer();
};
#endif
 
} // namespace NeoN