field_decomposition.hpp

Go to the documentation of this file.

#ifndef FIELD_DECOMPOSITION_HPP
#define FIELD_DECOMPOSITION_HPP
#include "space_settings.hpp"
#include "NamelistReader.hpp"
#include "globals.hpp"
#ifdef USE_MPI
#  include "my_mpi.hpp"
#endif

// Field decomposition class
template<class Device>
class FieldDecomposition {
    public:

#ifdef USE_MPI
    // MPI communicators etc
    MyMPI mpi;
#endif

    // Constants
    int n_ranks;   
    int n_phi_domains;     
    int n_pol_domains; 
    int n_ghost_planes; 
    int n_ghost_vertices; 

    int first_owned_node; 
    int nnodes_owned; 
    int first_owned_plane; 
    int nplanes_owned; 

    int first_node; 
    int last_node; 
    int n_nodes; 
    int first_plane; 
    int last_plane; 
    int n_planes; 

    // Views
    View<int*,CLayout,HostType> map_from_global_intpl; 
    View<int*,CLayout,HostType> all_first_node; 
    View<int*,CLayout,HostType> all_last_node; 
    View<int*,CLayout,HostType> all_first_plane; 
    View<int*,CLayout,HostType> all_last_plane; 

    FieldDecomposition(){}

    // Constructor
    FieldDecomposition(NLReader::NamelistReader& nlr, int nplanes, int nnodes){
#ifdef USE_MPI
        nlr.use_namelist("field_decomp_param");
        n_ranks = nlr.get<int>("n_ranks", 6);
        n_phi_domains = nlr.get<int>("n_phi_domains",n_ranks); // Default: no poloidal decomposition
        n_ghost_planes = nlr.get<int>("n_ghost_planes", 1);
        n_ghost_vertices = nlr.get<int>("n_ghost_vertices", 3000);

        // If there is no phi decomposition, don't use ghost planes
        if(n_phi_domains==1) n_ghost_planes = 0;

        n_pol_domains = n_ranks/n_phi_domains; // Number of poloidal domains is inferred

        // Temporarily create a new com which is SML_COMM split into contiguous sets of "n_ranks" processes
        MPI_Comm comm;
        MPI_Comm_split(SML_COMM_WORLD, SML_COMM_RANK/n_ranks, SML_COMM_RANK, &comm);
        // Split up comm using same class as standard decomposition (mpi.comm, mpi.plane_comm and mpi.intpl_comm)
        mpi = MyMPI(comm, n_phi_domains);
printf("\nrank %d mpi.my_rank=%d", SML_COMM_RANK, mpi.my_rank);

        all_first_node = View<int*,CLayout,HostType>("all_first_node",n_ranks);
        all_last_node = View<int*,CLayout,HostType>("all_last_node",n_ranks);
        all_first_plane = View<int*,CLayout,HostType>("all_first_plane",n_ranks);
        all_last_plane = View<int*,CLayout,HostType>("all_last_plane",n_ranks);
        
        // Some checks
        if(n_pol_domains*n_phi_domains != n_ranks)
            exit_XGC("\nError in field_decomposition: n_ranks must be divisible by n_phi_domains");

        // This restriction could be relaxed
        int nplanes_per_phi_domain = nplanes/n_phi_domains;
        if(nplanes_per_phi_domain*n_phi_domains != nplanes)
            exit_XGC("\nError in field_decomposition: n_phi_domains must divide evenly into nplanes");

        if(mpi.nranks != n_ranks)
            exit_XGC("\nError in field_decomposition: total XGC ranks must be divisible by field_decomp_param n_ranks");

        // Go through each domain and determine its size/offset
        // Planes are contiguous in the communicator so they are the inner loop
        // Everything here is ZERO-INDEXED
        int nodes_per_pol_domain = nnodes/n_pol_domains; // FLOOR
        int nodes_per_phi_domain = nplanes/n_phi_domains; // FLOOR
        int i = 0;
        for(int i_pol=0; i_pol<n_pol_domains; i_pol++){
            int first_owned_node_r = nodes_per_pol_domain*i_pol;
            int n_owned_nodes_r = (i_pol==n_pol_domains-1 ? (nnodes - first_owned_node_r) : nodes_per_pol_domain);
            for(int i_phi=0; i_phi<n_phi_domains; i_phi++){
                int first_owned_plane_r = nodes_per_phi_domain*i_phi;
                int n_owned_planes_r = (i_phi==n_phi_domains-1 ? (nplanes - first_owned_plane_r) : nodes_per_phi_domain);

                // Add modulo'd ghost cells for phi, cut off ghost cells for pol
                all_first_node(i) = std::max(first_owned_node_r - n_ghost_vertices, 0);
                all_last_node(i) = std::min(first_owned_node_r + n_owned_nodes_r - 1 + n_ghost_vertices, nnodes-1);

                // If nplanes ends up being the entire domain
                if((n_owned_planes_r + 2*n_ghost_planes) >= nplanes){
                    all_first_plane(i) = 0;
                    all_last_plane(i) = nplanes-1;
                }else{
                    all_first_plane(i) = positive_modulo(first_owned_plane_r - n_ghost_planes,nplanes);
                    all_last_plane(i) = (first_owned_plane_r + n_owned_planes_r - 1 + n_ghost_planes)%nplanes;
                }
printf("\nall_first_plane(%d) = (first_owned_plane_r - n_ghost_planes)modnplanes = (%d - %d)mod %d = %d \n", i, first_owned_plane_r, n_ghost_planes, nplanes, all_first_plane(i));


                // Finally, set for this rank
                if(i==mpi.my_rank){
                    first_owned_node = first_owned_node_r;
                    nnodes_owned = n_owned_nodes_r;
                    first_owned_plane = first_owned_plane_r;
                    nplanes_owned = n_owned_planes_r;

                    first_node = all_first_node(i);
                    last_node = all_last_node(i);
                    n_nodes = last_node - first_node + 1;
                    first_plane = all_first_plane(i);
                    last_plane = all_last_plane(i);
                    n_planes = positive_modulo(last_plane - first_plane,nplanes) + 1;
                }

                // Advance to next rank
                i++;
            }
        }
#endif
    }

    KOKKOS_INLINE_FUNCTION int find_domain_owner(int global_plane_index, int nplanes_total, int global_node_index, int nnodes_total) const{
        // These are enforced to be divisible
        int planes_per_phi_domain = nplanes_total/n_phi_domains;
        int intpl_pid = global_plane_index/planes_per_phi_domain;

        // Poloidal decomposition should work out despite not necessarily being divisible
        int n_vertices_per_pol_domain = nnodes_total/n_pol_domains;
        int plane_pid = global_node_index/n_vertices_per_pol_domain;

        // calculate global pid from plane and intpl coordinates
        return (intpl_pid + n_phi_domains*plane_pid);
    }

    int all_n_nodes(int local_pid) const{
        return all_last_node(local_pid) - all_first_node(local_pid) + 1;
    }

    int all_n_planes(int local_pid, int nplanes) const{
        return positive_modulo(all_last_plane(local_pid) - all_first_plane(local_pid), nplanes) + 1;
    }
};

#endif