operators.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2026 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <complex>
#ifdef NF_WITH_MPI_SUPPORT
#include <mpi.h>
#endif
#include <type_traits>
#include <vector>
 
#include "NeoN/core/error.hpp"
#include "NeoN/core/primitives/vec3.hpp"
 
namespace NeoN
{
 
#ifdef NF_WITH_MPI_SUPPORT
 
namespace mpi
{
enum class ReduceOp
{
    Max,
    Min,
    Sum,
    Prod,
    Land,
    Band,
    Lor,
    Bor,
    Maxloc,
    Minloc
};
 
constexpr MPI_Op getOp(const ReduceOp op)
{
    switch (op)
    {
    case ReduceOp::Max:
        return MPI_MAX;
    case ReduceOp::Min:
        return MPI_MIN;
    case ReduceOp::Sum:
        return MPI_SUM;
    case ReduceOp::Prod:
        return MPI_PROD;
    case ReduceOp::Land:
        return MPI_LAND;
    case ReduceOp::Band:
        return MPI_BAND;
    case ReduceOp::Lor:
        return MPI_LOR;
    case ReduceOp::Bor:
        return MPI_BOR;
    case ReduceOp::Maxloc:
        return MPI_MAXLOC;
    case ReduceOp::Minloc:
        return MPI_MINLOC;
    default:
        NF_ERROR_EXIT("Invalid MPI reduce operation requested.");
        return MPI_LOR; // This is to suppress the warning
    }
}
 
template<typename valueType>
constexpr MPI_Datatype getType()
{
    if constexpr (std::is_same_v<valueType, char>) return MPI_CHAR;
    else if constexpr (std::is_same_v<valueType, wchar_t>)
        return MPI_WCHAR;
    else if constexpr (std::is_same_v<valueType, short>)
        return MPI_SHORT;
    else if constexpr (std::is_same_v<valueType, int>)
        return MPI_INT;
    else if constexpr (std::is_same_v<valueType, long>)
        return MPI_LONG;
    else if constexpr (std::is_same_v<valueType, long long>)
        return MPI_LONG_LONG;
    else if constexpr (std::is_same_v<valueType, unsigned short>)
        return MPI_UNSIGNED_SHORT;
    else if constexpr (std::is_same_v<valueType, unsigned>)
        return MPI_UNSIGNED;
    else if constexpr (std::is_same_v<valueType, unsigned long>)
        return MPI_UNSIGNED_LONG;
    else if constexpr (std::is_same_v<valueType, unsigned long long>)
        return MPI_UNSIGNED_LONG_LONG;
    else if constexpr (std::is_same_v<valueType, float>)
        return MPI_FLOAT;
    else if constexpr (std::is_same_v<valueType, double>)
        return MPI_DOUBLE;
    else if constexpr (std::is_same_v<valueType, long double>)
        return MPI_LONG_DOUBLE;
    else if constexpr (std::is_same_v<valueType, bool>)
        return MPI_CXX_BOOL;
    else if constexpr (std::is_same_v<valueType, std::complex<float>>)
        return MPI_CXX_FLOAT_COMPLEX;
    else if constexpr (std::is_same_v<valueType, std::complex<double>>)
        return MPI_CXX_DOUBLE_COMPLEX;
    else if constexpr (std::is_same_v<valueType, std::complex<long double>>)
        return MPI_CXX_LONG_DOUBLE_COMPLEX;
    else
        NF_ERROR_EXIT("Invalid MPI datatype requested.");
    return MPI_CHAR; // This is to suppress the warning
}
 
template<typename valueType>
void allReduce(valueType& value, const ReduceOp op, MPI_Comm comm)
{
    MPI_Allreduce(
        MPI_IN_PLACE, reinterpret_cast<void*>(&value), 1, getType<valueType>(), getOp(op), comm
    );
}
 
template<>
inline void allReduce(Vec3& vector, const ReduceOp op, MPI_Comm comm)
{
    MPI_Allreduce(
        MPI_IN_PLACE,
        reinterpret_cast<void*>(vector.data()),
        static_cast<mpi_label_t>(vector.size()),
        getType<scalar>(),
        getOp(op),
        comm
    );
}
 
template<typename valueType>
void isend(
    const valueType* buffer,
    const mpi_label_t size,
    mpi_label_t rankReceive,
    mpi_label_t tag,
    MPI_Comm comm,
    MPI_Request* request
)
{
    mpi_label_t err =
        MPI_Isend(buffer, size, getType<valueType>(), rankReceive, tag, comm, request);
    NF_DEBUG_ASSERT(err == MPI_SUCCESS, "MPI_Isend failed.");
}
 
template<typename valueType>
void irecv(
    valueType* buffer,
    const mpi_label_t size,
    mpi_label_t rankSend,
    mpi_label_t tag,
    MPI_Comm comm,
    MPI_Request* request
)
{
    mpi_label_t err = MPI_Irecv(buffer, size, getType<valueType>(), rankSend, tag, comm, request);
    NF_DEBUG_ASSERT(err == MPI_SUCCESS, "MPI_Irecv failed.");
}
 
inline bool test(MPI_Request* request)
{
    mpi_label_t flag;
    mpi_label_t err = MPI_Test(request, &flag, MPI_STATUS_IGNORE);
    NF_DEBUG_ASSERT(err == MPI_SUCCESS, "MPI_Test failed.");
    return static_cast<bool>(flag);
}
 
inline void waitAll(MPI_Request* requests, const mpi_label_t count)
{
    if (count == 0) return;
    mpi_label_t err = MPI_Waitall(count, requests, MPI_STATUSES_IGNORE);
    NF_DEBUG_ASSERT(err == MPI_SUCCESS, "MPI_Waitall failed.");
}
 
inline void waitAll(std::vector<MPI_Request>& requests)
{
    waitAll(requests.data(), static_cast<mpi_label_t>(requests.size()));
}
 
template<typename valueType>
void allToAll(
    const valueType* sendBuf,
    mpi_label_t sendCount,
    valueType* recvBuf,
    mpi_label_t recvCount,
    MPI_Comm comm
)
{
    mpi_label_t err = MPI_Alltoall(
        sendBuf, sendCount, getType<valueType>(), recvBuf, recvCount, getType<valueType>(), comm
    );
    NF_DEBUG_ASSERT(err == MPI_SUCCESS, "MPI_Alltoall failed.");
}
 
template<typename valueType>
void allToAllV(
    const valueType* sendBuf,
    const mpi_label_t* sendCounts,
    const mpi_label_t* sendDispls,
    valueType* recvBuf,
    const mpi_label_t* recvCounts,
    const mpi_label_t* recvDispls,
    MPI_Comm comm
)
{
    mpi_label_t err = MPI_Alltoallv(
        sendBuf,
        sendCounts,
        sendDispls,
        getType<valueType>(),
        recvBuf,
        recvCounts,
        recvDispls,
        getType<valueType>(),
        comm
    );
    NF_DEBUG_ASSERT(err == MPI_SUCCESS, "MPI_Alltoallv failed.");
}
 
} // namespace mpi
 
#endif
 
}