grid_halo.hpp

Go to the documentation of this file.

// core — portable (Kokkos) GPU-resident ghost-layer exchange.
//
// Portable (Kokkos) GPU-resident grid halo: the field lives on the device as a
// peclet::core::View<T>; pack (gather send cells), unpack (scatter into ghost cells) and the
// periodic self-copy run as Kokkos::parallel_for on the default execution space (CUDA / HIP /
// OpenMP), so the full field never crosses the bus — only the compact halo buffers are staged to
// the host for MPI. The GPU-aware-MPI path (hand device pointers straight to MPI, dropping both
// staging copies) is the DEFAULT on device builds when a loopback probe confirms MPI accepts
// device pointers; env PECLET_CORE_GPU_AWARE_MPI=0 forces host staging (kill switch for stacks
// that segfault instead of erroring), =1 forces the device path without probing (legacy
// PECLET_CORE_CUDA_AWARE_MPI honoured). exchangeBegin/exchangeEnd expose the split so callers can
// overlap interior compute with the messages in flight; exchange() is the blocking composition.
// Topology comes from a host-built GridHaloTopology<Dim>::flatten(), and the result is bit-for-bit
// identical to the CPU exchange.
#ifndef PECLET_CORE_HALO_GRID_HALO_HPP
#define PECLET_CORE_HALO_GRID_HALO_HPP
 
#include <cstdio>
#include <cstdlib>
#include <string>
#include <type_traits>
#include <vector>
 
#include "peclet/core/common/mpi.hpp"
#if __has_include(<mpi-ext.h>)
#include <mpi-ext.h>  // OpenMPI: MPIX_CUDA_AWARE_SUPPORT / MPIX_Query_cuda_support
#endif
#include "peclet/core/common/types.hpp"
#include "peclet/core/common/view.hpp"
#include "peclet/core/halo/grid_halo_topology.hpp"
 
namespace peclet::core::halo {
 
namespace detail {
inline bool mpiReportsCudaAware() {
#if defined(MPIX_CUDA_AWARE_SUPPORT)
#if MPIX_CUDA_AWARE_SUPPORT
  return MPIX_Query_cuda_support() == 1;
#else
  return false;
#endif
#else
  return false;  // no query available: no signal
#endif
}
 
inline bool probeGpuAwareMpi() {
#ifdef KOKKOS_ENABLE_SERIAL
  if (std::is_same_v<ExecSpace, Kokkos::Serial>)
    return false;  // host backend: device pointers are host pointers; staging copies alias anyway
#endif
#ifdef KOKKOS_ENABLE_OPENMP
  if (std::is_same_v<ExecSpace, Kokkos::OpenMP>)
    return false;
#endif
#ifdef KOKKOS_ENABLE_THREADS
  if (std::is_same_v<ExecSpace, Kokkos::Threads>)
    return false;  // host backend, same as the two above — a wheel without a usable OpenMP gets
                   // this
#endif
  int inited = 0;
  MPI_Initialized(&inited);
  if (!inited)
    return false;
  MPI_Comm self;
  if (MPI_Comm_dup(MPI_COMM_SELF, &self) != MPI_SUCCESS)
    return false;
  MPI_Comm_set_errhandler(self, MPI_ERRORS_RETURN);
  constexpr int N = 64;
  View<int> src(Kokkos::view_alloc("peclet::core::halo::probeSrc", Kokkos::WithoutInitializing), N);
  View<int> dst("peclet::core::halo::probeDst", N);
  Kokkos::parallel_for(
      "peclet::core::halo::probeFill", Kokkos::RangePolicy<ExecSpace>(0, N),
      KOKKOS_LAMBDA(const int i) { src(i) = 0x5EC0DE + i * 7919; });
  Kokkos::fence();
  const int rc = MPI_Sendrecv(src.data(), N, MPI_INT, 0, 77, dst.data(), N, MPI_INT, 0, 77, self,
                              MPI_STATUS_IGNORE);
  MPI_Comm_free(&self);
  if (rc != MPI_SUCCESS)
    return false;
  int bad = 1;
  Kokkos::parallel_reduce(
      "peclet::core::halo::probeCheck", Kokkos::RangePolicy<ExecSpace>(0, N),
      KOKKOS_LAMBDA(const int i, int& b) {
        if (dst(i) != 0x5EC0DE + i * 7919)
          b += 1;
      },
      bad);
  return bad == 0;
}
 
inline bool gpuAwareMpi() {
  static const bool v = [] {
    const char* e = std::getenv("PECLET_CORE_GPU_AWARE_MPI");
    if (!e)
      e = std::getenv("PECLET_CORE_CUDA_AWARE_MPI");
    const bool aware = e ? (std::atoi(e) != 0) : (mpiReportsCudaAware() && probeGpuAwareMpi());
    if (std::getenv("PECLET_CORE_HALO_VERBOSE"))
      std::fprintf(stderr, "peclet.core halo: %s MPI buffers (%s)\n",
                   aware ? "DEVICE" : "host-staged", e ? "env-forced" : "auto");
    return aware;
  }();
  return v;
}
 
inline double haloTimeout() {
  static const double v = [] {
    const char* e = std::getenv("PECLET_CORE_HALO_TIMEOUT");
    return e ? std::atof(e) : 0.0;
  }();
  return v;
}
}  // namespace detail
 
template <class T>
class GridHalo {
 public:
  GridHalo() = default;
  GridHalo(const GridHalo&) = delete;
  GridHalo& operator=(const GridHalo&) = delete;
 
  template <int Dim>
  void init(const GridHaloTopology<Dim>& halo) {
    auto t = halo.flatten();
    comm_ = t.comm;
    sendRanks_ = t.sendRanks;
    recvRanks_ = t.recvRanks;
    sendCounts_ = t.sendCounts;
    recvCounts_ = t.recvCounts;
 
    sendOff_.assign(sendCounts_.size() + 1, 0);
    for (std::size_t k = 0; k < sendCounts_.size(); ++k)
      sendOff_[k + 1] = sendOff_[k] + sendCounts_[k];
    recvOff_.assign(recvCounts_.size() + 1, 0);
    for (std::size_t k = 0; k < recvCounts_.size(); ++k)
      recvOff_[k + 1] = recvOff_[k] + recvCounts_[k];
 
    nSend_ = static_cast<Index>(t.sendIdx.size());
    nRecv_ = static_cast<Index>(t.recvIdx.size());
    nSelf_ = static_cast<Index>(t.selfSrc.size());
 
    d_sendIdx_ = toDevice(t.sendIdx, "peclet::core::halo::sendIdx");
    d_recvIdx_ = toDevice(t.recvIdx, "peclet::core::halo::recvIdx");
    d_selfSrc_ = toDevice(t.selfSrc, "peclet::core::halo::selfSrc");
    d_selfDst_ = toDevice(t.selfDst, "peclet::core::halo::selfDst");
    d_sendBuf_ =
        View<T>(Kokkos::view_alloc("peclet::core::halo::sendBuf", Kokkos::WithoutInitializing),
                static_cast<std::size_t>(nSend_));
    d_recvBuf_ =
        View<T>(Kokkos::view_alloc("peclet::core::halo::recvBuf", Kokkos::WithoutInitializing),
                static_cast<std::size_t>(nRecv_));
    h_sendBuf_ = Kokkos::create_mirror_view(d_sendBuf_);
    h_recvBuf_ = Kokkos::create_mirror_view(d_recvBuf_);
  }
 
  void setLabel(std::string label) { label_ = std::move(label); }
 
  void exchange(const View<T>& field, int tag = 0) {
    exchangeBegin(field, tag);
    exchangeEnd(field);
  }
 
  void exchangeBegin(const View<T>& field, int tag = 0) {
    const bool aware = detail::gpuAwareMpi();
    // Periodic copy within our own block (read inner cell, write ghost cell).
    if (nSelf_) {
      View<T> f = field;
      IndexView src = d_selfSrc_, dst = d_selfDst_;
      Kokkos::parallel_for(
          "peclet::core::halo::selfCopy", Kokkos::RangePolicy<ExecSpace>(0, nSelf_),
          KOKKOS_LAMBDA(const Index i) { f(dst(i)) = f(src(i)); });
    }
    // Gather the cells we send into the contiguous send buffer.
    if (nSend_) {
      View<T> f = field;
      IndexView idx = d_sendIdx_;
      View<T> buf = d_sendBuf_;
      Kokkos::parallel_for(
          "peclet::core::halo::pack", Kokkos::RangePolicy<ExecSpace>(0, nSend_),
          KOKKOS_LAMBDA(const Index i) { buf(i) = f(idx(i)); });
      if (!aware)
        Kokkos::deep_copy(h_sendBuf_, d_sendBuf_);
    }
    // The send buffer (host-staged, or device for the aware path) must be ready before MPI reads
    // it.
    Kokkos::fence();
 
    T* sendBase = aware ? d_sendBuf_.data() : h_sendBuf_.data();
    T* recvBase = aware ? d_recvBuf_.data() : h_recvBuf_.data();
    reqs_.clear();
    reqs_.reserve(recvRanks_.size() + sendRanks_.size());
    for (std::size_t k = 0; k < recvRanks_.size(); ++k) {
      reqs_.emplace_back();
      MPI_Irecv(recvBase + recvOff_[k], recvCounts_[k] * static_cast<int>(sizeof(T)), MPI_BYTE,
                recvRanks_[k], tag, comm_, &reqs_.back());
    }
    for (std::size_t k = 0; k < sendRanks_.size(); ++k) {
      reqs_.emplace_back();
      MPI_Isend(sendBase + sendOff_[k], sendCounts_[k] * static_cast<int>(sizeof(T)), MPI_BYTE,
                sendRanks_[k], tag, comm_, &reqs_.back());
    }
    inFlightAware_ = aware;
  }
 
  void exchangeEnd(const View<T>& field) {
    if (!reqs_.empty()) {
      if (detail::haloTimeout() > 0)
        waitAllDiagnosed(detail::haloTimeout());
      else
        MPI_Waitall(static_cast<int>(reqs_.size()), reqs_.data(), MPI_STATUSES_IGNORE);
    }
    reqs_.clear();
    if (nRecv_) {
      if (!inFlightAware_)
        Kokkos::deep_copy(d_recvBuf_, h_recvBuf_);
      View<T> f = field;
      IndexView idx = d_recvIdx_;
      View<T> buf = d_recvBuf_;
      Kokkos::parallel_for(
          "peclet::core::halo::unpack", Kokkos::RangePolicy<ExecSpace>(0, nRecv_),
          KOKKOS_LAMBDA(const Index i) { f(idx(i)) = buf(i); });
    }
    Kokkos::fence();
  }
 
 private:
  MPI_Comm comm_ = MPI_COMM_NULL;
  std::vector<int> sendRanks_, recvRanks_, sendCounts_, recvCounts_;
  std::vector<int> sendOff_, recvOff_;
  std::vector<MPI_Request> reqs_;  // in-flight requests between exchangeBegin and exchangeEnd
  bool inFlightAware_ = false;     // staging mode of the in-flight exchange
  Index nSend_ = 0, nRecv_ = 0, nSelf_ = 0;
  std::string label_;
  IndexView d_sendIdx_, d_recvIdx_, d_selfSrc_, d_selfDst_;
 
  void waitAllDiagnosed(double timeout) {
    const double t0 = MPI_Wtime();
    const int n = static_cast<int>(reqs_.size());
    int done = 0;
    while (!done) {
      MPI_Testall(n, reqs_.data(), &done, MPI_STATUSES_IGNORE);
      if (done || MPI_Wtime() - t0 < timeout)
        continue;
      int rank = -1, size = -1, wrank = -1, wsize = -1;
      MPI_Comm_rank(comm_, &rank);
      MPI_Comm_size(comm_, &size);
      MPI_Comm_rank(MPI_COMM_WORLD, &wrank);
      MPI_Comm_size(MPI_COMM_WORLD, &wsize);
      MPI_Group g, wg;
      MPI_Comm_group(comm_, &g);
      MPI_Comm_group(MPI_COMM_WORLD, &wg);
      auto world = [&](int r) {
        int w = -1;
        MPI_Group_translate_ranks(g, 1, &r, wg, &w);
        return w;
      };
      std::string msg = "peclet.core halo TIMEOUT after " + std::to_string(int(timeout)) + " s [" +
                        label_ + "] world " + std::to_string(wrank) + "/" + std::to_string(wsize) +
                        " comm " + std::to_string(rank) + "/" + std::to_string(size) + ": " +
                        std::to_string(recvRanks_.size()) + " recv + " +
                        std::to_string(sendRanks_.size()) + " send partners\n";
      const std::size_t nr = recvRanks_.size();
      for (std::size_t k = 0; k < reqs_.size(); ++k) {
        int flag = 0;
        MPI_Request_get_status(reqs_[k], &flag, MPI_STATUS_IGNORE);
        if (flag)
          continue;
        const bool recv = k < nr;
        const int q = recv ? recvRanks_[k] : sendRanks_[k - nr];
        const int c = recv ? recvCounts_[k] : sendCounts_[k - nr];
        msg += std::string("   pending ") + (recv ? "RECV from" : "SEND to  ") + " comm " +
               std::to_string(q) + " (world " + std::to_string(world(q)) + ") " +
               std::to_string(c * static_cast<long>(sizeof(T))) + " bytes\n";
      }
      std::fputs(msg.c_str(), stderr);
      std::fflush(stderr);
      MPI_Group_free(&g);
      MPI_Group_free(&wg);
      MPI_Abort(MPI_COMM_WORLD, 77);
    }
  }
  View<T> d_sendBuf_, d_recvBuf_;
  HostView<T> h_sendBuf_, h_recvBuf_;
};
 
}  // namespace peclet::core::halo
 
#endif  // PECLET_CORE_HALO_GRID_HALO_HPP