broadcast_views.hpp

Go to the documentation of this file.

#ifndef BROADCAST_VIEWS_HPP
#define BROADCAST_VIEWS_HPP

#include "space_settings.hpp"
#include "my_mpi.hpp"

// Utility to resize a view based on extents
template <typename ViewType>
void resize_view(ViewType& view, const std::vector<int>& extents) {
    if constexpr (ViewType::rank == 1) {
        view = ViewType(view.label(), extents[0]);
    } else if constexpr (ViewType::rank == 2) {
        view = ViewType(view.label(), extents[0], extents[1]);
    } else if constexpr (ViewType::rank == 3) {
        view = ViewType(view.label(), extents[0], extents[1], extents[2]);
    } else if constexpr (ViewType::rank == 4) {
        view = ViewType(view.label(), extents[0], extents[1], extents[2], extents[3]);
    } else if constexpr (ViewType::rank == 5) {
        view = ViewType(view.label(), extents[0], extents[1], extents[2], extents[3], extents[4]);
    } else if constexpr (ViewType::rank == 6) {
        view = ViewType(view.label(), extents[0], extents[1], extents[2], extents[3], extents[4], extents[5]);
    }
}

// Function template to broadcast and resize an arbitrary number of Kokkos Views
template <typename... Views>
void broadcast_views(const MPI_Comm& comm, const int ROOT_RANK, Views&... views) {
    int my_rank;
    int nranks;
    MPI_Comm_rank( comm, &my_rank);
    MPI_Comm_size( comm, &nranks);

    // Allocate extents and use extents found on root rank
    std::vector<std::vector<int>> extents_list;
    ([&](){
        std::vector<int> extents(Views::rank);
        if(my_rank==ROOT_RANK){
            for (int i = 0; i < Views::rank; i++) {
                extents[i] = views.extent(i);
            }
        }
        extents_list.push_back(extents);
    }(), ...); // Fold expression for multiple views
    
    // Loop through each view, broadcast its extents
    for (int v = 0; v < (sizeof...(Views)); v++) {
        MPI_Bcast(extents_list[v].data(), extents_list[v].size(), MPI_INT, ROOT_RANK, comm);
    }
    
    // Resize each view on all other ranks to match extents from Rank ROOT_RANK
    if(my_rank!=ROOT_RANK){
        int view_index = 0;
        ([&](){
            resize_view(views, extents_list[view_index]);
            view_index++;
        }(), ...);
    }

    // Broadcast view data from ROOT_RANK
    ([&](){
        MPI_Bcast(views.data(), views.size(), get_mpi_type(views), ROOT_RANK, comm);
    }(), ...);
}

#endif