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); // Number of ranks across which the fields are decomposed. The larger this number, the less memory is used per rank. However, for best performance, this number should be as small as possible.
        n_phi_domains = nlr.get<int>("n_phi_domains",n_ranks); // Number of subdomains in the toroidal direction. By default, there is no poloidal decomposition.
        n_ghost_planes = nlr.get<int>("n_ghost_planes", 1); // Number of ghost planes on each side of a subdomain. Toroidal overlap is not required, but may be useful for performance since it reduces the need to shift particles between ranks.
        n_ghost_vertices = nlr.get<int>("n_ghost_vertices", 3000); // Number of ghost vertices on each side of the subdomain within the poloidal plane. It must be large enough to ensure that all neighboring vertices of the subdomain are present. However, this depends entirely on the grid since there is no guarantee on the order of vertices in an XGC grid.

        // 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