/**
 * @file xt_exchanger_irecv_isend_ddt_packed.c
 *
 * @copyright Copyright  (C)  2022 Jörg Behrens <behrens@dkrz.de>
 *                                 Moritz Hanke <hanke@dkrz.de>
 *                                 Thomas Jahns <jahns@dkrz.de>
 *
 * @author Jörg Behrens <behrens@dkrz.de>
 *         Moritz Hanke <hanke@dkrz.de>
 *         Thomas Jahns <jahns@dkrz.de>
 */
/*
 * Keywords:
 * Maintainer: Jörg Behrens <behrens@dkrz.de>
 *             Moritz Hanke <hanke@dkrz.de>
 *             Thomas Jahns <jahns@dkrz.de>
 * URL: https://dkrz-sw.gitlab-pages.dkrz.de/yaxt/
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are  permitted provided that the following conditions are
 * met:
 *
 * Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * Neither the name of the DKRZ GmbH nor the names of its contributors
 * may be used to endorse or promote products derived from this software
 * without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "core/ppm_xfuncs.h"
#include "xt/xt_mpi.h"
#include "xt_request_msgs_ddt_packed.h"
#include "xt_mpi_internal.h"
#include "xt_redist_internal.h"
#include "xt_exchanger_irecv_isend_ddt_packed.h"
#include "xt_exchanger_irecv_isend_ddt_packed_internal.h"
#include "xt_exchanger_simple_base.h"
#include "xt_exchanger_simple_base_internal.h"
#include "xt_exchanger_internal.h"
#include "xt_ddt_internal.h"

/* unfortunately GCC 11 cannot handle the literal constants used for
 * MPI_STATUSES_IGNORE by MPICH */
#if __GNUC__ >= 11 && __GNUC__ <= 13 && defined MPICH
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wstringop-overread"
#pragma GCC diagnostic ignored "-Wstringop-overflow"
#endif

enum {
  nvtab = 1,
};

struct xt_exchanger_vtable
xt_exchanger_irecv_isend_ddt_packed_vtable[nvtab];

static void
xt_exchanger_irecv_isend_ddt_packed_s_exchange(
  Xt_exchanger exchanger,
  const void *src_data, void *dst_data);

static void
xt_exchanger_irecv_isend_ddt_packed_a_exchange(
  Xt_exchanger exchangr,
  const void *src_data,
  void *dst_data,
  Xt_request *request);

void
xt_exchanger_irecv_isend_ddt_packed_init(void)
{
  xt_exchanger_irecv_isend_ddt_packed_vtable[0]
    = xt_exchanger_simple_base_vtable;
  xt_exchanger_irecv_isend_ddt_packed_vtable[0].s_exchange
    = xt_exchanger_irecv_isend_ddt_packed_s_exchange;
  xt_exchanger_irecv_isend_ddt_packed_vtable[0].a_exchange
    = xt_exchanger_irecv_isend_ddt_packed_a_exchange;
}

static void
xt_exchanger_irecv_isend_ddt_packed_s_exchange(
  Xt_exchanger exchanger,
  const void *src_data, void *dst_data)
{
  XT_GPU_INSTR_PUSH(xt_exchanger_irecv_isend_ddt_packed_s_exchange);
  Xt_exchanger_simple_base base = (Xt_exchanger_simple_base)exchanger;
  int nsend = base->nmsg[xt_direction_send],
    nrecv = base->nmsg[xt_direction_recv];
  enum { AUTO_ALLOC_SIZE = 32, };
  MPI_Request *requests, requests_auto[AUTO_ALLOC_SIZE];
  size_t *buffer_sizes, buffer_sizes_auto[AUTO_ALLOC_SIZE];
  Xt_ddt *restrict ddts, ddts_auto[AUTO_ALLOC_SIZE];

  size_t num_tx = (size_t)nrecv + (size_t)nsend;
  if (num_tx <= AUTO_ALLOC_SIZE) {
    requests = requests_auto;
    buffer_sizes = buffer_sizes_auto;
    ddts = ddts_auto;
  } else {
    requests = xmalloc(num_tx * sizeof (*requests));
    buffer_sizes = xmalloc(num_tx * sizeof (*buffer_sizes));
    ddts = xmalloc(num_tx * sizeof(*ddts));
  }

  size_t recv_buffer_size = 0;
  size_t send_buffer_size = 0;
  const struct Xt_msg_param *send_msgs
    = xt_exchanger_simple_base_get_send_msg(base),
    *recv_msgs = xt_exchanger_simple_base_get_recv_msg(base);
  for (int i = 0; i < nrecv; ++i)
    recv_buffer_size +=
      ((buffer_sizes[i] =
        xt_ddt_get_pack_size_internal(
          (size_t)recv_msgs[i].count,
          ((ddts[i] = xt_ddt_from_mpi_ddt(recv_msgs[i].datatype))))));
  for (int i = 0; i < nsend; ++i)
    send_buffer_size +=
      ((buffer_sizes[nrecv+i] =
        xt_ddt_get_pack_size_internal(
          (size_t)send_msgs[i].count,
          ((ddts[nrecv+i] = xt_ddt_from_mpi_ddt(send_msgs[i].datatype))))));

  enum xt_memtype src_data_memtype = xt_gpu_get_memtype(src_data);
  enum xt_memtype dst_data_memtype = xt_gpu_get_memtype(dst_data);
  unsigned char * send_buffer =
    xt_gpu_malloc(send_buffer_size, src_data_memtype);
  unsigned char * recv_buffer =
    xt_gpu_malloc(recv_buffer_size, dst_data_memtype);

  size_t ofs = 0;
  for (int i = 0; i < nrecv; ++i) {
    int recv_size = (int)buffer_sizes[i];
    xt_mpi_call(MPI_Irecv(
                  recv_buffer + ofs, recv_size, MPI_BYTE, recv_msgs[i].rank,
                  base->tag_offset + xt_mpi_tag_exchange_msg, base->comm,
                  requests+i), base->comm);
    ofs += (size_t)recv_size;
  }

#if XT_ENABLE_ALTERNATIVE_EXCHANGER_SEND_MODES
  typedef int (*ifp)(const void *buf, int count, MPI_Datatype datatype,
                     int dest, int tag, MPI_Comm comm, MPI_Request *request);
  enum Xt_exch_send_comm_mode send_mode
    = XT_CONFIG_FLAGS_GET_EXCH_SEND_COMM_MODE(base->config_flags);
  ifp a_send = (ifp)xt_send_func_tab[0][send_mode];
#else
#  define a_send MPI_Isend
#endif
  ofs = 0;
  for (int i = 0; i < nsend; ++i) {
    size_t send_size = buffer_sizes[nrecv+i];
    MPI_Aint extent, lb;
    xt_mpi_call(MPI_Type_get_extent(send_msgs[i].datatype, &lb, &extent), base->comm);
    xt_ddt_pack_internal(
      CAST_MPI_SEND_BUF(src_data), (size_t)send_msgs[i].count, ddts[nrecv+i],
      send_buffer + ofs, extent, src_data_memtype);
    xt_mpi_call(a_send(send_buffer + ofs, (int)send_size, MPI_BYTE,
                       send_msgs[i].rank,
                       base->tag_offset + xt_mpi_tag_exchange_msg, base->comm,
                       requests+nrecv+i), base->comm);
    ofs += send_size;
  }
#if ! XT_ENABLE_ALTERNATIVE_EXCHANGER_SEND_MODES
#  undef a_send
#endif

  xt_mpi_call(MPI_Waitall(nrecv + nsend, requests, MPI_STATUSES_IGNORE), base->comm);

  ofs = 0;
  for (int i = 0; i < nrecv; ++i) {
    size_t recv_size = buffer_sizes[i];
    MPI_Aint extent, lb;
    xt_mpi_call(MPI_Type_get_extent(recv_msgs[i].datatype, &lb, &extent), base->comm);
    xt_ddt_unpack_internal(
      recv_buffer + ofs, (size_t)recv_msgs[i].count, ddts[i],
      (unsigned char *)dst_data + recv_msgs[i].displacement,
      extent, dst_data_memtype);
    ofs += recv_size;
  }

  xt_gpu_free(recv_buffer, dst_data_memtype);
  xt_gpu_free(send_buffer, src_data_memtype);
  if (num_tx > AUTO_ALLOC_SIZE) {
    free(ddts);
    free(buffer_sizes);
    free(requests);
  }
  XT_GPU_INSTR_POP; // xt_exchanger_irecv_isend_ddt_packed_s_exchange
}

static void
xt_exchanger_irecv_isend_ddt_packed_a_exchange(
  Xt_exchanger exchanger,
  const void *src_data,
  void *dst_data,
  Xt_request *request)
{
  XT_GPU_INSTR_PUSH(xt_exchanger_irecv_isend_ddt_packed_a_exchange);
  Xt_exchanger_simple_base base = (Xt_exchanger_simple_base)exchanger;
  int nsend = base->nmsg[xt_direction_send],
    nrecv = base->nmsg[xt_direction_recv];
  /* todo: manually inline xt_request_msgs_ddt_packed_new to save
   * redundant allocations */
  MPI_Request * tmp_requests =
    xmalloc((size_t)(nrecv + nsend) * sizeof (*tmp_requests));
  void ** buffers =
    xmalloc((size_t)(nrecv + nsend) * sizeof (*buffers));
  MPI_Aint *recv_extents
    = xmalloc(2 * (size_t)nrecv * sizeof (*recv_extents)),
    *recv_displacements = recv_extents+nrecv;
  int *recv_counts = xmalloc(sizeof (*recv_counts) * (size_t)nrecv);
  enum xt_memtype src_data_memtype = xt_gpu_get_memtype(src_data);
  enum xt_memtype dst_data_memtype = xt_gpu_get_memtype(dst_data);

  Xt_ddt *recv_ddts = xmalloc((size_t)nrecv * sizeof(*recv_ddts));
  const struct Xt_msg_param *recv_msgs
    = xt_exchanger_simple_base_get_recv_msg(base);
  for (int i = 0; i < nrecv; ++i) {
    recv_ddts[i] = xt_ddt_from_mpi_ddt(recv_msgs[i].datatype);
    size_t buffer_size
      = xt_ddt_get_pack_size_internal((size_t)recv_msgs[i].count, recv_ddts[i]);
    buffers[i] = xt_gpu_malloc(buffer_size, dst_data_memtype);
    xt_mpi_call(MPI_Irecv(
                  buffers[i], (int)buffer_size, MPI_BYTE,
                  recv_msgs[i].rank,
                  base->tag_offset + xt_mpi_tag_exchange_msg, base->comm,
                  tmp_requests+i), base->comm);
  }

#if XT_ENABLE_ALTERNATIVE_EXCHANGER_SEND_MODES
  typedef int (*ifp)(const void *buf, int count, MPI_Datatype datatype,
                     int dest, int tag, MPI_Comm comm, MPI_Request *request);
  enum Xt_exch_send_comm_mode send_mode
    = XT_CONFIG_FLAGS_GET_EXCH_SEND_COMM_MODE(base->config_flags);
  ifp a_send = (ifp)xt_send_func_tab[0][send_mode];
#else
#  define a_send MPI_Isend
#endif
  const struct Xt_msg_param *send_msgs
    = xt_exchanger_simple_base_get_send_msg(base);
  for (int i = 0; i < nsend; ++i) {
    Xt_ddt send_ddt = xt_ddt_from_mpi_ddt(send_msgs[i].datatype);
    size_t buffer_size
      = xt_ddt_get_pack_size_internal((size_t)send_msgs[i].count, send_ddt);
    buffers[nrecv + i] = xt_gpu_malloc(buffer_size, src_data_memtype);
//! \todo merge all packing kernels into single kernel call -> less overhead,
//!       but not overlapping of packing and sending
    MPI_Aint extent, lb;
    xt_mpi_call(MPI_Type_get_extent(send_msgs[i].datatype, &lb, &extent), base->comm);
    xt_ddt_pack_internal(
      src_data, (size_t)send_msgs[i].count, send_ddt, buffers[nrecv + i],
      extent, src_data_memtype);
    xt_mpi_call(a_send(
                  buffers[nrecv + i], (int)buffer_size, MPI_BYTE,
                  send_msgs[i].rank,
                  base->tag_offset + xt_mpi_tag_exchange_msg, base->comm,
                  tmp_requests+nrecv+i), base->comm);
  }
#if ! XT_ENABLE_ALTERNATIVE_EXCHANGER_SEND_MODES
#  undef a_send
#endif
  for (int i = 0; i < nrecv; ++i) {
    MPI_Aint lb;
    xt_mpi_call(MPI_Type_get_extent(recv_msgs[i].datatype, &lb, recv_extents+i),
                base->comm);
    recv_displacements[i] = recv_msgs[i].displacement;
    recv_counts[i] = recv_msgs[i].count;
  }

  Xt_request requests =
    xt_request_msgs_ddt_packed_new(
      nrecv + nsend, tmp_requests, base->comm, nrecv, nsend,
      recv_ddts, recv_extents, recv_counts, buffers, buffers + nrecv, dst_data,
      src_data_memtype, dst_data_memtype);

  free(recv_counts);
  free(recv_extents);
  free(recv_ddts);
  free(buffers);
  free(tmp_requests);

  *request = requests;

  XT_GPU_INSTR_POP; // xt_exchanger_irecv_isend_ddt_packed_a_exchange
}

Xt_exchanger
xt_exchanger_irecv_isend_ddt_packed_new(int nsend, int nrecv,
                                        const struct Xt_msg_param *send_msgs,
                                        const struct Xt_msg_param *recv_msgs,
                                        MPI_Comm comm, int tag_offset,
                                        Xt_config config)
{
  /** note: tag_offset + xt_mpi_tag_exchange_msg must not
   *        be used on @a comm by any other part of the program during the
   *        lifetime of the created exchanger object
   */
  Xt_exchanger_simple_base exchanger =
    xt_exchanger_simple_base_alloc((size_t)nsend+(size_t)nrecv, 0);
  xt_exchanger_simple_base_init(
    exchanger, nsend, nrecv, send_msgs, recv_msgs,
    comm, tag_offset, config);
  exchanger->vtable
    = xt_exchanger_irecv_isend_ddt_packed_vtable;
  return (Xt_exchanger)exchanger;
}

/*
 * Local Variables:
 * c-basic-offset: 2
 * coding: utf-8
 * indent-tabs-mode: nil
 * show-trailing-whitespace: t
 * require-trailing-newline: t
 * End:
 */
