text: /*
   * (C) Copyright 2013 ECMWF.
   *
   * This software is licensed under the terms of the Apache Licence Version 2.0
   * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
   * In applying this licence, ECMWF does not waive the privileges and immunities
   * granted to it by virtue of its status as an intergovernmental organisation
   * nor does it submit to any jurisdiction.
   */
  
  /// @file HaloExchange.cc
  /// @author Willem Deconinck
  /// @date   Nov 2013
  
  #include <memory>
  #include <numeric>
  #include <sstream>
  #include <stdexcept>
  
  #include "atlas/array/Array.h"
  #include "atlas/parallel/HaloExchange.h"
  #include "atlas/parallel/mpi/Statistics.h"
  #include "atlas/util/vector.h"
  
  namespace atlas {
  namespace parallel {
  
  namespace {
  struct IsGhostPoint {
      IsGhostPoint(const int part[], const idx_t ridx[], const idx_t base, const int /*N*/) {
          part_   = part;
          ridx_   = ridx;
          base_   = base;
          mypart_ = mpi::rank();
      }
  
      bool operator()(idx_t idx) {
          if (part_[idx] != mypart_) {
              return true;
          }
          if (ridx_[idx] != base_ + idx) {
              return true;
          }
          return false;
      }
      int mypart_;
      const int* part_;
      const idx_t* ridx_;
      idx_t base_;
  };
  }  // namespace
  
  HaloExchange::HaloExchange(): name_(), is_setup_(false) {
      myproc = mpi::rank();
      nproc  = mpi::size();
  }
  
  HaloExchange::HaloExchange(const std::string& name): name_(name), is_setup_(false) {
      myproc = mpi::rank();
      nproc  = mpi::size();
  }
  
  HaloExchange::~HaloExchange() = default;
  
  void HaloExchange::setup(const int part[], const idx_t remote_idx[], const int base, const idx_t size) {
      setup(part, remote_idx, base, size, 0);
  }
  
  void HaloExchange::setup(const int part[], const idx_t remote_idx[], const int base, idx_t parsize, idx_t halo_begin) {
      ATLAS_TRACE("HaloExchange::setup");
  
      parsize_ = parsize;
      sendcounts_.resize(nproc);
      sendcounts_.assign(nproc, 0);
      recvcounts_.resize(nproc);
      recvcounts_.assign(nproc, 0);
      senddispls_.resize(nproc);
      senddispls_.assign(nproc, 0);
      recvdispls_.resize(nproc);
      recvdispls_.assign(nproc, 0);
  
      /*
      Find the amount of nodes this proc has to receive from each other proc
      */
  
      IsGhostPoint is_ghost(part, remote_idx, base, parsize_);
      atlas::vector<idx_t> ghost_points(parsize_);
      idx_t nghost = 0;
  
      atlas_omp_parallel_for(int jj = halo_begin; jj < parsize_; ++jj) {
          if (is_ghost(jj)) {
              int p = part[jj];
              atlas_omp_critical {
                  ++recvcounts_[p];
                  ghost_points[nghost] = jj;
                  nghost++;
              }
          }
      }
  
      recvcnt_ = std::accumulate(recvcounts_.begin(), recvcounts_.end(), 0);
  
      /*
      Find the amount of nodes this proc has to send to each other proc
      */
      ATLAS_TRACE_MPI(ALLTOALL) { mpi::comm().allToAll(recvcounts_, sendcounts_); }
  
      sendcnt_ = std::accumulate(sendcounts_.begin(), sendcounts_.end(), 0);
  
      recvdispls_[0] = 0;
      senddispls_[0] = 0;
      for (int jproc = 1; jproc < nproc; ++jproc)  // start at 1
      {
          recvdispls_[jproc] = recvcounts_[jproc - 1] + recvdispls_[jproc - 1];
          senddispls_[jproc] = sendcounts_[jproc - 1] + senddispls_[jproc - 1];
      }
      /*
      Fill vector "send_requests" with remote index of nodes needed, but are on
      other procs
      We can also fill in the vector "recvmap_" which holds local indices of
      requested nodes
      */
      std::vector<int> send_requests(recvcnt_);
      std::vector<int> recv_requests(sendcnt_);
      std::vector<int> cnt(nproc, 0);
      recvmap_.resize(recvcnt_);
  #ifdef __PGI
      // No idea why PGI compiler (20.7) in Release build ( -fast -O3 ) decides to vectorize following loop
  #pragma loop novector
  #endif
      for (idx_t jghost = 0; jghost < nghost; ++jghost) {
          const idx_t jj         = ghost_points[jghost];
          const int p            = part[jj];
          const int req_idx      = recvdispls_[p] + cnt[p];
          send_requests[req_idx] = remote_idx[jj] - base;
          recvmap_[req_idx]      = jj;
          cnt[p]++;
      }
  
      /*
      Fill vector "recv_requests" with what is needed by other procs
      */
  
      ATLAS_TRACE_MPI(ALLTOALL) {
          mpi::comm().allToAllv(send_requests.data(), recvcounts_.data(), recvdispls_.data(), recv_requests.data(),
                                sendcounts_.data(), senddispls_.data());
      }
  
      /*
      What needs to be sent to other procs is asked by remote_idx, which is local
      here
      */
      sendmap_.resize(sendcnt_);
      for (int jj = 0; jj < sendcnt_; ++jj) {
          sendmap_[jj] = recv_requests[jj];
      }
  
      is_setup_        = true;
      backdoor.parsize = parsize_;
  }
  
  void HaloExchange::wait_for_send(std::vector<int>& send_counts_init, std::vector<eckit::mpi::Request>& send_req) const {
      ATLAS_TRACE_MPI(WAIT, "mpi-wait send") {
          for (int jproc = 0; jproc < nproc; ++jproc) {
              if (send_counts_init[jproc] > 0) {
                  mpi::comm().wait(send_req[jproc]);
              }
          }
      }
  }
  
  namespace {
  
  template <typename Value>
  void execute_halo_exchange(HaloExchange* This, Value field[], int var_strides[], int var_extents[], int var_rank) {
      // WARNING: Only works if there is only one parallel dimension AND being
      // slowest moving
  
      array::ArrayShape shape{This->backdoor.parsize};
      for (int j = 0; j < var_rank; ++j) {
          shape.push_back(var_extents[j]);
      }
  
      array::ArrayStrides strides{var_extents[0] * var_strides[0]};
      for (int j = 0; j < var_rank; ++j) {
          strides.push_back(var_strides[j]);
      }
  
      std::unique_ptr<array::Array> arr(array::Array::wrap(field, array::ArraySpec{shape, strides}));
  
      switch (arr->rank()) {
          case 1: {
              This->execute<Value, 1>(*arr);
              break;
          }
          case 2: {
              This->execute<Value, 2>(*arr);
              break;
          }
          case 3: {
              This->execute<Value, 3>(*arr);
              break;
          }
          case 4: {
              This->execute<Value, 4>(*arr);
              break;
          }
          default:
              throw_NotImplemented("Rank not supported in halo exchange", Here());
      }
  }
  
  template <typename Value>
  void execute_adjoint_halo_exchange(HaloExchange* This, Value field[], int var_strides[], int var_extents[],
                                     int var_rank) {
      // WARNING: Only works if there is only one parallel dimension AND being
      // slowest moving
  
      array::ArrayShape shape{This->backdoor.parsize};
      for (int j = 0; j < var_rank; ++j) {
          shape.push_back(var_extents[j]);
      }
  
      array::ArrayStrides strides{var_extents[0] * var_strides[0]};
      for (int j = 0; j < var_rank; ++j) {
          strides.push_back(var_strides[j]);
      }
  
      std::unique_ptr<array::Array> arr(array::Array::wrap(field, array::ArraySpec{shape, strides}));
  
      switch (arr->rank()) {
          case 1: {
              This->execute_adjoint<Value, 1>(*arr);
              break;
          }
          case 2: {
              This->execute_adjoint<Value, 2>(*arr);
              break;
          }
          case 3: {
              This->execute_adjoint<Value, 3>(*arr);
              break;
          }
          case 4: {
              This->execute_adjoint<Value, 4>(*arr);
              break;
          }
          default:
              throw_NotImplemented("Rank not supported in halo exchange", Here());
      }
  }
  }  // namespace
  
  extern "C" {
  
  HaloExchange* atlas__HaloExchange__new() {
      return new HaloExchange();
  }
  
  void atlas__HaloExchange__delete(HaloExchange* This) {
      delete This;
  }
  
  void atlas__HaloExchange__setup(HaloExchange* This, int part[], idx_t remote_idx[], int base, int size) {
      This->setup(part, remote_idx, base, size);
  }
  
  
  void atlas__HaloExchange__execute_adjoint_strided_int(HaloExchange* This, int field[], int var_strides[],
                                                        int var_extents[], int var_rank) {
      execute_adjoint_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  
  void atlas__HaloExchange__execute_adjoint_strided_long(HaloExchange* This, long field[], int var_strides[],
                                                         int var_extents[], int var_rank) {
      execute_adjoint_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  
  void atlas__HaloExchange__execute_adjoint_strided_float(HaloExchange* This, float field[], int var_strides[],
                                                          int var_extents[], int var_rank) {
      execute_adjoint_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  
  void atlas__HaloExchange__execute_adjoint_strided_double(HaloExchange* This, double field[], int var_strides[],
                                                           int var_extents[], int var_rank) {
      execute_adjoint_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  
  void atlas__HaloExchange__execute_strided_int(HaloExchange* This, int field[], int var_strides[], int var_extents[],
                                                int var_rank) {
      execute_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  
  void atlas__HaloExchange__execute_strided_long(HaloExchange* This, long field[], int var_strides[], int var_extents[],
                                                 int var_rank) {
      execute_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  
  void atlas__HaloExchange__execute_strided_float(HaloExchange* This, float field[], int var_strides[], int var_extents[],
                                                  int var_rank) {
      execute_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  
  void atlas__HaloExchange__execute_strided_double(HaloExchange* This, double field[], int var_strides[],
                                                   int var_extents[], int var_rank) {
      execute_halo_exchange(This, field, var_strides, var_extents, var_rank);
  }
  }
  
  /////////////////////
  
  }  // namespace parallel
  }  // namespace atlas
  