poloidal_sum.hpp

Go to the documentation of this file.

#ifndef POLOIDAL_SUM_HPP
#define POLOIDAL_SUM_HPP

// Does a copy to the MPI memory space if different from the View's memory space
template<class T>
void poloidal_sum_in_place(const DomainDecomposition<DeviceType>& pol_decomp, T& f){
    // No operation if n_planes is 1
    if(pol_decomp.mpi.n_plane_ranks==1) return;

    // Cast to 1D with an unmanaged view for a generic element-wise operation
    View<double*,CLayout, typename T::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_1D((double*)(f.data()), f.size());

    using MPIHostType = HostType; // Do this MPI_Allreduce on Host; GPU-aware MPI consistently hangs or aborts here

    // Copy to host if on GPU and GPU-aware MPI is not available
    auto f_mpi_d = my_mirror_scratch_view(f_1D,MPIHostType());
    auto f_mpi = my_mirror_view(f_1D, MPIHostType(), f_mpi_d);
    mirror_copy(f_mpi, f_1D);

    /*** Sum all planes ***/
    MPI_Allreduce(MPI_IN_PLACE, f_mpi.data(), f_mpi.size(), MPI_DOUBLE, MPI_SUM, pol_decomp.mpi.plane_comm);

    // Copy back
    mirror_copy(f_1D, f_mpi);

    Kokkos::fence();
}


void poloidal_gather(const DomainDecomposition<DeviceType>& pol_decomp, const View<double***, CLayout, DeviceType>& full_view){
    // No operation if n_planes is 1
    if(pol_decomp.mpi.n_plane_ranks==1) return;

    GPTLstart("POL_GATHER_SETUP");
    View<int*,CLayout,HostType> cnts(NoInit("cnts"), pol_decomp.mpi.n_plane_ranks);
    View<int*,CLayout,HostType> displs(NoInit("displs"), pol_decomp.mpi.n_plane_ranks);

    // Create cnts array with # vertex nodes on each processor
    int nv = full_view.extent(0)*full_view.extent(2); // nvr*nvz
    for(int i = 0; i<cnts.size(); i++){
        displs(i) = nv*(pol_decomp.gvid0_pid_h(i) - 1); // gvid0_pid_h is 1-indexed
        cnts(i)   = nv*(pol_decomp.gvid0_pid_h(i+1) - pol_decomp.gvid0_pid_h(i));
    }

    // Note: These are currently not the same as pol_decomp.node_offset and pol_decomp.nnodes. That's confusing! ALS
    int my_node_send_count = pol_decomp.gvid0_pid_h(pol_decomp.mpi.my_plane_rank+1) - pol_decomp.gvid0_pid_h(pol_decomp.mpi.my_plane_rank);
    int my_send_count = my_node_send_count*nv;
    int my_node_offset = pol_decomp.gvid0_pid_h(pol_decomp.mpi.my_plane_rank) - 1;

    GPTLstop("POL_GATHER_SETUP");

    // Need to transpose first; do this in a temporary array
    GPTLstart("POL_GATHER_ALLOC1");
    View<double***, CLayout, DeviceType> tmp(NoInit("tmp"), my_node_send_count, full_view.extent(0),full_view.extent(2));
    GPTLstop("POL_GATHER_ALLOC1");

    GPTLstart("POL_GATHER_TRANSP1");
    Kokkos::parallel_for("transpose", Kokkos::RangePolicy<ExSpace>( 0, tmp.extent(0)), KOKKOS_LAMBDA(const int inode){
        for(int imu = 0; imu < tmp.extent(1); imu++){
            for(int ivp = 0; ivp < tmp.extent(2); ivp++){
                tmp(inode, imu, ivp) = full_view(imu, inode+my_node_offset, ivp);
            }
        }
    });
    Kokkos::fence();
    GPTLstop("POL_GATHER_TRANSP1");

    using MPIHostType = HostType; // Do this MPI_Allreduce on Host; GPU-aware MPI consistently hangs or aborts here

    GPTLstart("POL_GATHER_MIRROR");
    // Cast to 1D with an unmanaged view for a generic element-wise operation, then copy to host if on GPU and GPU-aware MPI is not available
    View<double*,CLayout, DeviceType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_1D((double*)(tmp.data()), tmp.size());
    auto f_mpi_d = my_mirror_scratch_view(f_1D,MPIHostType());
    auto f_mpi = my_mirror_view(f_1D, MPIHostType(), f_mpi_d);
    mirror_copy(f_mpi, f_1D);

    // Now do the same for the output view, copying into a temporary view for the transpose
    View<double***, CLayout, DeviceType> full_tmp(NoInit("full_tmp"), full_view.extent(1), full_view.extent(0), full_view.extent(2));
    View<double*,CLayout, DeviceType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_1D_out((double*)(full_tmp.data()), full_tmp.size());
    auto f_mpi_d_out = my_mirror_scratch_view(f_1D_out,MPIHostType());
    auto f_mpi_out = my_mirror_view(f_1D_out, MPIHostType(), f_mpi_d_out);
    // No copy since output view will be overwritten
    GPTLstop("POL_GATHER_MIRROR");

    GPTLstart("POL_GATHER_MPI");
    MPI_Allgatherv(f_mpi.data(), my_send_count, MPI_DOUBLE, f_mpi_out.data(), cnts.data(), displs.data(), MPI_DOUBLE, pol_decomp.mpi.plane_comm);
    GPTLstop("POL_GATHER_MPI");

    GPTLstart("POL_GATHER_MIRRORBACK");
    // Copy back
    mirror_copy(f_1D_out, f_mpi_out);
    GPTLstop("POL_GATHER_MIRRORBACK");

    // Transpose back into full_view
    GPTLstart("POL_GATHER_TRANSP2");
    Kokkos::parallel_for("transpose_back", Kokkos::RangePolicy<ExSpace>( 0, full_tmp.extent(0)), KOKKOS_LAMBDA(const int inode){
        for(int imu = 0; imu < full_tmp.extent(1); imu++){
            for(int ivp = 0; ivp < full_tmp.extent(2); ivp++){
                full_view(imu, inode, ivp) = full_tmp(inode, imu, ivp);
            }
        }
    });
    Kokkos::fence();
    GPTLstop("POL_GATHER_TRANSP2");
}

#endif