impl_op.hpp

#ifndef __IMPL_MPI_OP_HPP__
#define __IMPL_MPI_OP_HPP__
 
#include "common/common.hpp"
#include <mpi_w/message_t.hpp>
#include <mpi_w/mpi_types.hpp>
 
#include <common/execinfo.hpp>
#include <cstddef>
#include <limits>
#include <math.h>
#include <mpi.h>
#include <optional>
#include <span>
#include <stdexcept>
#include <vector>

namespace WrapMPI
{
 
  // SENDING

  template <POD_t DataType>
  [[nodiscard]] static int _send_unsafe(DataType* buf,
                                        size_t buf_size,
                                        size_t dest,
                                        size_t tag = 0) noexcept;

  template <POD_t DataType>
  int send(DataType data, size_t dest, size_t tag = 0);

  template <POD_t DataType>
  int send_v(std::span<const DataType> data,
             size_t dest,
             size_t tag = 0,
             bool send_size = true) noexcept;
 
  // RECEIVE

  template <POD_t DataType>
  std::optional<DataType>
  recv(size_t src, MPI_Status* status = nullptr, size_t tag = 0) noexcept;

  template <POD_t DataType>
  int recv_span(std::span<DataType> buf,
                size_t src,
                MPI_Status* status = nullptr,
                size_t tag = 0) noexcept;

  template <POD_t DataType>
  DataType try_recv(size_t src, MPI_Status* status = nullptr, size_t tag = 0);

  template <POD_t T>
  std::optional<std::vector<T>>
  recv_v(size_t source, MPI_Status* status = nullptr, size_t tag = 0) noexcept;

  template <POD_t T>
  std::vector<T>
  try_recv_v(size_t src, MPI_Status* status = nullptr, size_t tag = 0);
 
  // BROADCASTING

  template <POD_t T>
  [[nodiscard]] int _broadcast_unsafe(T* data, size_t _size, size_t root);

  template <POD_t DataType> int broadcast(DataType& data, size_t root) noexcept;

  template <POD_t T> int broadcast_span(std::span<T> data, size_t root);

  template <POD_t... Args>
  void host_dispatch(const ExecInfo& info, SIGNALS sign, Args&&... args);

  template <POD_t T>
  std::vector<T>
  _gather_unsafe(T* src_data, size_t size, size_t n_rank, size_t root = 0);
 
  template <POD_t T>
  int _gather_unsafe_to_buffer(T* dest,
                               T* src_data,
                               size_t size,
                               size_t root = 0) noexcept;

  template <POD_t T>
  std::vector<T>
  gather(std::span<T> local_data, size_t n_rank, size_t root = 0);
 
  template <POD_t T>
  void gather_span(std::span<T> dest,
                   std::span<const T> local_data,
                   size_t root = 0);
 
  template <POD_t T>
  std::vector<T>
  gather(std::span<const T> local_data, size_t n_rank, size_t root = 0);

  template <NumberType T> T gather_reduce(T data, size_t root = 0);
 
  template <NumberType T> T all_reduce(T data)
  {
    T global_sum{}; // global sum across all ran
    MPI_Allreduce(
        &data, &global_sum, 1, get_type<T>(), MPI_SUM, MPI_COMM_WORLD);
    return global_sum;
  }

  template <POD_t T>
  std::vector<T>
  gather_v(const std::vector<T>& local_data, size_t n_rank, size_t root = 0);
 
  //**
  // IMPL
  //**
  template <POD_t DataType>
  static int
  _send_unsafe(DataType* buf, size_t buf_size, size_t dest, size_t tag) noexcept
  {
    return MPI_Send(
        buf, buf_size, get_type<DataType>(), dest, tag, MPI_COMM_WORLD);
  }
 
  template <POD_t DataType>
  DataType try_recv(size_t src, MPI_Status* status, size_t tag)
  {
    auto opt_data = WrapMPI::recv<DataType>(src, status, tag);
    if (!opt_data.has_value())
    {
      WrapMPI::critical_error();
      exit(-1); // critical_error should exit before reaching this statement
    }
    else
    {
      return opt_data.value();
    }
  }
 
  template <POD_t T>
  std::optional<std::vector<T>>
  recv_v(size_t source, MPI_Status* status, size_t tag) noexcept
  {
    std::vector<T> buf;
 
    // Receive the size of the vector
    auto opt_size = recv<size_t>(source, status, tag);
    if (!opt_size.has_value())
    {
      return std::nullopt; // Return early if size reception fails
    }
    size_t buf_size = opt_size.value();
 
    // Resize the buffer
    buf.resize(buf_size);
 
    MPI_Datatype datatype = get_type<T>();
 
    // Receive the vector data
    int recv_status = MPI_Recv(buf.data(),
                               static_cast<int>(buf_size),
                               datatype,
                               source,
                               tag,
                               MPI_COMM_WORLD,
                               status);
 
    if (recv_status != MPI_SUCCESS)
    {
      return std::nullopt; // Return early if MPI_Recv fails
    }
 
    return buf; // Return the received vector
  }
 
  template <POD_t DataType>
  int recv_span(std::span<DataType> buf,
                size_t src,
                MPI_Status* status,
                size_t tag) noexcept
  {
    return MPI_Recv(buf.data(),
                    buf.size(),
                    get_type<DataType>(),
                    src,
                    tag,
                    MPI_COMM_WORLD,
                    status);
  }
 
  template <POD_t T>
  std::vector<T> try_recv_v(size_t src, MPI_Status* status, size_t tag)
  {
    auto opt_data = WrapMPI::recv_v<T>(src, status, tag);
    if (!opt_data.has_value())
    {
      WrapMPI::critical_error();
    }
    else
    {
      return opt_data.value();
    }
  }
 
  template <POD_t DataType> int broadcast(DataType& data, size_t root) noexcept
  {
    return MPI_Bcast(
        &data, 1, get_type<DataType>(), static_cast<int>(root), MPI_COMM_WORLD);
  }
 
  // template <> int broadcast(size_t &data, size_t root)
  // {
  //   return MPI_Bcast(&data,
  //                    sizeof(size_t),
  //                    MPI_UNSIGNED_LONG,
  //                    static_cast<int>(root),
  //                    MPI_COMM_WORLD);
  // }
 
  template <POD_t T>
  int broadcast(std::vector<T>& data, size_t root, size_t current_rank)
  {
 
    size_t data_size = 0;
    if (current_rank == root)
    {
      data_size = data.size();
    }
    broadcast(data_size, root);
    if (current_rank != root)
    {
      data.resize(data_size);
    }
 
    return MPI_Bcast(data.data(),
                     data_size,
                     get_type<T>(),
                     static_cast<int>(root),
                     MPI_COMM_WORLD);
  }
 
  template <POD_t T> int _broadcast_unsafe(T* data, size_t _size, size_t root)
  {
 
    if (data == nullptr)
    {
      throw std::invalid_argument("Data pointer is null");
    }
    if (_size == 0 || _size > std::numeric_limits<size_t>::max())
    {
      throw std::invalid_argument("Error size");
    }
 
    int comm_size = 0;
    MPI_Comm_size(MPI_COMM_WORLD, &comm_size);
    if (root >= static_cast<size_t>(comm_size))
    {
      throw std::invalid_argument("Root process rank is out of range");
    }
 
    // Broadcast operation
    return MPI_Bcast(
        data, _size, get_type<T>(), static_cast<int>(root), MPI_COMM_WORLD);
  }
 
  template <POD_t T> int broadcast_span(std::span<T> data, size_t root)
  {
    return _broadcast_unsafe(data.data(), data.size(), root);
  }
 
  template <POD_t T>
  std::vector<T>
  _gather_unsafe(T* src_data, size_t size, size_t n_rank, size_t root)
  {
    int src_size = static_cast<int>(size);
    std::vector<T> total_data(src_size * n_rank);
    T* dest_data = total_data.data();
    // auto mpi_type = get_type<T>();
 
    // int gather_result = MPI_Gather(
    //     src_data, src_size, mpi_type, dest_data, src_size, mpi_type, root,
    //     MPI_COMM_WORLD);
    // if (gather_result != MPI_SUCCESS)
    // {
    //   throw std::runtime_error("MPI_Gather failed");
    // }
    if (_gather_unsafe_to_buffer(dest_data, src_data, size, root) !=
        MPI_SUCCESS)
    {
      throw std::runtime_error("MPI_Gather failed");
    }
 
    return total_data;
  }
 
  template <POD_t T>
  int _gather_unsafe_to_buffer(T* const dest,
                               T* src_data,
                               size_t size,
                               size_t root) noexcept
  {
    int src_size = static_cast<int>(size);
    auto mpi_type = get_type<T>();
 
    return MPI_Gather(src_data,
                      src_size,
                      mpi_type,
                      dest,
                      src_size,
                      mpi_type,
                      root,
                      MPI_COMM_WORLD);
  }
 
  template <POD_t T>
  void
  gather_span(std::span<T> dest, std::span<const T> local_data, size_t root)
  {
    T* dest_data = dest.data();
#ifndef NDEBUG
    int size{};
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    assert(dest.size() == local_data.size() * size);
#endif
    _gather_unsafe_to_buffer(
        dest_data, const_cast<T*>(local_data.data()), local_data.size(), root);
  }
 
  template <POD_t T>
  std::vector<T> gather(std::span<T> local_data, size_t n_rank, size_t root)
  {
    return _gather_unsafe(local_data.data(), local_data.size(), n_rank, root);
  }
  // FIXME
  template <POD_t T>
  std::vector<T>
  gather(std::span<const T> local_data, size_t n_rank, size_t root)
  {
    return _gather_unsafe(
        const_cast<T*>(local_data.data()), local_data.size(), n_rank, root);
  }
 
  template <NumberType T> T gather_reduce(T data, size_t root)
  {
 
    T result{};
 
    MPI_Reduce(&data,
               &result,
               1,
               WrapMPI::get_type<T>(),
               MPI_SUM,
               root,
               MPI_COMM_WORLD);
 
    return result;
  }
 
  template <POD_t T>
  std::vector<T>
  gather_v(const std::vector<T>& local_data, size_t n_rank, size_t root)
  {
    return _gather_unsafe(local_data.data(), local_data.size(), n_rank, root);
  }
 
  template <POD_t... Args>
  void host_dispatch(const ExecInfo& info, SIGNALS sign, Args&&... args)
  {
 
    for (int j = 1; j < static_cast<int>(info.n_rank); ++j)
    {
      if (sign != WrapMPI::SIGNALS::NOP)
      {
        MPI_Send(&sign, sizeof(sign), MPI_CHAR, j, 0, MPI_COMM_WORLD);
      }
 
      (
          [&]<POD_t T>(T& arg)
          {
            size_t s = 0;
            void* buf = nullptr;
            if constexpr (std::is_same_v<std::decay_t<T>, std::span<double>>)
            {
              s = arg.size();
              buf = arg.data();
              MPI_Send(&s, 1, MPI_UNSIGNED_LONG, j, 0, MPI_COMM_WORLD);
            }
            else
            {
              s = sizeof(T);
              buf = &arg;
            }
 
            MPI_Send(
                buf, static_cast<int>(s), MPI_DOUBLE, j, 0, MPI_COMM_WORLD);
          }(std::forward<Args>(args)),
          ...);
    }
  }
 
  template <POD_t DataType> int send(DataType data, size_t dest, size_t tag)
  {
    return _send_unsafe<DataType>(&data, 1, dest, tag);
  }
 
  template <POD_t DataType>
  int send_v(std::span<const DataType> data,
             size_t dest,
             size_t tag,
             bool send_size) noexcept
  {
    int send_status = MPI_SUCCESS;
 
    if (send_size)
    {
      send_status = send<size_t>(data.size(), dest, tag);
    }
 
    if (send_status == MPI_SUCCESS)
    {
      send_status = _send_unsafe(data.data(), data.size(), dest, tag);
    }
 
    return send_status;
  }
 
  template <POD_t DataType>
  std::optional<DataType>
  recv(size_t src, MPI_Status* status, size_t tag) noexcept
  {
    DataType buf;
 
    int recv_status = MPI_Recv(
        &buf, sizeof(DataType), MPI_BYTE, src, tag, MPI_COMM_WORLD, status);
    if (recv_status != MPI_SUCCESS)
    {
      return std::nullopt;
    }
    return buf;
  }
 
} // namespace WrapMPI
 
#endif //__IMPL_MPI_OP_HPP__