yaxt_c_plugin.c

Go to the documentation of this file.

/*
  --------------------------------------------------------------------
  Example plugin for the ICON Community Interface (ComIn)
  to demonstrate the usage of YAXT.
 
  @authors 08/2021 :: ICON Community Interface  <icon@dwd.de>
 
  Note that in order to demonstrate ComIn's language interoperability,
  a similary plugin has been implemented in Fortran, see the subdirectory
  "yaxt_fortran".
  --------------------------------------------------------------------
 */
 
#include <stdio.h>
#include <string.h>
 
#include <comin.h>
#include <mpi.h>
#include <yaxt.h>
 
#define ASSERT(COND)                                                           \
  if (!(COND)) {                                                               \
    fprintf(stderr, "assertion failed: %s\n", #COND);                          \
    MPI_Abort(MPI_COMM_WORLD, 1);                                              \
  }
 
const int jg                 = 1;
int rank                     = -1;
MPI_Comm plugin_comm         = MPI_COMM_NULL;
t_comin_var_handle* temp_var = NULL;
Xt_int* idxlist              = NULL;
double* buffer               = NULL;
int ncells                   = 0;
int prognostic_cells         = 0;
Xt_redist yaxt_redist;
const int* decomp_domain = NULL;
 
void secondary_constructor();
void offload();
void destructor();
 
void comin_main() {
  rank = comin_parallel_get_host_mpi_rank();
  comin_callback_register(EP_SECONDARY_CONSTRUCTOR, secondary_constructor);
 
  comin_callback_register(EP_ATM_WRITE_OUTPUT_BEFORE, offload);
 
  comin_callback_register(EP_DESTRUCTOR, destructor);
 
  plugin_comm = MPI_Comm_f2c(comin_parallel_get_plugin_mpi_comm());
 
  if (!xt_initialized())
    xt_initialize(plugin_comm);
 
  ncells = comin_descrdata_get_domain_cells_ncells(jg);
 
  int decomp_domain_size[2] = {0, 0};
  comin_descrdata_get_domain_cells_decomp_domain(jg, &decomp_domain,
                                                 decomp_domain_size);
 
  const int* glb_index = NULL;
  int glb_index_size   = 0;
  comin_descrdata_get_domain_cells_glb_index(jg, &glb_index, &glb_index_size);
 
  ASSERT(decomp_domain_size[0] * decomp_domain_size[1] > glb_index_size);
 
  idxlist          = malloc(ncells * sizeof(*idxlist));
  prognostic_cells = 0;
  for (int i = 0; i < ncells; ++i) {
    if (decomp_domain[i] == 0) {
      idxlist[prognostic_cells] = glb_index[i];
      prognostic_cells++;
    }
  }
  Xt_idxlist idxvec = xt_idxvec_new(idxlist, prognostic_cells);
 
  Xt_idxlist empty_idxlist = xt_idxempty_new();
 
  Xt_xmap xmap = xt_xmap_all2all_new(idxvec, empty_idxlist, plugin_comm);
 
  yaxt_redist = xt_redist_p2p_new(xmap, MPI_DOUBLE);
  xt_xmap_delete(xmap);
  xt_idxlist_delete(empty_idxlist);
  xt_idxlist_delete(idxvec);
 
  buffer = malloc(prognostic_cells * sizeof(*buffer));
}
 
void secondary_constructor() {
  t_comin_entry_point ep = EP_ATM_WRITE_OUTPUT_BEFORE;
  t_comin_var_descriptor var_desc;
  strcpy(var_desc.name, "temp");
  var_desc.id = 1;
  temp_var    = comin_var_get(1, &ep, var_desc, COMIN_FLAG_READ);
}
 
void offload() {
  // this callback is only called in entrypoints outside of domain loops
  ASSERT(comin_current_get_domain_id() == COMIN_DOMAIN_OUTSIDE_LOOP);
  double* temp_ptr = comin_var_get_ptr_double(temp_var); // ensures [jc, jb, jk]
  int shape[5];
  comin_var_get_shape(temp_var, shape);
  int nlev         = comin_descrdata_get_domain_nlev(jg);
  prognostic_cells = 0;
  for (int i = 0; i < ncells; ++i) {
    if (decomp_domain[i] == 0) {
      int idx0 = i % shape[0];
      int idx2 = i / shape[0];
      buffer[prognostic_cells] =
          temp_ptr[idx0 + nlev * shape[0] + idx2 * shape[0] * shape[1]];
      prognostic_cells++;
    }
  }
 
  xt_redist_s_exchange1(yaxt_redist, buffer, NULL);
}
 
void destructor() {
  free(idxlist);
  free(buffer);
}