objectComms.h

// Copyright 2009-2025 NTESS. Under the terms
// of Contract DE-NA0003525 with NTESS, the U.S.
// Government retains certain rights in this software.
//
// Copyright (c) 2009-2025, NTESS
// All rights reserved.
//
// This file is part of the SST software package. For license
// information, see the LICENSE file in the top level directory of the
// distribution.

#ifndef SST_CORE_OBJECTCOMMS_H
#define SST_CORE_OBJECTCOMMS_H

#include "sst/core/objectSerialization.h"
#include "sst/core/sst_mpi.h"

#include <memory>
#include <typeinfo>

namespace SST::Comms {

#ifdef SST_CONFIG_HAVE_MPI
template <typename dataType>
void
broadcast(dataType& data, int root)
{
    int rank = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    if ( root == rank ) {
        // Serialize the data
        std::vector<char> buffer = Comms::serialize(data);

        // Now broadcast the size of the data
        int size = buffer.size();
        MPI_Bcast(&size, 1, MPI_INT, root, MPI_COMM_WORLD);

        // Now broadcast the data
        MPI_Bcast(buffer.data(), buffer.size(), MPI_BYTE, root, MPI_COMM_WORLD);
    }
    else {
        // Get the size of the broadcast
        int size = 0;
        MPI_Bcast(&size, 1, MPI_INT, root, MPI_COMM_WORLD);

        // Now get the data
        auto buffer = std::unique_ptr<char[]>(new char[size]);
        MPI_Bcast(buffer.get(), size, MPI_BYTE, root, MPI_COMM_WORLD);

        // Now deserialize data
        Comms::deserialize(buffer.get(), size, data);
    }
}

template <typename dataType>
void
send(int dest, int tag, dataType& data)
{
    // Serialize the data
    std::vector<char> buffer = Comms::serialize<dataType>(data);

    // Now send the data.  Send size first, then payload
    // std::cout<< sizeof(buffer.size()) << std::endl;
    int64_t size = buffer.size();
    MPI_Send(&size, 1, MPI_INT64_T, dest, tag, MPI_COMM_WORLD);

    int32_t fragment_size = 1000000000;
    int64_t offset        = 0;

    while ( size >= fragment_size ) {
        MPI_Send(buffer.data() + offset, fragment_size, MPI_BYTE, dest, tag, MPI_COMM_WORLD);
        size -= fragment_size;
        offset += fragment_size;
    }
    MPI_Send(buffer.data() + offset, size, MPI_BYTE, dest, tag, MPI_COMM_WORLD);
}

template <typename dataType>
void
recv(int src, int tag, dataType& data)
{
    // Get the size of the broadcast
    int64_t    size = 0;
    MPI_Status status;
    MPI_Recv(&size, 1, MPI_INT64_T, src, tag, MPI_COMM_WORLD, &status);

    // Now get the data
    auto    buffer        = std::unique_ptr<char[]>(new char[size]);
    int64_t offset        = 0;
    int32_t fragment_size = 1000000000;
    int64_t rem_size      = size;

    while ( rem_size >= fragment_size ) {
        MPI_Recv(buffer.get() + offset, fragment_size, MPI_BYTE, src, tag, MPI_COMM_WORLD, &status);
        rem_size -= fragment_size;
        offset += fragment_size;
    }
    MPI_Recv(buffer.get() + offset, rem_size, MPI_BYTE, src, tag, MPI_COMM_WORLD, &status);

    // Now deserialize data
    Comms::deserialize(buffer.get(), size, data);
}

template <typename dataType>
void
all_gather(dataType& data, std::vector<dataType>& out_data)
{
    int rank = 0, world = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &world);

    // Serialize the data
    std::vector<char> buffer = Comms::serialize(data);

    size_t sendSize = buffer.size();

    auto allSizes      = std::make_unique<int[]>(world);
    auto displacements = std::make_unique<int[]>(world);

    memset(allSizes.get(), '\0', world * sizeof(int));
    memset(displacements.get(), '\0', world * sizeof(int));

    MPI_Allgather(&sendSize, sizeof(int), MPI_BYTE, allSizes.get(), sizeof(int), MPI_BYTE, MPI_COMM_WORLD);

    int totalBuf = 0;
    for ( int i = 0; i < world; i++ ) {
        totalBuf += allSizes[i];
        if ( i > 0 ) displacements[i] = displacements[i - 1] + allSizes[i - 1];
    }

    auto bigBuff = std::unique_ptr<char[]>(new char[totalBuf]);

    MPI_Allgatherv(buffer.data(), buffer.size(), MPI_BYTE, bigBuff.get(), allSizes.get(), displacements.get(), MPI_BYTE,
        MPI_COMM_WORLD);

    out_data.resize(world);
    for ( int i = 0; i < world; i++ ) {
        auto* bbuf = bigBuff.get();
        Comms::deserialize(&bbuf[displacements[i]], allSizes[i], out_data[i]);
    }
}

#endif

} // namespace SST::Comms

#endif // SST_CORE_OBJECTCOMMS_H