mpi_interface.c

Go to the documentation of this file.

 
#ifdef MPI
#include <mpi.h>
#endif
#include <stddef.h>
#include <stdlib.h>
#include "ascot5.h"
#include "diag.h"
#include "mpi_interface.h"
#include "particle.h"
#include "simulate.h"
 
void mpi_interface_init(int argc, char** argv, int* mpi_rank, int* mpi_size,
                        int* mpi_root) {
#ifdef MPI
    /* MPI is used. Init MPI and set rank and size as determined by MPI.
     * Rank of the root process is zero */
    int provided;
    int initialized;
    MPI_Initialized(&initialized);
    if(initialized == 0)
        MPI_Init_thread(&argc, &argv, MPI_THREAD_FUNNELED, &provided);
    MPI_Comm_rank(MPI_COMM_WORLD, mpi_rank);
    MPI_Comm_size(MPI_COMM_WORLD, mpi_size);
    *mpi_root = 0;
#else
    /* MPI is not used and no command line arguments are given. Assume that
     * the entire simulation is a single process. */
    *mpi_rank = 0;
    *mpi_size = 1;
    *mpi_root = 0;
#endif
}
 
void mpi_interface_finalize() {
#ifdef MPI
    MPI_Finalize();
#endif
}
 
void mpi_interface_barrier() {
#ifdef MPI
    MPI_Barrier(MPI_COMM_WORLD);
#endif
}
 
void mpi_my_particles(int* start_index, int* n, int n_tot, int mpi_rank,
                      int mpi_size) {
    if(mpi_rank == mpi_size-1) {
        *n = n_tot - mpi_rank * (n_tot / mpi_size);
    }
    else {
        *n = n_tot / mpi_size;
    }
 
    *start_index = mpi_rank * (n_tot / mpi_size);
}
 
void mpi_gather_particlestate(
    particle_state* ps, particle_state** ps_gather, int* n_gather, int n_tot,
    int mpi_rank, int mpi_size, int mpi_root) {
#ifdef MPI
    const int n_real = 32;
    const int n_int = 5;
    const int n_err = 1;
 
    particle_state* ps_all = malloc(n_tot * sizeof(particle_state));
 
    if(mpi_rank == mpi_root) {
        int start_index, n;
        mpi_my_particles(&start_index, &n, n_tot, mpi_rank, mpi_size);
 
        for(int j = 0; j < n; j++) {
            ps_all[j] = ps[j];
        }
 
        for(int i = 1; i < mpi_size; i++) {
            mpi_my_particles(&start_index, &n, n_tot, i, mpi_size);
 
            real* realdata;
            realdata = malloc(n_real * n * sizeof(realdata));
            integer* intdata;
            intdata = malloc(n_int * n * sizeof(intdata));
            a5err* errdata;
            errdata = malloc(n_err * n * sizeof(intdata));
 
            MPI_Recv(realdata, n_real*n, mpi_type_real, i, 0,
                MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Recv(intdata, n_int*n, mpi_type_integer, i, 0,
                MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Recv(errdata, n_err*n, mpi_type_a5err, i, 0,
                MPI_COMM_WORLD, MPI_STATUS_IGNORE);
 
            for(int j = 0; j < n; j++) {
                ps_all[start_index+j].r          = realdata[0*n+j];
                ps_all[start_index+j].phi        = realdata[1*n+j];
                ps_all[start_index+j].z          = realdata[2*n+j];
                ps_all[start_index+j].ppar       = realdata[3*n+j];
                ps_all[start_index+j].mu         = realdata[4*n+j];
                ps_all[start_index+j].zeta       = realdata[5*n+j];
                ps_all[start_index+j].rprt       = realdata[6*n+j];
                ps_all[start_index+j].phiprt     = realdata[7*n+j];
                ps_all[start_index+j].zprt       = realdata[8*n+j];
                ps_all[start_index+j].p_r        = realdata[9*n+j];
                ps_all[start_index+j].p_phi      = realdata[10*n+j];
                ps_all[start_index+j].p_z        = realdata[11*n+j];
                ps_all[start_index+j].mass       = realdata[12*n+j];
                ps_all[start_index+j].charge     = realdata[13*n+j];
                ps_all[start_index+j].anum       = intdata[0*n+j];
                ps_all[start_index+j].znum       = intdata[1*n+j];
                ps_all[start_index+j].weight     = realdata[14*n+j];
                ps_all[start_index+j].time       = realdata[15*n+j];
                ps_all[start_index+j].cputime    = realdata[16*n+j];
                ps_all[start_index+j].rho        = realdata[17*n+j];
                ps_all[start_index+j].theta      = realdata[18*n+j];
                ps_all[start_index+j].id         = intdata[2*n+j];
                ps_all[start_index+j].endcond    = intdata[3*n+j];
                ps_all[start_index+j].walltile   = intdata[4*n+j];
                ps_all[start_index+j].B_r        = realdata[19*n+j];
                ps_all[start_index+j].B_phi      = realdata[20*n+j];
                ps_all[start_index+j].B_z        = realdata[21*n+j];
                ps_all[start_index+j].B_r_dr     = realdata[22*n+j];
                ps_all[start_index+j].B_phi_dr   = realdata[23*n+j];
                ps_all[start_index+j].B_z_dr     = realdata[24*n+j];
                ps_all[start_index+j].B_r_dphi   = realdata[25*n+j];
                ps_all[start_index+j].B_phi_dphi = realdata[26*n+j];
                ps_all[start_index+j].B_z_dphi   = realdata[27*n+j];
                ps_all[start_index+j].B_r_dz     = realdata[28*n+j];
                ps_all[start_index+j].B_phi_dz   = realdata[29*n+j];
                ps_all[start_index+j].B_z_dz     = realdata[30*n+j];
                ps_all[start_index+j].mileage    = realdata[31*n+j];
                ps_all[start_index+j].err        = errdata[j];
            }
 
            free(realdata);
            free(intdata);
            free(errdata);
        }
    }
    else {
 
        int start_index, n;
        mpi_my_particles(&start_index, &n, n_tot, mpi_rank, mpi_size);
 
        real* realdata;
        realdata = malloc(n_real * n * sizeof(realdata));
        integer* intdata;
        intdata = malloc(n_int * n * sizeof(intdata));
        a5err* errdata;
        errdata = malloc(n_err * n * sizeof(intdata));
 
        for(int j = 0; j < n; j++) {
            realdata[0*n+j]  = ps[j].r;
            realdata[1*n+j]  = ps[j].phi;
            realdata[2*n+j]  = ps[j].z;
            realdata[3*n+j]  = ps[j].ppar;
            realdata[4*n+j]  = ps[j].mu;
            realdata[5*n+j]  = ps[j].zeta;
            realdata[6*n+j]  = ps[j].rprt;
            realdata[7*n+j]  = ps[j].phiprt;
            realdata[8*n+j]  = ps[j].zprt;
            realdata[9*n+j]  = ps[j].p_r;
            realdata[10*n+j] = ps[j].p_phi;
            realdata[11*n+j] = ps[j].p_z;
            realdata[12*n+j] = ps[j].mass;
            realdata[13*n+j] = ps[j].charge;
            intdata[0*n+j]   = ps[j].anum;
            intdata[1*n+j]   = ps[j].znum;
            realdata[14*n+j] = ps[j].weight;
            realdata[15*n+j] = ps[j].time;
            realdata[16*n+j] = ps[j].cputime;
            realdata[17*n+j] = ps[j].rho;
            realdata[18*n+j] = ps[j].theta;
            intdata[2*n+j]   = ps[j].id;
            intdata[3*n+j]   = ps[j].endcond;
            intdata[4*n+j]   = ps[j].walltile;
            realdata[19*n+j] = ps[j].B_r;
            realdata[20*n+j] = ps[j].B_phi;
            realdata[21*n+j] = ps[j].B_z;
            realdata[22*n+j] = ps[j].B_r_dr;
            realdata[23*n+j] = ps[j].B_phi_dr;
            realdata[24*n+j] = ps[j].B_z_dr;
            realdata[25*n+j] = ps[j].B_r_dphi;
            realdata[26*n+j] = ps[j].B_phi_dphi;
            realdata[27*n+j] = ps[j].B_z_dphi;
            realdata[28*n+j] = ps[j].B_r_dz;
            realdata[29*n+j] = ps[j].B_phi_dz;
            realdata[30*n+j] = ps[j].B_z_dz;
            realdata[31*n+j] = ps[j].mileage;
            errdata[j] = ps[j].err;
        }
 
        MPI_Send(realdata, n_real*n, mpi_type_real, 0, 0, MPI_COMM_WORLD);
        MPI_Send(intdata, n_int*n, mpi_type_integer, 0, 0, MPI_COMM_WORLD);
        MPI_Send(errdata, n_err*n, mpi_type_a5err, 0, 0, MPI_COMM_WORLD);
 
        free(realdata);
        free(intdata);
        free(errdata);
    }
 
    *ps_gather = ps_all;
    *n_gather = n_tot;
 
#else
 
    int start_index, n;
    mpi_my_particles(&start_index, &n, n_tot, mpi_rank, mpi_size);
 
    /* Only store particles for this process in Condor-style execution */
    particle_state* ps_all = malloc(n * sizeof(particle_state));
 
    for(int j = 0; j < n; j++) {
        ps_all[j] = ps[j];
    }
 
    *ps_gather = ps_all;
    *n_gather = n;
 
#endif
}
 
void mpi_gather_diag(diag_offload_data* data, real* offload_array, int ntotal,
                     int mpi_rank, int mpi_size, int mpi_root) {
#ifdef MPI
 
    if(data->dist5D_collect || data->distrho5D_collect
       || data->dist6D_collect || data->distrho6D_collect) {
        if(mpi_rank == mpi_root) {
            MPI_Reduce(MPI_IN_PLACE, offload_array,
                data->offload_dist_length, mpi_type_real, MPI_SUM,
                0, MPI_COMM_WORLD);
        } else {
            MPI_Reduce(offload_array, offload_array,
                data->offload_dist_length, mpi_type_real, MPI_SUM,
                0, MPI_COMM_WORLD);
        }
    }
 
    if(data->diagorb_collect) {
        if(mpi_rank == mpi_root) {
            for(int i = 1; i < mpi_size; i++) {
                int start_index, n;
                mpi_my_particles(&start_index, &n, ntotal, i, mpi_size);
 
                for(int j = 0; j < data->diagorb.Nfld; j++) {
                    MPI_Recv(&offload_array[data->offload_diagorb_index
                                        +j*data->diagorb.Nmrk*data->diagorb.Npnt
                                        +start_index*data->diagorb.Npnt],
                        n*data->diagorb.Npnt, mpi_type_real, i, 0,
                        MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                }
            }
        }
        else {
            int start_index, n;
            mpi_my_particles(&start_index, &n, ntotal, mpi_rank, mpi_size);
 
            for(int j = 0; j < data->diagorb.Nfld; j++) {
                MPI_Send(&offload_array[data->offload_diagorb_index
                                      +j*data->diagorb.Nmrk*data->diagorb.Npnt],
                n*data->diagorb.Npnt, mpi_type_real, 0, 0, MPI_COMM_WORLD);
            }
        }
    }
 
    if(data->diagtrcof_collect) {
        /* 3 fields for transport coefficients, id, D, K */
        int nfield = 3;
 
        if(mpi_rank == mpi_root) {
            for(int i = 1; i < mpi_size; i++) {
                int start_index, n;
                mpi_my_particles(&start_index, &n, ntotal, i, mpi_size);
 
                for(int j = 0; j < nfield; j++) {
                    MPI_Recv(&offload_array[data->offload_diagtrcof_index
                                        +j*data->diagtrcof.Nmrk
                                        +start_index],
                        n, mpi_type_real, i, 0,
                        MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                }
            }
        }
        else {
            int start_index, n;
            mpi_my_particles(&start_index, &n, ntotal, mpi_rank, mpi_size);
 
            for(int j = 0; j < nfield; j++) {
                MPI_Send(&offload_array[data->offload_diagtrcof_index
                                      +j*data->diagtrcof.Nmrk],
                n, mpi_type_real, 0, 0, MPI_COMM_WORLD);
            }
        }
    }
 
#endif
}