halfDuplexCommBuffer.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2025 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <string>
#include <typeindex>
#include <vector>
 
#include "NeoN/core/error.hpp"
#include "NeoN/core/mpi/environment.hpp"
#include "NeoN/core/mpi/operators.hpp"
#include "NeoN/core/view.hpp"
 
namespace NeoN
{
 
#ifdef NF_WITH_MPI_SUPPORT
 
namespace mpi
{
 
inline int bufferHash(const std::string& str)
{
    // There is also an MPI environment value for that, but somehow it doesn't work using that
    // Also reserve 10 tags for other uses
    constexpr int maxTagValue = 32767 - 10;
    std::size_t tag = std::hash<std::string> {}(str);
    tag &= 0x7FFFFFFF; // turn 'int' signed bit to 0, MPI does not like negative tags.
    return (static_cast<int>(tag) % maxTagValue) + 10;
}
 
class HalfDuplexCommBuffer
{
 
public:
 
    HalfDuplexCommBuffer() = default;
 
    ~HalfDuplexCommBuffer() = default;
 
    HalfDuplexCommBuffer(MPIEnvironment mpiEnviron, std::vector<size_t> rankCommSize)
        : mpiEnviron_(mpiEnviron)
    {
        setCommRankSize<char>(rankCommSize);
    }
 
    inline void setMPIEnvironment(MPIEnvironment mpiEnviron) { mpiEnviron_ = mpiEnviron; }
 
    inline bool isCommInit() const { return tag_ != -1; }
 
    template<typename valueType>
    void setCommRankSize(std::vector<size_t> rankCommSize)
    {
        NF_DEBUG_ASSERT(
            !isCommInit(), "Communication buffer was initialised by name: " << commName_ << "."
        );
        NF_DEBUG_ASSERT(
            rankCommSize.size() == mpiEnviron_.sizeRank(),
            "Rank size mismatch. " << rankCommSize.size() << " vs. " << mpiEnviron_.sizeRank()
        );
        typeSize_ = sizeof(valueType);
        rankOffset_.resize(rankCommSize.size() + 1);
        request_.resize(rankCommSize.size(), MPI_REQUEST_NULL);
        updateDataSize([&](const size_t rank) { return rankCommSize[rank]; }, sizeof(valueType));
    }
 
    template<typename valueType>
    void initComm(std::string commName)
    {
        NF_DEBUG_ASSERT(
            !isCommInit(), "Communication buffer was initialised by name: " << commName_ << "."
        );
        setType<valueType>();
        commName_ = commName;
        tag_ = bufferHash(commName);
    }
 
    inline const std::string& getCommName() const { return commName_; }
 
    bool isComplete();
 
    void send();
 
    void receive();
 
    void waitComplete();
 
    void finaliseComm();
 
    template<typename valueType>
    View<valueType> get(const size_t rank)
    {
        NF_DEBUG_ASSERT(isCommInit(), "Communication buffer is not initialised.");
        NF_DEBUG_ASSERT(typeSize_ == sizeof(valueType), "Data type (size) mismatch.");
        return View<valueType>(
            reinterpret_cast<valueType*>(rankBuffer_.data() + rankOffset_[rank]),
            (rankOffset_[rank + 1] - rankOffset_[rank]) / sizeof(valueType)
        );
    }
 
    template<typename valueType>
    View<const valueType> get(const size_t rank) const
    {
        NF_DEBUG_ASSERT(isCommInit(), "Communication buffer is not initialised.");
        NF_DEBUG_ASSERT(typeSize_ == sizeof(valueType), "Data type (size) mismatch.");
        return View<const valueType>(
            reinterpret_cast<const valueType*>(rankBuffer_.data() + rankOffset_[rank]),
            (rankOffset_[rank + 1] - rankOffset_[rank]) / sizeof(valueType)
        );
    }
 
private:
 
    int tag_ {-1};                        /*< The tag for the communication. */
    std::string commName_ {"unassigned"}; /*< The name of the communication. */
    std::size_t typeSize_ {sizeof(char)}; /*< The data type currently stored in the buffer. */
    MPIEnvironment mpiEnviron_;           /*< The MPI environment. */
    std::vector<MPI_Request> request_;    /*< The MPI request for communication with each rank. */
    std::vector<char> rankBuffer_;        /*< The buffer data for all ranks. Never shrinks. */
    std::vector<std::size_t>
        rankOffset_; /*< The offset (in bytes) for a rank data in the buffer. */
 
    template<typename valueType>
    void setType()
    {
        NF_DEBUG_ASSERT(
            !isCommInit(), "Communication buffer was initialised by name: " << commName_ << "."
        );
        if (0 == (typeSize_ - sizeof(valueType))) return;
        updateDataSize(
            [rankOffset = rankOffset_, typeSize = typeSize_](const size_t rank)
            { return (rankOffset[rank + 1] - rankOffset[rank]) / typeSize; },
            sizeof(valueType)
        );
 
        typeSize_ = sizeof(valueType);
    }
 
    template<typename func>
    void updateDataSize(func rankSize, std::size_t newSize)
    {
        std::size_t dataSize = 0;
        for (size_t rank = 0; rank < mpiEnviron_.sizeRank(); ++rank)
        {
            rankOffset_[rank] = dataSize;
            dataSize += rankSize(rank) * newSize;
        }
        rankOffset_.back() = dataSize;
        if (rankBuffer_.size() < dataSize) rankBuffer_.resize(dataSize); // we never size down.
    }
};
 
} // namespace mpi
 
#endif
 
}