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()
      : my_rank(0),
        n_ranks(1){}
 
#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() const{
        return my_rank;
    }
 
    int get_n_local_tasks(int n_tasks) const{
        return ((n_tasks + n_ranks - my_rank - 1)/n_ranks);
    }
 
    int get_max_n_local_tasks(int n_tasks) const{
        return ((n_tasks + n_ranks - 1)/n_ranks);
    }
 
#ifdef USE_MPI
    MPI_Comm get_comm() const{
        return comm;
    }
#endif
 
    template<typename F>
    void execute_task(int ntasks, F func) const{
        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) const{
        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