File mpiallreduce.h

File List > code_source > templat > include > TempLat > parallel > mpi > comm > exchange > mpiallreduce.h

Go to the documentation of this file

#ifndef TEMPLAT_PARALLEL_MPI_COMM_EXCHANGE_MPIALLREDUCE_H
#define TEMPLAT_PARALLEL_MPI_COMM_EXCHANGE_MPIALLREDUCE_H

/* This file is part of TempLat, available at https://cosmolattice.github.io/templat .
   Copyright 2021-2026 The TempLat authors, see AUTHORS.md.
   Released under the MIT license, see LICENSE.md. */

// File info: Main contributor(s): Wessel Valkenburg, Year: 2019

#include <vector>
#include <array>

#include "TempLat/parallel/devices/kokkos/kokkos.h"
#include "TempLat/util/exception.h"
#include "TempLat/util/log/saycomplete.h"
#include "TempLat/parallel/mpi/mpitypeconstants.h"
#include "TempLat/parallel/device.h"

namespace TempLat
{
  MakeException(MPIAllReduceException);

  class MPIAllReduce
  {
  public:
    // Put public methods here. These should change very little over time.
    MPIAllReduce(MPI_Comm comm) : mComm(comm) {}

#ifdef HAVE_MPI
    template <typename T>
      requires requires { MPITypeSelect<T>(); }
    T Allreduce(const T &value, MPI_Op operation, ptrdiff_t size = 1, int *error = nullptr)
    {
      T copyValue = value;
      int myerror = MPI_SUCCESS != MPI_Allreduce(MPI_IN_PLACE, &copyValue, size, MPITypeSelect<T>(), operation, mComm);
      if (error) {
        *error = myerror;
      } else {
        combineAllError(myerror);
      }
      return copyValue;
    }

    template <typename T>
    std::vector<T> &Allreduce(std::vector<T> *value, MPI_Op operation, ptrdiff_t size = 0, int *error = nullptr)
    {
      int myerror = MPI_SUCCESS !=
                    MPI_Allreduce(MPI_IN_PLACE, value->data(), value->size(), MPITypeSelect<T>(), operation, mComm);
      if (error) {
        *error = myerror;
      } else {
        combineAllError(myerror);
      }
      return *value;
    }

    template <typename View>
      requires requires(View view) {
        typename View::value_type;
        typename View::size_type;
        view.data();
        MPITypeSelect<typename View::value_type>();
      }
    View Allreduce(View value, MPI_Op operation, ptrdiff_t size = 0, int *error = nullptr)
    {
      using T = typename View::value_type;
      int myerror =
          MPI_SUCCESS != MPI_Allreduce(MPI_IN_PLACE, value.data(), value.size(), MPITypeSelect<T>(), operation, mComm);
      if (error) {
        *error = myerror;
      } else {
        combineAllError(myerror);
      }
      return value;
    }

    template <typename T, size_t N>
    std::vector<std::array<T, N>> &Allreduce(std::vector<std::array<T, N>> *value, MPI_Op operation, ptrdiff_t size = N,
                                             int *error = nullptr)
    {
      int myerror = MPI_SUCCESS !=
                    MPI_Allreduce(MPI_IN_PLACE, value->data(), value->size() * N, MPITypeSelect<T>(), operation, mComm);
      if (error) {
        *error = myerror;
      } else {
        combineAllError(myerror);
      }
      return *value;
    }

    template <typename T, size_t N>
    std::array<T, N> &Allreduce(std::array<T, N> *value, MPI_Op operation, ptrdiff_t size = N, int *error = nullptr)
    {
      int myerror = MPI_SUCCESS != MPI_Allreduce(MPI_IN_PLACE, value->data(), N, MPITypeSelect<T>(), operation, mComm);
      if (error) {
        *error = myerror;
      } else {
        combineAllError(myerror);
      }
      return *value;
    }

    template <typename T>
    complex<T> Allreduce(complex<T> value, MPI_Op operation, ptrdiff_t size = 2, int *error = nullptr)
    {
      std::array<double, 2u> realView{{value.real(), value.imag()}};
      Allreduce(&realView, operation);
      return complex<T>(realView[0], realView[1]);
    }

#else
    template <typename T>
    const T &Allreduce(const T &whatever, MPI_Op operation, ptrdiff_t size = 1, int *error = nullptr)
    {
      return whatever;
    }

    template <typename T> T &Allreduce(T *whatever, MPI_Op operation, int *error = nullptr) { return *whatever; }
#endif
    template <typename T> T computeAllSum(const T &value, size_t size = 1, int *error = nullptr)
    {
      return Allreduce(value, MPI_SUM, size, error);
    }

    template <typename T> T computeAllOr(const T &value, size_t size = 1, int *error = nullptr)
    {
      return Allreduce(value, MPI_LOR, size, error);
    }

    template <typename T> T computeAllAnd(const T &value, size_t size = 1, int *error = nullptr)
    {
      return Allreduce(value, MPI_LAND, size, error);
    }

    template <typename T> T computeAllMin(const T &value, size_t size = 1, int *error = nullptr)
    {
      return Allreduce(value, MPI_MIN, size, error);
    }

    template <typename T> T computeAllMax(const T &value, size_t size = 1, int *error = nullptr)
    {
      return Allreduce(value, MPI_MAX, size, error);
    }

    int combineAllError(int error)
    {
      int nextError;
      computeAllOr(error, 1, &nextError);
      int result = (error || nextError);
      if (result) throw MPIAllReduceException("Received error from MPI after call to MPI_Allreduce.");
      return result;
    }

  private:
    /* Put all member variables and private methods here. These may change arbitrarily. */
    MPI_Comm mComm;

    void useMCommToAvoidWarningUnused()
    {
      if (!mComm) {
        say << "Hello world.\n This function is never called.\n";
      }
    }
  };
} // namespace TempLat

#endif