mpi_space_wrapper.hpp

Go to the documentation of this file.

#ifndef MPI_SPACE_WRAPPER_HPP
#define MPI_SPACE_WRAPPER_HPP
 
#include "space_settings.hpp"
#include "my_mirror_view.hpp"
 
enum class MPISpaceOptions { // If supported, send/recv between device memory...
    AlwaysPreferGPUMPI,      // ... always (even if both views are on host)
    MixedPreferGPUMPI,       // ... if at least one of the views is on device
    MixedPreferHostMPI,      // ... if both views are on device
    AlwaysPreferHostMPI,     // ... never (even if both views are on device)
};
 
enum class MPISpaceWrapperOptions { // If input view is not in the same space as the selected MPI space...
    NotInPlace,                     // ... allocate a send buffer
    InPlace,                        // ... copy the input view to the recv buffer and use MPI_IN_PLACE
};
 
 
// Check whether some space is HostSpace
template<class Space>
struct is_host_space : std::false_type {};
 
template<>
struct is_host_space<Kokkos::HostSpace> : std::true_type {};
 
// Copy data to the selected space and execute the supplied MPI call
template<class SendView, class RecvView, class Device, class F>
void MPI_call_space_wrapper(const SendView& send, const RecvView& recv,
        F&& user_mpi_call, MPISpaceWrapperOptions option_in_place, Device nd){
 
    // Check whether to try in-place optimization (i.e. reuse recv buffer as send buffer)
    bool use_in_place = (option_in_place==MPISpaceWrapperOptions::InPlace &&
                         !std::is_same_v<typename SendView::memory_space, typename Device::memory_space>);
 
    // Allocate buffer for one or both views in MPI space as needed
    // If send buffer is not in the right space, we can try to do MPI_IN_PLACE on recv_mpi
    auto recv_mpi = my_mirror_view(recv, Device());
    auto send_mpi = use_in_place ? recv_mpi : my_mirror_view(send, Device());
 
    // Copy send if needed
    mirror_copy(send_mpi, send);
 
    // MPI call
    user_mpi_call(use_in_place ? MPI_IN_PLACE : send_mpi.data(), recv_mpi.data());
 
    // Copy recv back if needed
    mirror_copy(recv, recv_mpi);
 
}
 
// (MPI_IN_PLACE) Copy data to the selected space and execute the supplied MPI call
template<class SendView, class Device, class F>
void MPI_call_space_wrapper(const SendView& send,
        F&& user_mpi_call, Device nd){
 
    // Allocate buffer in MPI space as needed
    auto send_mpi = my_mirror_view(send, Device());
    
    // Copy send if needed
    mirror_copy(send_mpi, send);
    
    // MPI call
    user_mpi_call(send_mpi.data());
    
    // Copy recv back if needed
    mirror_copy(send, send_mpi);
        
}
 
// Choose which space to do the MPI operation in
template<class SendView, class RecvView>
bool MPI_choose_space(const SendView& send, const RecvView& recv, MPISpaceOptions option){
    using SendSpace  = typename SendView::memory_space;
    using RecvSpace  = typename RecvView::memory_space;
    using MPISpace   = typename MPIDeviceType::memory_space;
 
    // Decide whether MPI should see device (MPISpace) or host pointers
    constexpr bool gpu_aware_mpi_available = !is_host_space<MPISpace>::value;
    constexpr bool send_view_is_on_device = !is_host_space<SendSpace>::value;
    constexpr bool recv_view_is_on_device = !is_host_space<RecvSpace>::value;
    bool use_device_for_mpi = false; // Default to host MPI
    if(option==MPISpaceOptions::AlwaysPreferGPUMPI){
        use_device_for_mpi = gpu_aware_mpi_available;
    }else if(option==MPISpaceOptions::MixedPreferGPUMPI){
        use_device_for_mpi = gpu_aware_mpi_available && (send_view_is_on_device || recv_view_is_on_device);
    }else if(option==MPISpaceOptions::MixedPreferHostMPI){
        use_device_for_mpi = gpu_aware_mpi_available && (send_view_is_on_device && recv_view_is_on_device);
    }
    return use_device_for_mpi;
}
 
// (MPI_IN_PLACE) Choose which space to do the MPI operation in
template<class SendView>
bool MPI_choose_space(const SendView& send, MPISpaceOptions option){
    using SendSpace  = typename SendView::memory_space;
    using MPISpace   = typename MPIDeviceType::memory_space;
            
    // Decide whether MPI should see device (MPISpace) or host pointers
    constexpr bool gpu_aware_mpi_available = !is_host_space<MPISpace>::value;
    bool use_device_for_mpi = false; // Default to host MPI
    if(option==MPISpaceOptions::AlwaysPreferGPUMPI){
        use_device_for_mpi = gpu_aware_mpi_available;
    }else if(option==MPISpaceOptions::MixedPreferGPUMPI){ // These options should be disallowed for in-place
        use_device_for_mpi = gpu_aware_mpi_available;
    }else if(option==MPISpaceOptions::MixedPreferHostMPI){ // These options should be disallowed for in-place
        use_device_for_mpi = false;
    }
    return use_device_for_mpi;
}
 
/* Automatically copy views to/from device as needed to use MPI in the specified manner
   Wrap your intended MPI call, simply replacing f_send.data() and f_recv.data() with send_ptr and recv_ptr respectively.
   Default space option is MPISpaceOptions::MixedPreferGPUMPI
   in_place_option is an optimization to be used with care since MPI_IN_PLACE may not improve performance and will not always
   work with the same syntax (e.g. MPI_Reduce requires a different input for the root rank)
 
   Usage example:
   MPI_space_wrapper(f_send, f_recv, [&](double* send_ptr, double* recv_ptr){
       MPI_Allreduce(send_ptr, recv_ptr, f_send.size(), MPI_DOUBLE, MPI_SUM, pol_decomp.mpi.intpl_comm);
   }, MPISpaceOptions::AlwaysPreferHostMPI);
 */
template<class SendView, class RecvView, class F>
void MPI_space_wrapper(const SendView& send_MD, const RecvView& recv_MD, F&& user_mpi_call,
        MPISpaceOptions option = MPISpaceOptions::MixedPreferGPUMPI,
        MPISpaceWrapperOptions in_place_option = MPISpaceWrapperOptions::NotInPlace){
 
    using value_type = typename SendView::non_const_value_type;
    using SendSpace  = typename SendView::memory_space;
    using RecvSpace  = typename RecvView::memory_space;
 
    static_assert(std::is_same_v<typename SendView::non_const_value_type,
            typename RecvView::non_const_value_type>,
            "Send/recv value types must match.");
 
    // Get 1D unmanaged send/recv views; remove unmanaged tag for simpler template deductions (okay since views are const)
    const View<value_type*,CLayout, typename SendView::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> send_unm(send_MD.data(), send_MD.size());
    const View<value_type*,CLayout, typename RecvView::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> recv_unm(recv_MD.data(), recv_MD.size());
    const View<value_type*,CLayout, typename SendView::device_type> send = send_unm;
    const View<value_type*,CLayout, typename RecvView::device_type> recv = recv_unm;
 
    // Choose whether to do operation on host or device
    bool use_device_for_mpi = MPI_choose_space(send, recv, option);
 
    // Execute MPI operation
    if(use_device_for_mpi){
        MPI_call_space_wrapper(send, recv, user_mpi_call, in_place_option, DeviceType());
    }else{
        MPI_call_space_wrapper(send, recv, user_mpi_call, in_place_option, HostType());
    }
}
 
/* For MPI_IN_PLACE usage: Automatically copy view to/from device as needed to use MPI in the specified manner
   Wrap your intended MPI call, simply replacing f.data() with mpi_ptr
   Default space option is MPISpaceOptions::AlwaysPreferGPUMPI
 
   Usage example:
   MPI_space_wrapper(f, [&](double* mpi_ptr){
       MPI_Allreduce(MPI_IN_PLACE, mpi_ptr, f.size(), MPI_DOUBLE, MPI_SUM, pol_decomp.mpi.intpl_comm);
   }, MPISpaceOptions::AlwaysPreferHostMPI);
 */
template<class SendView, class F>
void MPI_space_wrapper(const SendView& send_MD, F&& user_mpi_call,
        MPISpaceOptions option = MPISpaceOptions::AlwaysPreferGPUMPI){
 
    using value_type = typename SendView::non_const_value_type;
    using SendSpace  = typename SendView::memory_space;
 
    // Get 1D unmanaged send views; remove unmanaged tag for simpler template deductions (okay since views are const)
    const View<value_type*,CLayout, typename SendView::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> send_unm(send_MD.data(), send_MD.size());
    const View<value_type*,CLayout, typename SendView::device_type> send = send_unm;
 
    // Choose whether to do operation on host or device
    bool use_device_for_mpi = MPI_choose_space(send, option);
 
    // Execute MPI operation
    if(use_device_for_mpi){
        MPI_call_space_wrapper(send, user_mpi_call, DeviceType());
    }else{
        MPI_call_space_wrapper(send, user_mpi_call, HostType());
    }
}
 
 
#endif