task_group.hpp

Go to the documentation of this file.

#ifndef TASK_GROUP_HPP
#define TASK_GROUP_HPP
#ifdef USE_MPI
#include "my_mpi.hpp"
#endif

class TaskGroup{
#ifdef USE_MPI
    MPI_Comm comm; 
#endif
    int my_rank; 
    int n_ranks; 

    public:

    TaskGroup(){}

#ifdef USE_MPI

    TaskGroup(MPI_Comm comm_in)
      : comm(comm_in)
    {
       MPI_Comm_size( comm, &n_ranks );
       MPI_Comm_rank( comm, &my_rank );
    }
#endif

    int get_my_rank(){
        return my_rank;
    }

    int get_n_local_tasks(int n_tasks){
        return ((n_tasks + n_ranks - my_rank - 1)/n_ranks);
    }

    int get_max_n_local_tasks(int n_tasks){
        return ((n_tasks + n_ranks - 1)/n_ranks);
    }

    template<typename F>
    void execute_task(int ntasks, F func){
        for(int itask=0; itask<ntasks; itask++){
            int assigned_rank = itask%n_ranks;
            if(my_rank==assigned_rank){
                int i_local_task = itask/n_ranks;
                func(itask, i_local_task);
            }
        }
    }

#ifdef USE_MPI

    template<typename F>
    void communicate(int ntasks, F func){
        for(int itask=0; itask<ntasks; itask++){
            int i_local_task = itask/n_ranks;
            int assigned_rank = itask%n_ranks;
            bool is_assigned_rank = (my_rank==assigned_rank);
            func(itask, i_local_task, assigned_rank, is_assigned_rank, comm);
        }
    }
#endif
};

#endif