Tpetra_Details_DistributorActor.hpp

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

#ifndef TPETRA_DETAILS_DISTRIBUTOR_ACTOR_HPP
#define TPETRA_DETAILS_DISTRIBUTOR_ACTOR_HPP

#include <optional>
#include <vector>

#include "Teuchos_Assert.hpp"
#include "Tpetra_Details_DistributorPlan.hpp"
#include "Tpetra_Util.hpp"

#include "Teuchos_Array.hpp"
#include "Teuchos_Comm.hpp"
#include "Tpetra_Details_MpiTypeTraits.hpp"
#include "Tpetra_Details_Profiling.hpp"
#include "Teuchos_RCP.hpp"

#include "Kokkos_TeuchosCommAdapters.hpp"

#ifdef HAVE_TPETRA_MPI
#include "mpi.h"
#include "Tpetra_Details_Ialltofewv.hpp"
#endif

namespace Tpetra::Details {

template <class View>
constexpr bool isKokkosView = Kokkos::is_view<View>::value;

template <class View1, class View2>
constexpr bool areKokkosViews = Kokkos::is_view<View1>::value && Kokkos::is_view<View2>::value;

class DistributorActor {
  static constexpr int DEFAULT_MPI_TAG = 1;

  using IndexView = DistributorPlan::IndexView;
  using SubViewLimits = DistributorPlan::SubViewLimits;

public:
  DistributorActor();
  DistributorActor(const DistributorActor& otherActor) = default;

  template <class ExpView, class ImpView>
  void doPostsAndWaits(const DistributorPlan& plan,
                       const ExpView &exports,
                       size_t numPackets,
                       const ImpView &imports);

  template <class ExpView, class ImpView>
  void doPostsAndWaits(const DistributorPlan& plan,
                       const ExpView &exports,
                       const Teuchos::ArrayView<const size_t>& numExportPacketsPerLID,
                       const ImpView &imports,
                       const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID);

  template <class ImpView>
  void doPostRecvs(const DistributorPlan& plan,
                   size_t numPackets,
                   const ImpView& imports);

  template <class ImpView>
  void doPostRecvs(const DistributorPlan& plan,
                   const ImpView &imports,
                   const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID);

  template <class ExpView, class ImpView>
  void doPostSends(const DistributorPlan& plan,
                   const ExpView& exports,
                   size_t numPackets,
                   const ImpView& imports);

  template <class ExpView, class ImpView>
  void doPostSends(const DistributorPlan& plan,
                   const ExpView &exports,
                   const Teuchos::ArrayView<const size_t>& numExportPacketsPerLID,
                   const ImpView &imports,
                   const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID);

  template <class ExpView, class ImpView>
  void doPosts(const DistributorPlan& plan,
               const ExpView& exports,
               size_t numPackets,
               const ImpView& imports);

  template <class ExpView, class ImpView>
  void doPosts(const DistributorPlan& plan,
               const ExpView &exports,
               const Teuchos::ArrayView<const size_t>& numExportPacketsPerLID,
               const ImpView &imports,
               const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID);

  void doWaits(const DistributorPlan& plan);

  void doWaitsRecv(const DistributorPlan& plan);

  void doWaitsSend(const DistributorPlan& plan);

  void doWaitsIalltofewv(const DistributorPlan &plan);

  bool isReady() const;

private:

  template <class ImpView>
  void doPostRecvsImpl(const DistributorPlan& plan,
                       const ImpView& imports,
                       const SubViewLimits& totalPacketsFrom);

  template <class ExpView, class ImpView>
  void doPostSendsImpl(const DistributorPlan& plan,
                       const ExpView& exports,
                       const SubViewLimits& exportSubViewLimits,
                       const ImpView& imports,
                       const SubViewLimits& importSubViewLimits);

#ifdef HAVE_TPETRA_MPI
  template <class ExpView, class ImpView>
  void doPostsAllToAllImpl(const DistributorPlan &plan,
                           const ExpView &exports,
                           const SubViewLimits& exportSubViewLimits,
                           const ImpView &imports,
                           const SubViewLimits& importSubViewLimits);

#if defined(HAVE_TPETRACORE_MPI_ADVANCE)
  template <class ExpView, class ImpView>
  void doPostsNbrAllToAllVImpl(const DistributorPlan &plan,
                               const ExpView &exports,
                               const SubViewLimits& exportSubViewLimits,
                               const ImpView &imports,
                               const SubViewLimits& importSubViewLimits);
#endif // HAVE_TPETRACORE_MPI_ADVANCE

  template <typename ExpView, typename ImpView>
  void doPostsIalltofewvImpl(const DistributorPlan &plan,
                             const ExpView &exports,
                             const SubViewLimits& exportSubViewLimits,
                             const ImpView &imports,
                             const SubViewLimits& importSubViewLimits);

  // ialltofewv members
  struct {
    Details::Ialltofewv impl;
    std::optional<Details::Ialltofewv::Req> req;
    Teuchos::RCP<std::vector<int>> sendcounts;
    Teuchos::RCP<std::vector<int>> sdispls;
    Teuchos::RCP<std::vector<int>> recvcounts;
    Teuchos::RCP<std::vector<int>> rdispls;
    std::vector<int> roots;
  } ialltofewv_;

#endif // HAVE_TPETRA_MPI

  int mpiTag_;

  Teuchos::Array<Teuchos::RCP<Teuchos::CommRequest<int>>> requestsRecv_;
  Teuchos::Array<Teuchos::RCP<Teuchos::CommRequest<int>>> requestsSend_;
};

template <class ExpView, class ImpView>
void DistributorActor::doPosts(const DistributorPlan& plan,
                               const ExpView& exports,
                               size_t numPackets,
                               const ImpView& imports)
{
  doPostRecvs(plan, numPackets, imports);
  doPostSends(plan, exports, numPackets, imports);
}

template <class ExpView, class ImpView>
void DistributorActor::doPostsAndWaits(const DistributorPlan& plan,
                                       const ExpView& exports,
                                       size_t numPackets,
                                       const ImpView& imports)
{
  static_assert(areKokkosViews<ExpView, ImpView>,
      "Data arrays for DistributorActor::doPostsAndWaits must be Kokkos::Views");
  doPosts(plan, exports, numPackets, imports);
  doWaits(plan);
}

template <class ExpView, class ImpView>
void DistributorActor::doPosts(const DistributorPlan& plan,
                               const ExpView &exports,
                               const Teuchos::ArrayView<const size_t>& numExportPacketsPerLID,
                               const ImpView &imports,
                               const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID)
{
  doPostRecvs(plan, imports, numImportPacketsPerLID);
  doPostSends(plan, exports, numExportPacketsPerLID, imports, numImportPacketsPerLID);
}

template <class ExpView, class ImpView>
void DistributorActor::doPostsAndWaits(const DistributorPlan& plan,
                                       const ExpView& exports,
                                       const Teuchos::ArrayView<const size_t>& numExportPacketsPerLID,
                                       const ImpView& imports,
                                       const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID)
{
  static_assert(areKokkosViews<ExpView, ImpView>,
      "Data arrays for DistributorActor::doPostsAndWaits must be Kokkos::Views");
  doPosts(plan, exports, numExportPacketsPerLID, imports, numImportPacketsPerLID);
  doWaits(plan);
}

template <typename ViewType>
using HostAccessibility = Kokkos::SpaceAccessibility<Kokkos::DefaultHostExecutionSpace, typename ViewType::memory_space>;

template <typename DstViewType, typename SrcViewType>
using enableIfHostAccessible = std::enable_if_t<HostAccessibility<DstViewType>::accessible &&
                                                HostAccessibility<SrcViewType>::accessible>;

template <typename DstViewType, typename SrcViewType>
using enableIfNotHostAccessible = std::enable_if_t<!HostAccessibility<DstViewType>::accessible ||
                                                   !HostAccessibility<SrcViewType>::accessible>;

template <typename DstViewType, typename SrcViewType>
enableIfHostAccessible<DstViewType, SrcViewType>
packOffset(const DstViewType& dst,
           const SrcViewType& src,
           const size_t dst_offset,
           const size_t src_offset,
           const size_t size)
{
  memcpy((void*) (dst.data()+dst_offset), src.data()+src_offset, size*sizeof(typename DstViewType::value_type));
}

template <typename DstViewType, typename SrcViewType>
enableIfNotHostAccessible<DstViewType, SrcViewType>
packOffset(const DstViewType& dst,
           const SrcViewType& src,
           const size_t dst_offset,
           const size_t src_offset,
           const size_t size)
{
  Kokkos::Compat::deep_copy_offset(dst, src, dst_offset, src_offset, size);
}

#ifdef HAVE_TPETRA_MPI

template <class ExpView, class ImpView>
void DistributorActor::doPostsIalltofewvImpl(const DistributorPlan &plan,
                                           const ExpView &exports,
                                           const SubViewLimits& exportSubViewLimits,
                                           const ImpView &imports,
                                           const SubViewLimits& importSubViewLimits) {
  using size_type = Teuchos::Array<size_t>::size_type;
  using ExportValue = typename ExpView::non_const_value_type;

  ProfilingRegion pr("Tpetra::Distributor: doPostsIalltofewvImpl");

  TEUCHOS_TEST_FOR_EXCEPTION(
      !plan.getIndicesTo().is_null(), std::runtime_error,
      "Send Type=\"Ialltofewv\" only works for fast-path communication.");

  TEUCHOS_TEST_FOR_EXCEPTION(
    bool(ialltofewv_.req), std::runtime_error,
    "This actor has an active Ialltofewv already");

  TEUCHOS_TEST_FOR_EXCEPTION(
      bool(ialltofewv_.sendcounts), std::runtime_error,
      "This actor has an active Ialltofewv already");

  TEUCHOS_TEST_FOR_EXCEPTION(
    bool(ialltofewv_.sdispls), std::runtime_error,
      "This actor has an active Ialltofewv already");

  TEUCHOS_TEST_FOR_EXCEPTION(
    bool(ialltofewv_.recvcounts), std::runtime_error,
      "This actor has an active Ialltofewv already");

  TEUCHOS_TEST_FOR_EXCEPTION(
    bool(ialltofewv_.rdispls), std::runtime_error,
      "This actor has an active Ialltofewv already");

  auto comm = plan.getComm();

  const auto& [importStarts, importLengths] = importSubViewLimits;
  const auto& [exportStarts, exportLengths] = exportSubViewLimits;


  ialltofewv_.roots = plan.getRoots();
  const int nroots = ialltofewv_.roots.size();
  const int *roots = ialltofewv_.roots.data();
  ialltofewv_.req = std::make_optional<Details::Ialltofewv::Req>();
  ialltofewv_.sendcounts = Teuchos::rcp(new std::vector<int>(nroots));
  ialltofewv_.sdispls = Teuchos::rcp(new std::vector<int>(nroots));
  ialltofewv_.recvcounts = Teuchos::rcp(new std::vector<int>);
  ialltofewv_.rdispls = Teuchos::rcp(new std::vector<int>);

  for (int rootIdx = 0; rootIdx < nroots; ++rootIdx) {
    const int root = roots[rootIdx];

    //if we can't find the root proc index in our plan, it just means we send 0
    // also make sure root only appears once in getProcsTo()
    size_type rootProcIndex = plan.getProcsTo().size(); // sentinel value -> not found
    for (size_type pi = 0; pi < plan.getProcsTo().size(); ++pi) {
      if (plan.getProcsTo()[pi] == root) {
        rootProcIndex = pi;
        break;
      }
    }

    // am I sending to root?
    int sendcount = 0;
    if (rootProcIndex != plan.getProcsTo().size()) {
      sendcount = exportLengths[rootProcIndex];
    }
    (*ialltofewv_.sendcounts)[rootIdx] = sendcount;

    int sdispl = 0;
    if (0 != sendcount) {
      sdispl = exportStarts[rootProcIndex];
    }
    (*ialltofewv_.sdispls)[rootIdx] = sdispl;

    // If I happen to be this root, set up my receive metadata
    if (comm->getRank() == root) {

      // don't recv anything from anywhere by default
      ialltofewv_.recvcounts->resize(comm->getSize());
      std::fill(ialltofewv_.recvcounts->begin(), ialltofewv_.recvcounts->end(), 0);
      ialltofewv_.rdispls->resize(comm->getSize());
      std::fill(ialltofewv_.rdispls->begin(), ialltofewv_.rdispls->end(), 0);

      const size_type actualNumReceives =
      Teuchos::as<size_type>(plan.getNumReceives()) +
      Teuchos::as<size_type>(plan.hasSelfMessage() ? 1 : 0);

      for (size_type i = 0; i < actualNumReceives; ++i) {
        const int src = plan.getProcsFrom()[i];
        (*ialltofewv_.rdispls)[src] = importStarts[i];
        (*ialltofewv_.recvcounts)[src] = Teuchos::as<int>(importLengths[i]);
      }
    }

  } // rootIdx

  // TODO: do we need to pass ExportValue{} here?
  MPI_Datatype rawType = ::Tpetra::Details::MpiTypeTraits<ExportValue>::getType(ExportValue{});
  // FIXME: is there a better way to do this?
  Teuchos::RCP<const Teuchos::MpiComm<int>> tMpiComm =
    Teuchos::rcp_dynamic_cast<const Teuchos::MpiComm<int>>(comm);
  Teuchos::RCP<const Teuchos::OpaqueWrapper<MPI_Comm>> oMpiComm =
    tMpiComm->getRawMpiComm();
  MPI_Comm mpiComm = (*oMpiComm)();

  // don't care about send-side accessibility because it's not accessed through kokkos
  // rely on MPI to do the right thing
  constexpr bool recvDevAccess = Kokkos::SpaceAccessibility<
     Kokkos::DefaultExecutionSpace, typename ImpView::memory_space>::accessible;
  constexpr bool sendDevAccess = Kokkos::SpaceAccessibility<
     Kokkos::DefaultExecutionSpace, typename ExpView::memory_space>::accessible;
  static_assert(recvDevAccess == sendDevAccess, "sending across host/device");

  const int err = ialltofewv_.impl.post<recvDevAccess>(exports.data(), ialltofewv_.sendcounts->data(), ialltofewv_.sdispls->data(), rawType,
    imports.data(), ialltofewv_.recvcounts->data(), ialltofewv_.rdispls->data(), 
    roots, nroots,
    rawType,
    mpiTag_, mpiComm, &(*ialltofewv_.req));

  TEUCHOS_TEST_FOR_EXCEPTION(err != MPI_SUCCESS, std::runtime_error,
                              "ialltofewv failed with error \""
                              << Teuchos::mpiErrorCodeToString(err)
                              << "\".");
}





template <class ExpView, class ImpView>
void DistributorActor::doPostsAllToAllImpl(const DistributorPlan &plan,
                                           const ExpView &exports,
                                           const SubViewLimits& exportSubViewLimits,
                                           const ImpView &imports,
                                           const SubViewLimits& importSubViewLimits) {
  TEUCHOS_TEST_FOR_EXCEPTION(
      !plan.getIndicesTo().is_null(), std::runtime_error,
      "Send Type=\"Alltoall\" only works for fast-path communication.");

  using size_type = Teuchos::Array<size_t>::size_type;

  auto comm = plan.getComm();
  std::vector<int> sendcounts(comm->getSize(), 0);
  std::vector<int> sdispls(comm->getSize(), 0);
  std::vector<int> recvcounts(comm->getSize(), 0);
  std::vector<int> rdispls(comm->getSize(), 0);

  auto& [importStarts, importLengths] = importSubViewLimits;
  auto& [exportStarts, exportLengths] = exportSubViewLimits;

  for (size_t pp = 0; pp < plan.getNumSends(); ++pp) {
    sdispls[plan.getProcsTo()[pp]] = exportStarts[pp];
    size_t numPackets = exportLengths[pp];
    // numPackets is converted down to int, so make sure it can be represented
    TEUCHOS_TEST_FOR_EXCEPTION(numPackets > size_t(INT_MAX), std::logic_error,
                               "Tpetra::Distributor::doPostsAllToAll: "
                               "Send count for send "
                                   << pp << " (" << numPackets
                                   << ") is too large "
                                      "to be represented as int.");
    sendcounts[plan.getProcsTo()[pp]] = static_cast<int>(numPackets);
  }

  const size_type actualNumReceives =
      Teuchos::as<size_type>(plan.getNumReceives()) +
      Teuchos::as<size_type>(plan.hasSelfMessage() ? 1 : 0);

  for (size_type i = 0; i < actualNumReceives; ++i) {
    rdispls[plan.getProcsFrom()[i]] = importStarts[i];
    size_t totalPacketsFrom_i = importLengths[i];
    // totalPacketsFrom_i is converted down to int, so make sure it can be
    // represented
    TEUCHOS_TEST_FOR_EXCEPTION(totalPacketsFrom_i > size_t(INT_MAX),
                               std::logic_error,
                               "Tpetra::Distributor::doPostsAllToAll: "
                               "Recv count for receive "
                                   << i << " (" << totalPacketsFrom_i
                                   << ") is too large "
                                      "to be represented as int.");
    recvcounts[plan.getProcsFrom()[i]] = static_cast<int>(totalPacketsFrom_i);
  }

  Teuchos::RCP<const Teuchos::MpiComm<int>> mpiComm =
      Teuchos::rcp_dynamic_cast<const Teuchos::MpiComm<int>>(comm);
  Teuchos::RCP<const Teuchos::OpaqueWrapper<MPI_Comm>> rawComm =
      mpiComm->getRawMpiComm();
  using T = typename ExpView::non_const_value_type;
  MPI_Datatype rawType = ::Tpetra::Details::MpiTypeTraits<T>::getType(T());

#if defined(HAVE_TPETRACORE_MPI_ADVANCE)
  if (Details::DISTRIBUTOR_MPIADVANCE_ALLTOALL == plan.getSendType()) {
    MPIX_Comm *mpixComm = *plan.getMPIXComm();
    TEUCHOS_TEST_FOR_EXCEPTION(!mpixComm, std::runtime_error,
                               "MPIX_Comm is null in doPostsAllToAll \""
                                   << __FILE__ << ":" << __LINE__);

    const int err = MPIX_Alltoallv(
        exports.data(), sendcounts.data(), sdispls.data(), rawType,
        imports.data(), recvcounts.data(), rdispls.data(), rawType, mpixComm);

    TEUCHOS_TEST_FOR_EXCEPTION(err != MPI_SUCCESS, std::runtime_error,
                               "MPIX_Alltoallv failed with error \""
                                   << Teuchos::mpiErrorCodeToString(err)
                                   << "\".");

    return;
  }
#endif // HAVE_TPETRACORE_MPI_ADVANCE

  const int err = MPI_Alltoallv(
      exports.data(), sendcounts.data(), sdispls.data(), rawType,
      imports.data(), recvcounts.data(), rdispls.data(), rawType, (*rawComm)());

  TEUCHOS_TEST_FOR_EXCEPTION(err != MPI_SUCCESS, std::runtime_error,
                             "MPI_Alltoallv failed with error \""
                                 << Teuchos::mpiErrorCodeToString(err)
                                 << "\".");
}

#if defined(HAVE_TPETRACORE_MPI_ADVANCE)
template <class ExpView, class ImpView>
void DistributorActor::doPostsNbrAllToAllVImpl(const DistributorPlan &plan,
                                               const ExpView &exports,
                                               const SubViewLimits& exportSubViewLimits,
                                               const ImpView &imports,
                                               const SubViewLimits& importSubViewLimits) {
  TEUCHOS_TEST_FOR_EXCEPTION(
      !plan.getIndicesTo().is_null(), std::runtime_error,
      "Send Type=\"Alltoall\" only works for fast-path communication.");

  const int myRank = plan.getComm()->getRank();
  MPIX_Comm *mpixComm = *plan.getMPIXComm();
  using size_type = Teuchos::Array<size_t>::size_type;

  const size_t numSends = plan.getNumSends() + plan.hasSelfMessage();
  const size_t numRecvs = plan.getNumReceives() + plan.hasSelfMessage();
  std::vector<int> sendcounts(numSends, 0);
  std::vector<int> sdispls(numSends, 0);
  std::vector<int> recvcounts(numRecvs, 0);
  std::vector<int> rdispls(numRecvs, 0);

  auto& [importStarts, importLengths] = importSubViewLimits;
  auto& [exportStarts, exportLengths] = exportSubViewLimits;

  for (size_t pp = 0; pp < numSends; ++pp) {
    sdispls[pp] = exportStarts[pp];
    size_t numPackets = exportLengths[pp];
    // numPackets is converted down to int, so make sure it can be represented
    TEUCHOS_TEST_FOR_EXCEPTION(numPackets > size_t(INT_MAX), std::logic_error,
                               "Tpetra::Distributor::doPostsNbrAllToAllV: "
                               "Send count for send "
                                   << pp << " (" << numPackets
                                   << ") is too large "
                                      "to be represented as int.");
    sendcounts[pp] = static_cast<int>(numPackets);
  }

  for (size_type i = 0; i < numRecvs; ++i) {
    rdispls[i] = importStarts[i];
    size_t totalPacketsFrom_i = importLengths[i];
    // totalPacketsFrom_i is converted down to int, so make sure it can be
    // represented
    TEUCHOS_TEST_FOR_EXCEPTION(totalPacketsFrom_i > size_t(INT_MAX),
                               std::logic_error,
                               "Tpetra::Distributor::doPostsNbrAllToAllV: "
                               "Recv count for receive "
                                   << i << " (" << totalPacketsFrom_i
                                   << ") is too large "
                                      "to be represented as int.");
    recvcounts[i] = static_cast<int>(totalPacketsFrom_i);
  }

  using T = typename ExpView::non_const_value_type;
  MPI_Datatype rawType = ::Tpetra::Details::MpiTypeTraits<T>::getType(T());

  MPIX_Info* xinfo;
  MPIX_Topo* xtopo;
  MPIX_Info_init(&xinfo);
  MPIX_Topo_init(numRecvs, plan.getProcsFrom().data(), recvcounts.data(),
      numSends, plan.getProcsTo().data(), sendcounts.data(), xinfo, &xtopo);
  const int err = MPIX_Neighbor_alltoallv_topo(
      exports.data(), sendcounts.data(), sdispls.data(), rawType,
      imports.data(), recvcounts.data(), rdispls.data(), rawType, xtopo, mpixComm);
  MPIX_Topo_free(&xtopo);
  MPIX_Info_free(&xinfo);

  TEUCHOS_TEST_FOR_EXCEPTION(err != MPI_SUCCESS, std::runtime_error,
                             "MPIX_Neighbor_alltoallv failed with error \""
                                 << Teuchos::mpiErrorCodeToString(err)
                                 << "\".");
}
#endif // HAVE_TPETRACORE_MPI_ADVANCE
#endif // HAVE_TPETRA_MPI

template <class ImpView>
void DistributorActor::doPostRecvs(const DistributorPlan& plan,
                                   size_t numPackets,
                                   const ImpView& imports)
{
  auto importSubViewLimits = plan.getImportViewLimits(numPackets);
  doPostRecvsImpl(plan, imports, importSubViewLimits);
}

template <class ImpView>
void DistributorActor::doPostRecvs(const DistributorPlan& plan,
                                   const ImpView &imports,
                                   const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID)
{
  auto importSubViewLimits = plan.getImportViewLimits(numImportPacketsPerLID);
  doPostRecvsImpl(plan, imports, importSubViewLimits);
}

template <class ImpView>
void DistributorActor::doPostRecvsImpl(const DistributorPlan& plan,
                                       const ImpView &imports,
                                       const SubViewLimits& importSubViewLimits)
{
  static_assert(isKokkosView<ImpView>,
      "Data arrays for DistributorActor::doPostRecvs must be Kokkos::Views");
  using Teuchos::Array;
  using Teuchos::as;
  using Teuchos::ireceive;
  using Kokkos::Compat::subview_offset;
  using size_type = Array<size_t>::size_type;
  using imports_view_type = ImpView;

#ifdef KOKKOS_ENABLE_CUDA
  static_assert (! std::is_same<typename ImpView::memory_space, Kokkos::CudaUVMSpace>::value,
                 "Please do not use Tpetra::Distributor with UVM "
                 "allocations.  See GitHub issue #1088.");
#endif // KOKKOS_ENABLE_CUDA

#ifdef KOKKOS_ENABLE_SYCL
    static_assert (! std::is_same<typename ImpView::memory_space, Kokkos::Experimental::SYCLSharedUSMSpace>::value,
                   "Please do not use Tpetra::Distributor with SharedUSM "
                   "allocations.  See GitHub issue #1088 (corresponding to CUDA).");
#endif // KOKKOS_ENABLE_SYCL


#if defined(HAVE_TPETRA_MPI)
  // All-to-all communication layout is quite different from
  // point-to-point, so we handle it separately.

  // These send options require no matching receives, so we just return.
  const Details::EDistributorSendType sendType = plan.getSendType();
  if ((sendType == Details::DISTRIBUTOR_ALLTOALL)
      || (sendType == Details::DISTRIBUTOR_IALLTOFEWV)
#ifdef HAVE_TPETRACORE_MPI_ADVANCE
      || (sendType == Details::DISTRIBUTOR_MPIADVANCE_ALLTOALL)
      || (sendType == Details::DISTRIBUTOR_MPIADVANCE_NBRALLTOALLV)
#endif
      ) {
    return;
  }
#endif // HAVE_TPETRA_MPI

  ProfilingRegion pr("Tpetra::Distributor::doPostRecvs");

  const int myProcID = plan.getComm()->getRank ();

  auto& [importStarts, importLengths] = importSubViewLimits;

  // Distributor uses requestsRecv_.size() and requestsSend_.size()
  // as the number of outstanding nonblocking message requests, so
  // we resize to zero to maintain this invariant.
  //
  // getNumReceives() does _not_ include the self message, if there is
  // one.  Here, we do actually send a message to ourselves, so we
  // include any self message in the "actual" number of receives to
  // post.
  //
  // NOTE (mfh 19 Mar 2012): Epetra_MpiDistributor::DoPosts()
  // doesn't (re)allocate its array of requests.  That happens in
  // CreateFromSends(), ComputeRecvs_(), DoReversePosts() (on
  // demand), or Resize_().
  const size_type actualNumReceives = as<size_type> (plan.getNumReceives()) +
    as<size_type> (plan.hasSelfMessage() ? 1 : 0);

#ifdef HAVE_TPETRA_DEBUG
  size_t totalNumImportPackets = 0;
  for (size_t i = 0; i < Teuchos::as<size_t>(actualNumReceives); ++i) {
    totalNumImportPackets += importLengths[i];
  }
  TEUCHOS_TEST_FOR_EXCEPTION(
      imports.extent (0) < totalNumImportPackets, std::runtime_error,
      "Tpetra::Distributor::doPostRecvs: The 'imports' array must have "
      "enough entries to hold the expected number of import packets.  "
      "imports.extent(0) = " << imports.extent (0) << " < "
      "totalNumImportPackets = " << totalNumImportPackets << ".");
  TEUCHOS_TEST_FOR_EXCEPTION
    (!requestsRecv_.empty(), std::logic_error, "Tpetra::Distributor::"
     "doPostRecvs: Process " << myProcID << ": requestsRecv_.size () = "
     << requestsRecv_.size () << " != 0.");
#endif // HAVE_TPETRA_DEBUG

  requestsRecv_.resize (0);

  // Post the nonblocking receives.  It's common MPI wisdom to post
  // receives before sends.  In MPI terms, this means favoring
  // adding to the "posted queue" (of receive requests) over adding
  // to the "unexpected queue" (of arrived messages not yet matched
  // with a receive).
  {
    ProfilingRegion prr("Tpetra::Distributor::doPostRecvs MPI_Irecv");

    for (size_type i = 0; i < actualNumReceives; ++i) {
      size_t totalPacketsFrom_i = importLengths[Teuchos::as<size_t>(i)];
      TEUCHOS_TEST_FOR_EXCEPTION(totalPacketsFrom_i > size_t(INT_MAX),
                                 std::logic_error, "Tpetra::Distributor::doPostRecvs: "
                                 "Recv count for receive " << i << " (" << totalPacketsFrom_i << ") is too large "
                                 "to be represented as int.");
      if (plan.getProcsFrom()[i] != myProcID && totalPacketsFrom_i) {
        // If my process is receiving these packet(s) from another
        // process (not a self-receive), and if there is at least
        // one packet to receive:
        //
        // 1. Set up the persisting view (recvBuf) into the imports
        //    array, given the offset and size (total number of
        //    packets from process getProcsFrom()[i]).
        // 2. Start the Irecv and save the resulting request.
        imports_view_type recvBuf =
          subview_offset (imports, importStarts[i], totalPacketsFrom_i);
        requestsRecv_.push_back (ireceive<int> (recvBuf, plan.getProcsFrom()[i],
              mpiTag_, *plan.getComm()));
      }
    }
  }
}

template <class ExpView, class ImpView>
void DistributorActor::doPostSends(const DistributorPlan& plan,
                                   const ExpView& exports,
                                   size_t numPackets,
                                   const ImpView& imports)
{
  auto exportSubViewLimits = plan.getExportViewLimits(numPackets);
  auto importSubViewLimits = plan.getImportViewLimits(numPackets);
  doPostSendsImpl(plan, exports, exportSubViewLimits, imports, importSubViewLimits);
}

template <class ExpView, class ImpView>
void DistributorActor::doPostSends(const DistributorPlan& plan,
                                   const ExpView &exports,
                                   const Teuchos::ArrayView<const size_t>& numExportPacketsPerLID,
                                   const ImpView &imports,
                                   const Teuchos::ArrayView<const size_t>& numImportPacketsPerLID)
{
  auto exportSubViewLimits = plan.getExportViewLimits(numExportPacketsPerLID);
  auto importSubViewLimits = plan.getImportViewLimits(numImportPacketsPerLID);
  doPostSendsImpl(plan, exports, exportSubViewLimits, imports, importSubViewLimits);
}

template <class ExpView, class ImpView>
void DistributorActor::doPostSendsImpl(const DistributorPlan& plan,
                                       const ExpView& exports,
                                       const SubViewLimits& exportSubViewLimits,
                                       const ImpView& imports,
                                       const SubViewLimits& importSubViewLimits)
{
  static_assert(areKokkosViews<ExpView, ImpView>,
      "Data arrays for DistributorActor::doPostSends must be Kokkos::Views");
  using Teuchos::Array;
  using Teuchos::as;
  using Teuchos::isend;
  using Teuchos::send;
  using Kokkos::Compat::subview_offset;
  using Kokkos::Compat::deep_copy_offset;
  using size_type = Array<size_t>::size_type;
  using exports_view_type = ExpView;

#ifdef KOKKOS_ENABLE_CUDA
  static_assert
    (! std::is_same<typename ExpView::memory_space, Kokkos::CudaUVMSpace>::value &&
     ! std::is_same<typename ImpView::memory_space, Kokkos::CudaUVMSpace>::value,
     "Please do not use Tpetra::Distributor with UVM allocations.  "
     "See Trilinos GitHub issue #1088.");
#endif // KOKKOS_ENABLE_CUDA

#ifdef KOKKOS_ENABLE_SYCL
    static_assert
      (! std::is_same<typename ExpView::memory_space, Kokkos::Experimental::SYCLSharedUSMSpace>::value &&
       ! std::is_same<typename ImpView::memory_space, Kokkos::Experimental::SYCLSharedUSMSpace>::value,
       "Please do not use Tpetra::Distributor with SharedUSM allocations.  "
       "See Trilinos GitHub issue #1088 (corresponding to CUDA).");
#endif // KOKKOS_ENABLE_SYCL

  ProfilingRegion ps("Tpetra::Distributor::doPostSends");

  const int myRank = plan.getComm()->getRank ();
  // Run-time configurable parameters that come from the input
  // ParameterList set by setParameterList().
  const Details::EDistributorSendType sendType = plan.getSendType();

  auto& [exportStarts, exportLengths] = exportSubViewLimits;
  auto& [importStarts, importLengths] = importSubViewLimits;

#if defined(HAVE_TPETRA_MPI)
  //  All-to-all communication layout is quite different from
  //  point-to-point, so we handle it separately.

  if (sendType == Details::DISTRIBUTOR_ALLTOALL) {
    doPostsAllToAllImpl(plan, exports, exportSubViewLimits, imports, importSubViewLimits);
    return;
  } else if (sendType == Details::DISTRIBUTOR_IALLTOFEWV) {
    doPostsIalltofewvImpl(plan, exports, exportSubViewLimits, imports, importSubViewLimits);
    return;
  }
#ifdef HAVE_TPETRACORE_MPI_ADVANCE
  else if (sendType == Details::DISTRIBUTOR_MPIADVANCE_ALLTOALL) {
    doPostsAllToAllImpl(plan, exports, exportSubViewLimits, imports, importSubViewLimits);
    return;
  } else if (sendType == Details::DISTRIBUTOR_MPIADVANCE_NBRALLTOALLV) {
    doPostsNbrAllToAllVImpl(plan, exports, exportSubViewLimits, imports, importSubViewLimits);
    return;
  }
#endif // defined(HAVE_TPETRACORE_MPI_ADVANCE)


#else // HAVE_TPETRA_MPI
    if (plan.hasSelfMessage()) {
      // This is how we "send a message to ourself": we copy from
      // the export buffer to the import buffer.  That saves
      // Teuchos::Comm implementations other than MpiComm (in
      // particular, SerialComm) the trouble of implementing self
      // messages correctly.  (To do this right, SerialComm would
      // need internal buffer space for messages, keyed on the
      // message's tag.)
      size_t selfReceiveOffset = 0;
      deep_copy_offset(imports, exports, selfReceiveOffset,
                       exportStarts[0],
                       exportLengths[0]);
    }
    // should we just return here?
    // likely not as comm could be a serial comm
#endif // HAVE_TPETRA_MPI

  size_t selfReceiveOffset = 0;

#ifdef HAVE_TPETRA_DEBUG
  TEUCHOS_TEST_FOR_EXCEPTION
    (requestsSend_.size () != 0,
     std::logic_error,
     "Tpetra::Distributor::doPostSends: Process "
     << myRank << ": requestsSend_.size() = " << requestsSend_.size () << " != 0.");
#endif // HAVE_TPETRA_DEBUG

  // Distributor uses requestsRecv_.size() and requestsSend_.size()
  // as the number of outstanding nonblocking message requests, so
  // we resize to zero to maintain this invariant.
  //
  // getNumReceives() does _not_ include the self message, if there is
  // one.  Here, we do actually send a message to ourselves, so we
  // include any self message in the "actual" number of receives to
  // post.
  //
  // NOTE (mfh 19 Mar 2012): Epetra_MpiDistributor::DoPosts()
  // doesn't (re)allocate its array of requests.  That happens in
  // CreateFromSends(), ComputeRecvs_(), DoReversePosts() (on
  // demand), or Resize_().
  const size_type actualNumReceives = as<size_type> (plan.getNumReceives()) +
    as<size_type> (plan.hasSelfMessage() ? 1 : 0);
  requestsSend_.resize (0);

  {
    for (size_type i = 0; i < actualNumReceives; ++i) {
      if (plan.getProcsFrom()[i] == myRank) { // Receiving from myself
        selfReceiveOffset = importStarts[i];  // Remember the self-recv offset
      }
    }
  }

  ProfilingRegion pss("Tpetra::Distributor::doPostSends sends");

  // setup scan through getProcsTo() list starting with higher numbered procs
  // (should help balance message traffic)
  //
  // FIXME (mfh 20 Feb 2013) Why haven't we precomputed this?
  // It doesn't depend on the input at all.
  size_t numBlocks = plan.getNumSends() + plan.hasSelfMessage();
  size_t procIndex = 0;
  while ((procIndex < numBlocks) && (plan.getProcsTo()[procIndex] < myRank)) {
    ++procIndex;
  }
  if (procIndex == numBlocks) {
    procIndex = 0;
  }

  size_t selfNum = 0;
  size_t selfIndex = 0;

  if (plan.getIndicesTo().is_null()) {
    const char isend_region[] = "Tpetra::Distributor::doPostSends MPI_Isend FAST";
    const char send_region[] = "Tpetra::Distributor::doPostSends MPI_Send FAST";
    ProfilingRegion pssf((sendType == Details::DISTRIBUTOR_ISEND) ? isend_region : send_region);

    // Data are already blocked (laid out) by process, so we don't
    // need a separate send buffer (besides the exports array).
    for (size_t i = 0; i < numBlocks; ++i) {
      size_t p = i + procIndex;
      if (p > (numBlocks - 1)) {
        p -= numBlocks;
      }

      if (plan.getProcsTo()[p] != myRank) {
        if (exportLengths[p] == 0) {
          // Do not attempt to send messages of length 0.
          continue;
        }

        exports_view_type tmpSend = subview_offset(exports, exportStarts[p], exportLengths[p]);

        if (sendType == Details::DISTRIBUTOR_ISEND) {
          // NOTE: This looks very similar to the tmpSend above, but removing
          // tmpSendBuf and uses tmpSend leads to a performance hit on Arm
          // SerialNode builds
          exports_view_type tmpSendBuf =
            subview_offset (exports, exportStarts[p], exportLengths[p]);
          requestsSend_.push_back (isend<int> (tmpSendBuf, plan.getProcsTo()[p],
                mpiTag_, *plan.getComm()));
        }
        else {  // DISTRIBUTOR_SEND
          send<int> (tmpSend,
              as<int> (tmpSend.size ()),
              plan.getProcsTo()[p], mpiTag_, *plan.getComm());
        }
      }
      else { // "Sending" the message to myself
        selfNum = p;
      }
    }

    if (plan.hasSelfMessage()) {
      // This is how we "send a message to ourself": we copy from
      // the export buffer to the import buffer.  That saves
      // Teuchos::Comm implementations other than MpiComm (in
      // particular, SerialComm) the trouble of implementing self
      // messages correctly.  (To do this right, SerialComm would
      // need internal buffer space for messages, keyed on the
      // message's tag.)
      deep_copy_offset(imports, exports, selfReceiveOffset,
                       exportStarts[selfNum], exportLengths[selfNum]);
    }

  }
  else { // data are not blocked by proc, use send buffer
    ProfilingRegion psss("Tpetra::Distributor::doPostSends: MPI_Send SLOW");

    using Packet = typename ExpView::non_const_value_type;
    using Layout = typename ExpView::array_layout;
    using Device = typename ExpView::device_type;
    using Mem = typename ExpView::memory_traits;

    // This buffer is long enough for only one message at a time.
    // Thus, we use DISTRIBUTOR_SEND always in this case, regardless
    // of sendType requested by user.
    // This code path formerly errored out with message:
    //     Tpetra::Distributor::doPosts(3 args):
    //     The "send buffer" code path
    //     doesn't currently work with nonblocking sends.
    // Now, we opt to just do the communication in a way that works.
#ifdef HAVE_TPETRA_DEBUG
    if (sendType != Details::DISTRIBUTOR_SEND) {
      if (plan.getComm()->getRank() == 0)
        std::cout << "The requested Tpetra send type "
                  << DistributorSendTypeEnumToString(sendType)
                  << " requires Distributor data to be ordered by"
                  << " the receiving processor rank.  Since these"
                  << " data are not ordered, Tpetra will use Send"
                  << " instead." << std::endl;
    }
#endif

    size_t maxSendLength = 0;
    for (size_t i = 0; i < numBlocks; ++i) {
      size_t p = i + procIndex;
      if (p > (numBlocks - 1)) {
        p -= numBlocks;
      }

      size_t sendArrayOffset = 0;
        size_t j = plan.getStartsTo()[p];
        for (size_t k = 0; k < plan.getLengthsTo()[p]; ++k, ++j) {
        sendArrayOffset += exportLengths[j];
      }
      maxSendLength = std::max(maxSendLength, sendArrayOffset);
    }
    Kokkos::View<Packet*,Layout,Device,Mem> sendArray ("sendArray", maxSendLength);

    for (size_t i = 0; i < numBlocks; ++i) {
      size_t p = i + procIndex;
      if (p > (numBlocks - 1)) {
        p -= numBlocks;
      }

      if (plan.getProcsTo()[p] != myRank) {
        size_t sendArrayOffset = 0;
        size_t j = plan.getStartsTo()[p];
        for (size_t k = 0; k < plan.getLengthsTo()[p]; ++k, ++j) {
          packOffset(sendArray, exports, sendArrayOffset, exportStarts[j], exportLengths[j]);
          sendArrayOffset += exportLengths[j];
        }
        typename ExpView::execution_space().fence();

        ImpView tmpSend =
          subview_offset(sendArray, size_t(0), sendArrayOffset);

        send<int> (tmpSend,
            as<int> (tmpSend.size ()),
            plan.getProcsTo()[p], mpiTag_, *plan.getComm());
      }
      else { // "Sending" the message to myself
        selfNum = p;
        selfIndex = plan.getStartsTo()[p];
      }
    }

    if (plan.hasSelfMessage()) {
      for (size_t k = 0; k < plan.getLengthsTo()[selfNum]; ++k) {
        packOffset(imports, exports, selfReceiveOffset, exportStarts[selfIndex], exportLengths[selfIndex]);
        selfReceiveOffset += exportLengths[selfIndex];
        ++selfIndex;
      }
    }
  }
}

}

#endif