plane_field_gatherer.hpp

Go to the documentation of this file.

#ifndef PLANE_FIELD_GATHERER_HPP
#define PLANE_FIELD_GATHERER_HPP

#include "domain_decomposition.hpp"
#include "grid.hpp"
#include "my_subview.hpp"
#include "checkpoint.hpp"

class PlaneFieldGatherer{
    const DomainDecomposition<DeviceType> pol_decomp;
    const Grid<DeviceType> grid;

    bool gather_subset;
    int nplanes;
    int plane_offset;
    int nnode;
    int node_offset;

    int tmp_nphi;
    View<Field<VarType::Scalar,PhiInterpType::Planes>**, CLayout,HostType> tmp_s; // for scalar fields
    View<Field<VarType::Vector,PhiInterpType::Planes>**, CLayout,HostType> tmp_v; // for vector fields

    View<Field<VarType::Scalar,PhiInterpType::Planes>*, CLayout,HostType> tmp_s_full; // for scalar fields
    View<Field<VarType::Vector,PhiInterpType::Planes>*, CLayout,HostType> tmp_v_full; // for vector fields

    /* USE_GPU without USE_GPU_AWARE_MPI requires that a temporary array be allocated
     */
    int choose_tmp_nphi(){
#if defined(USE_GPU) && !defined(USE_GPU_AWARE_MPI)
        return nplanes;
#else
        return 0;
#endif
    }

    /* // Series of host MPI_Igatherv to populate phi_ff. An MPI_Allgatherv or other op (or intercommunicator?) might be possible/better.
     */
    template<typename T, typename FT>
    void allgather_to_local_ranks(View<FT*, CLayout,HostType>& tmp_full,T* destination){
#ifdef USE_MPI
        int one_obj_in_dbl = sizeof(FT)/8; // from bytes to doubles

        View<int*, HostType> recvcounts(NoInit("recvcounts"), pol_decomp.mpi.n_intpl_ranks);
        View<int*, HostType> displs(NoInit("displs"),pol_decomp.mpi.n_intpl_ranks);
        std::vector<MPI_Request> rrequests(pol_decomp.mpi.n_intpl_ranks);
//printf("\nrank %d, both map: %d, %d \n", pol_decomp.mpi.my_rank, pol_decomp.field_decomp.map_from_global_intpl(0), pol_decomp.field_decomp.map_from_global_intpl(1));
        // Loop over ranks in global inter-planar communicator, each taking turns as root
        for(int i_root=0; i_root<pol_decomp.mpi.n_intpl_ranks; i_root++){
            int root_local_rank = pol_decomp.field_decomp.map_from_global_intpl(i_root); // will range from e.g. 0 to 7
            int root_local_first_plane = pol_decomp.field_decomp.all_first_plane(root_local_rank);
            int root_local_last_plane = pol_decomp.field_decomp.all_last_plane(root_local_rank);
            int root_local_nplanes = positive_modulo(root_local_last_plane - root_local_first_plane,grid.nplanes) + 1;

            int root_local_first_node = pol_decomp.field_decomp.all_first_node(root_local_rank);
            int root_local_nnode = pol_decomp.field_decomp.all_last_node(root_local_rank) - pol_decomp.field_decomp.all_first_node(root_local_rank) + 1;
            int size = root_local_nnode * one_obj_in_dbl;

            // Now determine where each contributing rank will send its contribution
            // Initialize to zero
            Kokkos::deep_copy(recvcounts, 0);
            Kokkos::deep_copy(displs, 0);

            // Note that the data arrives OUT OF ORDER to handle the case where the ghost planes wrap around plane 0
            for(int i=0; i<root_local_nplanes; i++){
                int contributor = (i+root_local_first_plane)%grid.nplanes;
                recvcounts(contributor) = size;
                displs(contributor) = i*size;
                //if(i<displs.size()-1) displs(i+1) = displs(i) + recvcounts(contributor);
            }

            int my_size = recvcounts(pol_decomp.mpi.my_intpl_rank);

/*checkpoint("");
checkpoint("");
for(int ipr = 0; ipr<pol_decomp.mpi.nranks; ipr++){
            if(pol_decomp.mpi.my_rank==ipr){
                printf("\nRANK %d, plane_rank %d, intpl_rank %d:", ipr, pol_decomp.mpi.my_plane_rank, pol_decomp.mpi.my_intpl_rank);
                printf("\nnode_offset: %d, roots: %d, root_local_nnode: %d", node_offset, root_local_first_node, root_local_nnode);
                printf("\nsize: %d, my_size: %d", size, my_size);
                printf("\nroot(%d) rank: %d, first planes: %d, last plane: %d, n_planes: %d", i_root, root_local_rank, root_local_first_plane, root_local_last_plane, root_local_nplanes);
                printf("\nrecvcounts = :");
                for(int i=0; i<recvcounts.size(); i++){
                    printf(" %d", recvcounts(i));
                }
                printf("\ndispl = :");
                for(int i=0; i<displs.size(); i++){
                    printf(" %d", displs(i));
                }
            }
checkpoint("");
checkpoint("");
}*/

            MPI_Igatherv(tmp_full.data()+root_local_first_node, my_size, MPI_DOUBLE, destination, recvcounts.data(), displs.data(), MPI_DOUBLE, i_root, pol_decomp.mpi.intpl_comm, &(rrequests[i_root]));
//MPI_Wait(&(rrequests[i_root]), MPI_STATUS_IGNORE);

//checkpoint("Finished delivering to an intpl rank");
        }
        // Wait for Igatherv to complete
        for(int i_root=0; i_root<pol_decomp.mpi.n_intpl_ranks; i_root++) MPI_Wait(&(rrequests[i_root]), MPI_STATUS_IGNORE);
#endif
    }

    /* Copy the irho==0 index of a rho_ff field
     */
    template<typename FT>
    void copy_to_tmp_full(View<FT*, CLayout,HostType>& tmp_full, const View<FT**, CLayout,HostType>& rho_ff_h){
        constexpr int ZERO_GYRO = 0;

        // Check that the host array is allocated
        if(rho_ff_h.extent(0) != grid.nnode || rho_ff_h.extent(1)<1) exit_XGC("\nError in gather_phi_ff_on_device: expected host view of size (nnode,nrho)\n");

        // Copy this rank's contribution to tmp_full
        Kokkos::parallel_for("gather_field_info", Kokkos::RangePolicy<HostExSpace>( 0, grid.nnode), [=] (const int inode){
            tmp_full(inode)=rho_ff_h(inode,ZERO_GYRO);
        });
    }

    /* Copy a ff field
     */
    template<typename FT>
    void copy_to_tmp_full(View<FT*, CLayout,HostType>& tmp_full, const View<FT*, CLayout,HostType>& ff_h){
        // Check that the host array is allocated
        if(ff_h.extent(0) != grid.nnode) exit_XGC("\nError in gather_phi_ff_on_device: expected host view of size (nnode,nrho)\n");

        // Copy this rank's contribution to tmp_full
        Kokkos::parallel_for("gather_field_info", Kokkos::RangePolicy<HostExSpace>( 0, grid.nnode), [=] (const int inode){
            tmp_full(inode)=ff_h(inode);
        });
    }

    template<typename T_h, typename FT>
    void gather_phi_ff_on_device(View<FT**, CLayout,HostType>& tmp, View<FT*, CLayout,HostType>& tmp_full, const T_h& rho_ff_h, View<FT**, CLayout,DeviceType>& phi_ff){
#if defined(USE_GPU) && !defined(USE_GPU_AWARE_MPI)
        constexpr bool gpu_without_gpu_aware_MPI = true;
#else
        constexpr bool gpu_without_gpu_aware_MPI = false;
#endif

        // Copy this rank's contribution to tmp_full
        copy_to_tmp_full(tmp_full, rho_ff_h);

        // Allocate phi_ff device view
        phi_ff = View<FT**, CLayout,DeviceType>("gfpack_view",nplanes, nnode);

        // If using GPUs without GPU-aware MPI, need to do MPI gather into a temporary host allocation
        auto destination = (gpu_without_gpu_aware_MPI ? tmp.data() : phi_ff.data());
//checkpoint("");
//checkpoint("");
#ifdef USE_MPI
/*for(int ipr = 0; ipr<pol_decomp.mpi.nranks; ipr++){
            if(pol_decomp.mpi.my_rank==ipr){
                printf("\nRANK %d:", ipr);
                printf("\nphi_ff(%d, %d)\n", phi_ff.extent(0), phi_ff.extent(1));
                printf("\ngpu_without_gpu_aware_MPI=%d\n", gpu_without_gpu_aware_MPI?1 : 0);
                printf("\ntmp(%d, %d)\n", tmp.extent(0), tmp.extent(1));
                printf("\ntmp_full(%d)\n", tmp_full.extent(0));
            }
checkpoint("");
checkpoint("");
}*/
#endif
        // Gather contributions from each rank; if we're just gathering a subset of the field, the logic is more complicated
        if(gather_subset){
            allgather_to_local_ranks(tmp_full, destination);
        }else{
#ifdef USE_MPI
            int one_obj_in_dbl = sizeof(FT)/8; // from bytes to doubles
            int sz=nnode * one_obj_in_dbl;
            MPI_Allgather(tmp_full.data(), sz, MPI_DOUBLE, destination, sz, MPI_DOUBLE, pol_decomp.mpi.intpl_comm);
#else
            // Used for mini_appKernel; just replicate the rho field for now
            if(gpu_without_gpu_aware_MPI)
                for(int i=0; i<grid.nplanes; i++) Kokkos::deep_copy(my_subview(tmp, i), tmp_full);
            else
                for(int i=0; i<grid.nplanes; i++) Kokkos::deep_copy(my_subview(phi_ff, i), tmp_full);
#endif
        }
//checkpoint("Entering copy");
        // If using GPUs without GPU-aware MPI, copy from the temporary allocation to device memory
        if(gpu_without_gpu_aware_MPI){
            Kokkos::deep_copy(phi_ff, tmp);
        }
//checkpoint("Finished copy");
    }

    template<typename FT>
    inline View<FT**, CLayout,HostType>& which_tmp();

    template<typename FT>
    inline View<FT*, CLayout,HostType>& which_tmp_full();

    public:

    PlaneFieldGatherer(const DomainDecomposition<DeviceType>& pol_decomp, const Grid<DeviceType>& grid)
      : pol_decomp(pol_decomp),
        grid(grid),
        gather_subset(pol_decomp.decompose_fields),
        // If using field_decomposition, we don't copy the entire field
        nplanes(gather_subset ? pol_decomp.field_decomp.n_planes : grid.nplanes),
        plane_offset(gather_subset ? pol_decomp.field_decomp.first_plane : 0),
        nnode(gather_subset ? pol_decomp.field_decomp.n_nodes : grid.nnode),
        node_offset(gather_subset ? pol_decomp.field_decomp.first_node : 0),
        tmp_nphi(choose_tmp_nphi()),
        tmp_s("tmp_s", tmp_nphi, nnode),
        tmp_v("tmp_v", tmp_nphi, nnode),
        tmp_s_full("tmp_s_full", grid.nnode),
        tmp_v_full("tmp_v_full", grid.nnode)
        {
//printf("\nplane_offset: %d, node_offset: %d\n", plane_offset, node_offset);
//printf("\ntmp_nphi: %d, tmp_s(%d, %d), tmp_v(%d, %d)\n", tmp_nphi, tmp_s.extent(0), tmp_s.extent(1), tmp_v.extent(0), tmp_v.extent(1));
}

    // Interface to use correct tmp
    template<typename T_h, typename FT>
    void gather_phi_ff_on_device(const T_h& rho_ff_h, View<FT**, CLayout,DeviceType>& phi_ff){
        gather_phi_ff_on_device(which_tmp<FT>(), which_tmp_full<FT>(), rho_ff_h, phi_ff);
    }
};

template<>
inline View<Field<VarType::Scalar,PhiInterpType::Planes>**, CLayout,HostType>& PlaneFieldGatherer::which_tmp(){
    return tmp_s;
}

template<>
inline View<Field<VarType::Vector,PhiInterpType::Planes>**, CLayout,HostType>& PlaneFieldGatherer::which_tmp(){
    return tmp_v;
}

template<>
inline View<Field<VarType::Scalar,PhiInterpType::Planes>*, CLayout,HostType>& PlaneFieldGatherer::which_tmp_full(){
    return tmp_s_full;
}

template<>
inline View<Field<VarType::Vector,PhiInterpType::Planes>*, CLayout,HostType>& PlaneFieldGatherer::which_tmp_full(){
    return tmp_v_full;
}

#endif