Tpetra_Details_iallreduce.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef TPETRA_DETAILS_IALLREDUCE_HPP
#define TPETRA_DETAILS_IALLREDUCE_HPP


#include "TpetraCore_config.h"
#include "Teuchos_EReductionType.hpp"
#ifdef HAVE_TPETRACORE_MPI
#  include "Tpetra_Details_extractMpiCommFromTeuchos.hpp"
#  include "Tpetra_Details_MpiTypeTraits.hpp"
#endif // HAVE_TPETRACORE_MPI
#include "Tpetra_Details_temporaryViewUtils.hpp"
#include "Tpetra_Details_Behavior.hpp"
#include "Kokkos_Core.hpp"
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <functional>

#ifndef DOXYGEN_SHOULD_SKIP_THIS
namespace Teuchos {
  // forward declaration of Comm
  template<class OrdinalType> class Comm;
} // namespace Teuchos
#endif // NOT DOXYGEN_SHOULD_SKIP_THIS

namespace Tpetra {
namespace Details {

#ifdef HAVE_TPETRACORE_MPI
std::string getMpiErrorString (const int errCode);
#endif

class CommRequest {
public:
  virtual ~CommRequest () {}

  virtual void wait () {}

  virtual void cancel () {}
};

// Don't rely on anything in this namespace.
namespace Impl {

std::shared_ptr<CommRequest>
emptyCommRequest ();

#ifdef HAVE_TPETRACORE_MPI
#if MPI_VERSION >= 3
template<typename InputViewType, typename OutputViewType, typename ResultViewType>
struct MpiRequest : public CommRequest
{
  MpiRequest(const InputViewType& send, const OutputViewType& recv, const ResultViewType& result, MPI_Request req_)
    : sendBuf(send), recvBuf(recv), resultBuf(result), req(req_)
  {}

  ~MpiRequest()
  {
    //this is a no-op if wait() or cancel() have already been called
    cancel();
  }

  void wait () override
  {
    if (req != MPI_REQUEST_NULL) {
      const int err = MPI_Wait (&req, MPI_STATUS_IGNORE);
      TEUCHOS_TEST_FOR_EXCEPTION
        (err != MPI_SUCCESS, std::runtime_error,
         "MpiCommRequest::wait: MPI_Wait failed with error \""
         << getMpiErrorString (err));
      // MPI_Wait should set the MPI_Request to MPI_REQUEST_NULL on
      // success.  We'll do it here just to be conservative.
      req = MPI_REQUEST_NULL;
      //Since recvBuf contains the result, copy it to the user's resultBuf.
      Kokkos::deep_copy(resultBuf, recvBuf);
    }
  }

  void cancel () override
  {
    //BMK: Per https://www.mpi-forum.org/docs/mpi-3.1/mpi31-report/node126.htm,
    //MPI_Cancel cannot be used for collectives like iallreduce.
    req = MPI_REQUEST_NULL;
  }

private:
  InputViewType sendBuf;
  OutputViewType recvBuf;
  ResultViewType resultBuf;
  //This request is active if and only if req != MPI_REQUEST_NULL.
  MPI_Request req;
};

MPI_Request
iallreduceRaw (const void* sendbuf,
               void* recvbuf,
               const int count,
               MPI_Datatype mpiDatatype,
               const Teuchos::EReductionType op,
               MPI_Comm comm);
#endif

void
allreduceRaw  (const void* sendbuf,
               void* recvbuf,
               const int count,
               MPI_Datatype mpiDatatype,
               const Teuchos::EReductionType op,
               MPI_Comm comm);

template<class InputViewType, class OutputViewType>
std::shared_ptr<CommRequest>
iallreduceImpl (const InputViewType& sendbuf,
            const OutputViewType& recvbuf,
            const ::Teuchos::EReductionType op,
            const ::Teuchos::Comm<int>& comm)
{
  using Packet = typename InputViewType::non_const_value_type;
  if(comm.getSize() == 1)
  {
    Kokkos::deep_copy(recvbuf, sendbuf);
    return emptyCommRequest();
  }
  Packet examplePacket;
  MPI_Datatype mpiDatatype = sendbuf.extent(0) ?
    MpiTypeTraits<Packet>::getType (examplePacket) :
    MPI_BYTE;
  bool datatypeNeedsFree = MpiTypeTraits<Packet>::needsFree;
  MPI_Comm rawComm = ::Tpetra::Details::extractMpiCommFromTeuchos (comm);
  //Note BMK: Nonblocking collectives like iallreduce cannot use GPU buffers.
  //See https://www.open-mpi.org/faq/?category=runcuda#mpi-cuda-support
  auto sendMPI = Tpetra::Details::TempView::toMPISafe<InputViewType, false>(sendbuf);
  auto recvMPI = Tpetra::Details::TempView::toMPISafe<OutputViewType, false>(recvbuf);
  std::shared_ptr<CommRequest> req;
  //Next, if input/output alias and comm is an intercomm, make a deep copy of input.
  //Not possible to do in-place allreduce for intercomm.
  if(isInterComm(comm) && sendMPI.data() == recvMPI.data())
  {
    //Can't do in-place collective on an intercomm,
    //so use a separate 1D copy as the input.
    Kokkos::View<Packet*, Kokkos::HostSpace> tempInput(Kokkos::ViewAllocateWithoutInitializing("tempInput"), sendMPI.extent(0));
    for(size_t i = 0; i < sendMPI.extent(0); i++)
      tempInput(i) = sendMPI.data()[i];
#if MPI_VERSION >= 3
    //MPI 3+: use async allreduce
    MPI_Request mpiReq = iallreduceRaw((const void*) tempInput.data(), (void*) recvMPI.data(), tempInput.extent(0), mpiDatatype, op, rawComm);
    req = std::shared_ptr<CommRequest>(new MpiRequest<decltype(tempInput), decltype(recvMPI), OutputViewType>(tempInput, recvMPI, recvbuf, mpiReq));
#else
    //Older MPI: Iallreduce not available. Instead do blocking all-reduce and return empty request.
    allreduceRaw((const void*) sendMPI.data(), (void*) recvMPI.data(), sendMPI.extent(0), mpiDatatype, op, rawComm);
    Kokkos::deep_copy(recvbuf, recvMPI);
    req = emptyCommRequest();
#endif
  }
  else
  {
#if MPI_VERSION >= 3
    //MPI 3+: use async allreduce
    MPI_Request mpiReq = iallreduceRaw((const void*) sendMPI.data(), (void*) recvMPI.data(), sendMPI.extent(0), mpiDatatype, op, rawComm);
    req = std::shared_ptr<CommRequest>(new MpiRequest<decltype(sendMPI), decltype(recvMPI), OutputViewType>(sendMPI, recvMPI, recvbuf, mpiReq));
#else
    //Older MPI: Iallreduce not available. Instead do blocking all-reduce and return empty request.
    allreduceRaw((const void*) sendMPI.data(), (void*) recvMPI.data(), sendMPI.extent(0), mpiDatatype, op, rawComm);
    Kokkos::deep_copy(recvbuf, recvMPI);
    req = emptyCommRequest();
#endif
  }
  if(datatypeNeedsFree)
    MPI_Type_free(&mpiDatatype);
  return req;
}

#else

//No MPI: reduction is always the same as input.
template<class InputViewType, class OutputViewType>
std::shared_ptr<CommRequest>
iallreduceImpl (const InputViewType& sendbuf,
            const OutputViewType& recvbuf,
            const ::Teuchos::EReductionType,
            const ::Teuchos::Comm<int>&)
{
  Kokkos::deep_copy(recvbuf, sendbuf);
  return emptyCommRequest();
}

#endif // HAVE_TPETRACORE_MPI

} // namespace Impl

//
// SKIP DOWN TO HERE
//

template<class InputViewType, class OutputViewType>
std::shared_ptr<CommRequest>
iallreduce (const InputViewType& sendbuf,
            const OutputViewType& recvbuf,
            const ::Teuchos::EReductionType op,
            const ::Teuchos::Comm<int>& comm)
{
  static_assert (Kokkos::is_view<InputViewType>::value,
                 "InputViewType must be a Kokkos::View specialization.");
  static_assert (Kokkos::is_view<OutputViewType>::value,
                 "OutputViewType must be a Kokkos::View specialization.");
  constexpr int rank = static_cast<int> (OutputViewType::rank);
  static_assert (static_cast<int> (InputViewType::rank) == rank,
                 "InputViewType and OutputViewType must have the same rank.");
  static_assert (rank == 0 || rank == 1,
                 "InputViewType and OutputViewType must both have "
                 "rank 0 or rank 1.");
  typedef typename OutputViewType::non_const_value_type packet_type;
  static_assert (std::is_same<typename OutputViewType::value_type,
                   packet_type>::value,
                 "OutputViewType must be a nonconst Kokkos::View.");
  static_assert (std::is_same<typename InputViewType::non_const_value_type,
                   packet_type>::value,
                 "InputViewType and OutputViewType must be Views "
                 "whose entries have the same type.");
  //Make sure layouts are contiguous (don't accept strided 1D view)
  static_assert (!std::is_same<typename InputViewType::array_layout, Kokkos::LayoutStride>::value,
      "Input/Output views must be contiguous (not LayoutStride)");
  static_assert (!std::is_same<typename OutputViewType::array_layout, Kokkos::LayoutStride>::value,
      "Input/Output views must be contiguous (not LayoutStride)");

  return Impl::iallreduceImpl<InputViewType, OutputViewType> (sendbuf, recvbuf, op, comm);
}

std::shared_ptr<CommRequest>
iallreduce (const int localValue,
            int& globalValue,
            const ::Teuchos::EReductionType op,
            const ::Teuchos::Comm<int>& comm);

} // namespace Details
} // namespace Tpetra

#endif // TPETRA_DETAILS_IALLREDUCE_HPP