AMReX_ForkJoin.H

Go to the documentation of this file.

#ifndef AMREX_FORKJOIN_H
#define AMREX_FORKJOIN_H
#include <AMReX_Config.H>
 
#include <AMReX_ParallelContext.H>
#include <AMReX_MultiFab.H>
#include <utility>
 
namespace amrex {
 
class ForkJoin
{
  public:
 
    enum class Strategy {
        single,     
        duplicate,  
        split,      
    };
    enum class Intent { in, out, inout };
 
    struct ComponentSet {
        ComponentSet() = default;
        ComponentSet(int lo_, int hi_) : lo(lo_), hi(hi_) {}
        int lo; 
        int hi; 
    };
 
    ForkJoin (const Vector<int> &task_rank_n);
 
    ForkJoin (const Vector<double> &task_rank_pct);
 
    ForkJoin (int ntasks)
        : ForkJoin( Vector<double>(ntasks, 1.0 / ntasks) ) { }
 
    int NTasks () const { return static_cast<int>(split_bounds.size() - 1); }
 
    int MyTask () const { return task_me; }
 
    bool Verbose () const { return flag_verbose; }
 
    void SetVerbose (bool verbose_in) { flag_verbose = verbose_in; }
 
    ComponentSet ComponentBounds(const std::string& name, int idx=0) const;
 
    int NProcsTask (int task) const {
        AMREX_ASSERT(task + 1 < split_bounds.size());
        return split_bounds[task + 1] - split_bounds[task];
    }
 
    void reg_mf (MultiFab &mf, const std::string &name, int idx,
                 Strategy strategy, Intent intent, int owner = -1);
 
    void reg_mf (MultiFab &mf, const std::string &name,
                 Strategy strategy, Intent intent, int owner = -1) {
        reg_mf(mf, name, 0, strategy, intent, owner);
    }
 
    void reg_mf (const MultiFab &mf, const std::string &name, int idx,
                 Strategy strategy, Intent intent, int owner = -1) {
        AMREX_ASSERT_WITH_MESSAGE(intent == Intent::in,
                                  "const MultiFab must be registered read-only");
        reg_mf(const_cast<MultiFab&>(mf), name, idx, strategy, intent, owner);
    }
    void reg_mf (const MultiFab &mf, const std::string &name,
                 Strategy strategy, Intent intent, int owner = -1) {
        reg_mf(mf, name, 0, strategy, intent, owner);
    }
 
    void reg_mf_vec (const Vector<MultiFab *> &mfs, const std::string &name,
                     Strategy strategy, Intent intent, int owner = -1) {
        data[name].reserve(mfs.size());
        for (int i = 0; i < mfs.size(); ++i) {
            reg_mf(*mfs[i], name, i, strategy, intent, owner);
        }
    }
 
    void reg_mf_vec (const Vector<MultiFab const *> &mfs, const std::string &name,
                     Strategy strategy, Intent intent, int owner = -1) {
        data[name].reserve(mfs.size());
        for (int i = 0; i < mfs.size(); ++i) {
            reg_mf(*mfs[i], name, i, strategy, intent, owner);
        }
    }
 
    void modify_ngrow (const std::string &name, int idx, IntVect ngrow);
    void modify_ngrow (const std::string &name, IntVect ngrow) {
        modify_ngrow(name, 0, ngrow);
    }
 
    void modify_split (const std::string &name, int idx, Vector<ComponentSet> comp_split);
    void modify_split (const std::string &name, Vector<ComponentSet> comp_split) {
        modify_split(name, 0, std::move(comp_split));
    }
 
    // TODO: may want to de-register MFs if they change across invocations
 
    MultiFab &get_mf (const std::string &name, int idx = 0) {
        AMREX_ASSERT_WITH_MESSAGE(data.count(name) > 0 && idx < data[name].size(), "get_mf(): name or index not found");
        AMREX_ASSERT(task_me >= 0 && task_me < data[name][idx].forked.size());
        return data[name][idx].forked[task_me];
    }
 
    Vector<MultiFab *> get_mf_vec (const std::string &name) {
        int dim = static_cast<int>(data.at(name).size());
        Vector<MultiFab *> result(dim);
        for (int idx = 0; idx < dim; ++idx) {
            result[idx] = &get_mf(name, idx);
        }
        return result;
    }
 
    void set_task_output_dir (const std::string & dir)
    {
        task_output_dir = dir;
    }
 
    static void set_task_output_file (const std::string & filename)
    {
        ParallelContext::set_last_frame_ofs(filename);
    }
 
    template <class F>
    void fork_join (const F &fn)
    {
        flag_invoked = true; // set invoked flag
        const int io_rank = 0; // team's sub-rank 0 does IO
        create_task_output_dir();
        MPI_Comm task_comm = split_tasks();
        copy_data_to_tasks(); // move data to local tasks
        ParallelContext::push(task_comm, task_me, io_rank);
        set_task_output_file(get_io_filename());
        fn(*this);
        ParallelContext::pop();
#ifdef BL_USE_MPI
        MPI_Comm_free(&task_comm);
#endif
        copy_data_from_tasks(); // move local data back
    }
 
  private:
 
    struct MFFork
    {
        MultiFab *orig = nullptr;
        Strategy strategy;
        Intent intent;
        int owner_task; 
        IntVect ngrow;
        Vector<ComponentSet> comp_split; 
        Vector<MultiFab> forked; 
 
        MFFork () = default;
        ~MFFork () = default;
        MFFork (const MFFork&) = delete;
        MFFork& operator= (const MFFork&) = delete;
        MFFork (MFFork&&) = default;
        MFFork& operator= (MFFork&&) = default;
        MFFork (MultiFab *omf, Strategy s, Intent i, int own,
                const IntVect &ng,
                Vector<ComponentSet> cs)
            : orig(omf), strategy(s), intent(i), owner_task(own), ngrow(ng), comp_split(std::move(cs)) {}
 
        [[nodiscard]] bool empty() const { return orig == nullptr; }
    };
 
    bool flag_verbose = false; 
    bool flag_invoked = false; 
    Vector<int> split_bounds; 
    int task_me = -1; 
    std::map<BoxArray::RefID, Vector<std::unique_ptr<DistributionMapping>>> dms; 
    std::unordered_map<std::string, Vector<MFFork>> data;
    std::string task_output_dir; 
 
    void init(const Vector<int> &task_rank_n);
 
    const DistributionMapping &get_dm (const BoxArray& ba, int task_idx,
                                       const DistributionMapping& dm_orig);
 
    void copy_data_to_tasks ();
 
    void copy_data_from_tasks ();
 
    MPI_Comm split_tasks ();
 
    void create_task_output_dir ();
 
    std::string get_io_filename (bool flag_unique = false);
};
 
}
 
#endif // AMREX_FORKJOIN_H