load_balance.hpp

Go to the documentation of this file.

#ifndef LOAD_BALANCE_HPP
#define LOAD_BALANCE_HPP

#include "timer_macro.hpp"
#include "shift.hpp"
#include "count_ptl_per_node.hpp"
#include "view_arithmetic.hpp"
#include "f0_redistribute.hpp"

// Used for original fortran load balancer
extern "C" void calculate_load_imbalance(double f0_cost);
extern "C" int assess_whether_to_rebalance_load();
extern "C" void reset_cost_trackers();
extern "C" void set_weights(int* gvid0_pid, double* ptl_count, double* f0_node_cost);

class LoadRegion{
    public:
    enum class UpdateMethod{
        NoHistory=0,
        ExpHistory
    };

    private:

    View<double*,HostType> estimated_time_per_vertex;
    std::string region_name;
    std::vector<std::string> timer_names;
    bool verbose;
    bool model_is_initialized;
    bool model_has_history;
    double predicted_max_region_time;
    double observed_max_region_time;
    double observed_load_imbalance;
    double prediction_undershoot;

#ifdef USE_MPI
    MPI_Comm inter_period_comm;
    MPI_Comm period_comm;
#endif
    int n_periods;  
    int n_unique_ranks; 
    int my_period_rank;

    double time_accumulated;

    UpdateMethod update_method;

    void update_model_no_history(const View<int*,CLayout,HostType>& current_partition, const View<double*, HostType>& all_periods_timings){
        // Loop over ranks' timers
        for(int i=0; i<all_periods_timings.size(); i++){
            int node_offset = current_partition(i) - 1;
            int next_node_offset = current_partition(i+1) - 1;
            int nnodes = next_node_offset - node_offset;
            // Best 0th order estimate is to assume all vertices took the same time
            double time_per_node = all_periods_timings(i)/nnodes;
            for(int j=0; j<nnodes; j++){
                estimated_time_per_vertex(j + node_offset) = time_per_node;
            }
        }
    }

    void update_model_exp_history(const View<int*,CLayout,HostType>& current_partition, const View<double*, HostType>& all_periods_timings){
        // Loop over ranks' timers
        for(int i=0; i<all_periods_timings.size(); i++){
            // Get partition owned by this rank
            int node_offset = current_partition(i) - 1;
            int next_node_offset = current_partition(i+1) - 1;
            int nnodes = next_node_offset - node_offset;

            // Get average time per node for this rank
            double avg_time_per_node = all_periods_timings(i)/nnodes;

            // Get expected time based on existing model
            double expected_time = 0.0;
            for(int j=0; j<nnodes; j++){
                expected_time += estimated_time_per_vertex(j + node_offset);
            }
            double expected_time_per_node = expected_time/nnodes;

            // Difference between observed time and expected time
            double extra_time_per_node = avg_time_per_node - expected_time_per_node;

            for(int j=0; j<nnodes; j++){
                // Distribute extra time evenly
                constexpr double adjustment_rate = 0.5;

                estimated_time_per_vertex(j + node_offset) += extra_time_per_node*adjustment_rate;
            }
        }
    }

    double get_time_since_previous_call(){
        // Get time from camtimers (accumulated)
        double new_accumulated_time = 0.0;

        // Loop through all timers associated with this region, and add the accumulated time for each
        for(int i=0; i<timer_names.size(); i++){
            double accumulated_single_timer = 0;
            int ierr = GPTLget_wallclock(timer_names[i].c_str(), -1, &accumulated_single_timer);

            // If ierr != 0, timer is not yet defined, so the accumulated time is zero
            if(ierr != 0) accumulated_single_timer = 0.0;

            new_accumulated_time += accumulated_single_timer;
        }

        // Subtract previous accumulated time to get time spent since last update
        double time_spent = new_accumulated_time - time_accumulated;

        // Save new accumulated time
        time_accumulated = new_accumulated_time;

        return time_spent;
    }

    void touch_timers(){
        for(int i=0; i<timer_names.size(); i++){
            GPTLstart(timer_names[i].c_str());
            GPTLstop(timer_names[i].c_str());
        }
    }

    public:

#ifdef USE_MPI
    LoadRegion(int n_vertices, const MPI_Comm& inter_period_comm, const MPI_Comm& period_comm, UpdateMethod update_method, bool verbose, std::string region_name, std::vector<std::string> timer_names)
        : model_is_initialized(false),
          verbose(verbose),
          time_accumulated(0.0),
          inter_period_comm(inter_period_comm),
          period_comm(period_comm),
          estimated_time_per_vertex("estimated_time_per_vertex",n_vertices),
          prediction_undershoot(1.0),
          update_method(update_method),
          model_has_history(false),
          region_name(region_name),
          timer_names(timer_names)
    {
        // Touch GPTL timers to be sure they exist
        touch_timers();

        // Initialize timer
        reset_timer();

        // Get problem size from comms
        MPI_Comm_rank(period_comm, &my_period_rank);
        MPI_Comm_size(period_comm, &n_unique_ranks);
        MPI_Comm_size(inter_period_comm, &n_periods);
    }
#endif

    void reset_timer(){
        // Reset timer by getting the time since the previous call and suppressing the output
        double time_spent = get_time_since_previous_call();
    }

    double get_estimated_time_per_vertex(int i) const{
        return estimated_time_per_vertex(i);
    }

    View<double*,HostType> get_estimated_time_per_vertex() const{
        return estimated_time_per_vertex;
    }

    double get_prediction_undershoot() const{
        return prediction_undershoot;
    }

    double get_observed_max_region_time() const{
        return observed_max_region_time;
    }

    double get_observed_load_imbalance() const{
        return observed_load_imbalance;
    }

    bool get_model_is_initialized() const{
        return model_is_initialized;
    }

    // Predict performance of new partition based on model
    double get_largest_predicted_time(const View<int*,CLayout,HostType>& proposed_partition) const{
        double largest_time = 0.0;
        //int largest_time_ind = 0;
        for(int i=0; i<(proposed_partition.size()-1); i++){
            double proc_time = 0.0;
            for(int i_node=proposed_partition(i)-1; i_node<proposed_partition(i+1)-1; i_node++){
                proc_time += estimated_time_per_vertex(i_node);
            }
            if(proc_time>largest_time){
                //largest_time_ind = i;
                largest_time = proc_time;
            }
        }
        return largest_time;
    }

    void initialize_model(){
        reset_timer();

        // Model can now be used
        model_is_initialized = true;
    }

    void update_model(const View<int*,CLayout,HostType>& current_partition, double manual_time=-1.0){
        if(update_method==UpdateMethod::NoHistory || update_method==UpdateMethod::ExpHistory){

            // Allocate for timing of each
            View<double*, HostType> all_periods_timings(NoInit("all_periods_timings"), n_unique_ranks);

#ifdef USE_MPI
            // Get time spent since last model update
            // Time can be entered manually; it's used if positive (i.e. the input was provided)
            double time_spent_this_rank = (manual_time<0.0 ? get_time_since_previous_call()
                                                           : manual_time);

            // Reduce onto plane 0
            // Could try MPI_SUM rather than MPI_MAX
            double time_spent;
            if(n_periods>1){
                MPI_Reduce(&time_spent_this_rank, &time_spent, 1, MPI_DOUBLE, MPI_MAX, 0, inter_period_comm);
            }else{
                time_spent = time_spent_this_rank;
            }

            // Gather from all ranks in plane 0
            MPI_Gather(&time_spent, 1, MPI_DOUBLE, all_periods_timings.data(), 1, MPI_DOUBLE, 0, period_comm);
#endif

            if (is_rank_zero()){
                // Look at the timing prediction made last time
                predicted_max_region_time = get_largest_predicted_time(current_partition);
                if(verbose) printf("\nPredicted max time in this region was %1.5e", predicted_max_region_time);

                // Get max region time
                observed_max_region_time = 0.0;
                double observed_sum_region_time = 0.0;
                for(int i=0; i<all_periods_timings.size(); i++){
                    observed_max_region_time = std::max(observed_max_region_time, all_periods_timings(i));
                    observed_sum_region_time += all_periods_timings(i);
                }
                double observed_idle_region_time = observed_max_region_time*all_periods_timings.size() - observed_sum_region_time;
                observed_load_imbalance = observed_idle_region_time/observed_sum_region_time;
                if(verbose) printf("\nObserved max time in this region was %1.5e", observed_max_region_time);
                if(verbose) printf("\n - Load imbalance (T_idle/T_work): %1.2f%%", observed_load_imbalance*100);

                // Take the max here so that there is never an assumed overshoot
                if(predicted_max_region_time!=0.0){ // If time is zero, there hasn't been a prediction yet
                    prediction_undershoot = std::max(observed_max_region_time/predicted_max_region_time, 1.0);
                }
                if(verbose) printf("\nThe prediction undershot by a factor of %1.5e", observed_max_region_time/predicted_max_region_time);

                if(update_method==UpdateMethod::NoHistory){
                    update_model_no_history(current_partition, all_periods_timings);
                }else if(update_method==UpdateMethod::ExpHistory){
                    if(model_has_history){
                        update_model_exp_history(current_partition, all_periods_timings);
                    }else{
                        update_model_no_history(current_partition, all_periods_timings);
                        model_has_history = true;
                    }
                }
            }
        }else{
            exit_XGC("Error: Update method not available\n");
        }
    }
};

class LoadBalance{
    public:
    enum class WeightingAlgorithm{
        SingleRegionBalance=0,
        ParticleBalance,
        Fortran,
        Default
    };

    enum class ConstraintOption{
        ParticleCount=0
    };

    enum class ReweightOption{
        Always=0,
        IfBetter
    };

    private:

    WeightingAlgorithm default_weighting_algorithm;

    std::vector<LoadRegion> regions;

    View<int*,CLayout,HostType> proposed_partition; 

#ifdef USE_MPI
    MPI_Comm comm;
#endif

    double threshold_to_rebalance;
    bool verbose;

    double constraint1_max;

    double get_even_division(const View<double*,HostType>& input, int n) const{
        double total = 0.0;
        Kokkos::parallel_reduce("sum", Kokkos::RangePolicy<HostExSpace>(0,input.size()), [=]( const int i, double& l_total){
            l_total += input(i);
        }, total);

        return (total/n);
    }


    bool greedily_fill_partition(const View<double*,HostType>& weight, const View<double*,HostType>& constraint1, double target_weight_per_rank){
        int nnode = weight.size();
        int nproc = proposed_partition.size()-1;

        // Start from rank 0 of the plane; assign nodes until it meets the desired weight per rank, then move to next rank
        double constraint1_in_this_rank = 0.0;
        double weight_in_this_rank = 0.0;
        proposed_partition(0) = 1; // 1-indexed
        int pid = 0;
        bool assign_one_node_per_proc = false;
        for(int i_node = 0; i_node<nnode; i_node++){
            // Since every rank needs at last one node, switch to assigning at every
            // iteration if there are only as many nodes left as ranks
            if(nnode-i_node==nproc-pid) assign_one_node_per_proc = true;

            // Check if a criterion is met to consider the rank sufficiently loaded
            bool rank_is_loaded = false;

            // Criterion 1: We are out of nodes
            if(assign_one_node_per_proc) rank_is_loaded = true;

            // Criterion 2: We have hit a constraint
            if(constraint1_in_this_rank>=constraint1_max) rank_is_loaded = true;

            // Criterion 3: We have loaded the rank with an even amount of work
            if(weight_in_this_rank>=target_weight_per_rank) rank_is_loaded = true;

            // If there are enough particles assigned to this rank, move to the next rank
            if(rank_is_loaded){
                proposed_partition(pid+1) = i_node+1; // 1-indexed
                if(verbose) printf("\nRank %d is loaded (Nodes %d-%d); weight=%1.4e, ptl*nplanes=%d", pid, proposed_partition(pid), proposed_partition(pid+1)-1, weight_in_this_rank, int(constraint1_in_this_rank));
                constraint1_in_this_rank = 0.0;
                weight_in_this_rank = 0.0;
                pid++;
                if(pid==nproc-1) break;
            }
            constraint1_in_this_rank += constraint1(i_node);
            weight_in_this_rank += weight(i_node);
        }

        // Last value will always be nnode+1
        proposed_partition(nproc) = nnode+1;

        // Check n_ptl in final rank
        constraint1_in_this_rank = 0.0;
        for(int i_node=proposed_partition(nproc-1)-1; i_node<nnode; i_node++){
            constraint1_in_this_rank += constraint1(i_node);
        }
        if(constraint1_in_this_rank>=constraint1_max){
            // Return false if the constraint was not met
            return false;
        }else{
            return true;
        }
    }

    // Predict performance of new partition based on model
    double get_largest_predicted_time(const View<int*,CLayout,HostType>& partition, const View<double*,HostType>& weight) const{
        double largest_time = 0.0;
        //int largest_time_ind = 0;
        for(int i=0; i<(partition.size()-1); i++){
            double proc_time = 0.0;
            for(int i_node=partition(i)-1; i_node<partition(i+1)-1; i_node++){
                proc_time += weight(i_node);
            }
            if(proc_time>largest_time){
                //largest_time_ind = i;
                largest_time = proc_time;
            }
        }
        return largest_time;
    }

    void one_weight_balance(const View<double*,HostType>& weight, const View<double*,CLayout,HostType> constraint1){
        int nproc = proposed_partition.size()-1;

        // Ideally, each rank would get the same amount of work, i.e. the average
        double ideal_weight_per_rank = get_even_division(weight, nproc);

        // Make initial attempt, targeting even distribution of work
        bool meets_constraints = greedily_fill_partition(weight, constraint1, ideal_weight_per_rank);

        if(!meets_constraints){
            // An equal work distribution cannot be assigned due to a constraint.
            // First, determine the weights if only the constraint were followed
            meets_constraints = greedily_fill_partition(constraint1, constraint1, get_even_division(constraint1, nproc));
            if(!meets_constraints) exit_XGC("\nUnexpected issue in load balance: constraint-based partition doesn't satisfy constraints\n");

            // This partition meets the constraints, but the proposed timing is likely unacceptable
            double upper_limit_weight_per_rank = get_largest_predicted_time(proposed_partition, weight);
            if(verbose) printf("The ideal distribution (%1.3e) does not satisfy constraints. This distribution (%1.3e) does.\n", ideal_weight_per_rank, upper_limit_weight_per_rank);

            // Bisect the difference between the original target time and the upper limit, and see if that satisfies the constraints
            // Continue until the load balance has been minimized to a precision based on the the original desired time
            const double desired_precision = 0.01; // Fraction of imbalance that's worth even checking whether it can be removed
            double desired_step_size = ideal_weight_per_rank*desired_precision;
            double step_size = upper_limit_weight_per_rank - ideal_weight_per_rank;
            double compromise_weight_per_rank = upper_limit_weight_per_rank;
            if(verbose) printf("\nEmploying a binary search to narrow down on the best option.");
            while(step_size>desired_step_size){
                // Halve the size of the step
                step_size /= 2;

                if(meets_constraints){
                    // Try reducing the load imbalance since constraints were met
                    compromise_weight_per_rank -= step_size;
                }else{
                    // Try raising the load imbalance since constraints were broken
                    compromise_weight_per_rank += step_size;
                }

                // Get new partition
                meets_constraints = greedily_fill_partition(weight, constraint1, compromise_weight_per_rank);
                if(verbose) printf("\n  Stepped by %1.3e to %1.3e. The new partition does%s meet constraints.", step_size, compromise_weight_per_rank, meets_constraints?"" : "nt");
            }
            // In case we end at a partition that fails
            while(!meets_constraints){
                compromise_weight_per_rank += step_size;
                meets_constraints = greedily_fill_partition(weight, constraint1, compromise_weight_per_rank);
                if(verbose) printf("\n  Stepped by %1.3e UP to %1.3e. The new partition does%s meet constraints.", step_size, compromise_weight_per_rank, meets_constraints?"" : "nt");
            }
            if(verbose) printf("\n");
        }
    }

    // TODO
    // Evaluate overall model performance
    // MAIN_LOOP time vs sum of regions (% coverage)
    //

    /* This function assesses the proposed partition and decides whether it is worth switching to.
     * The simplest formulation is to recommend the new partition if the new partition is predicted to be faster
     * than the existing one. There is also a factor to account for historical inaccuracy of the previous estimate.
     * A future option would be to require a certain level of improvement before repartitioning. */
    bool recommend_proposed_partition(){
        // Total time spent in regions (assuming that global barriers demarcate them):
        double observed_max_all_rgns_time = 0.0;
        for(int i=0; i<regions.size(); i++){
            observed_max_all_rgns_time += regions[i].get_observed_max_region_time();
        }

        // Total predicted time, accounting for observed prediction undershoot for each region
        double predicted_max_all_rgns_time = 0.0;
        for(int i=0; i<regions.size(); i++){
            predicted_max_all_rgns_time += regions[i].get_largest_predicted_time(proposed_partition)*regions[i].get_prediction_undershoot();
        }

        double fractional_improvement = 1.0 - predicted_max_all_rgns_time/observed_max_all_rgns_time;

        if(verbose){
            printf("\nDetermining whether to adopt the proposed partition:");
            printf("\n Observed time with current partition was: %1.3e", observed_max_all_rgns_time);
            printf("\n Predicted time with proposed partition, adjusting for historical undershoot (of Rgn 0: %1.3e): %1.3e", regions[0].get_prediction_undershoot(), predicted_max_all_rgns_time);
            printf("\n Adopted if Fractional improvement of adjusted prediction (%1.3e) exceeds specified threshold (%1.3e)\n", fractional_improvement, threshold_to_rebalance);
        }

        // Recommend the new partition if the predicted time is less than the time observed with the current partition
        if(fractional_improvement > threshold_to_rebalance){
            return true;
        }else{
            return false;
        }
    }

    void print_new_partition(){
        // Print new partition
        for(int i = 0; i<proposed_partition.size()-1; i++){
            int nnodes_for_proc = proposed_partition(i+1) - proposed_partition(i);
            printf("process %d : nnodes = %d; range = (%d - %d)\n", i, nnodes_for_proc, proposed_partition(i), proposed_partition(i+1)-1);
        }
    }

    void update_model(const View<int*,CLayout,HostType>& current_partition){
        for(int i=0; i<regions.size(); i++){
            regions[i].update_model(current_partition);
        }
    }

    // Overload if providing timings manually rather than using camtimers. This is used for testing.
    void update_model(const View<int*,CLayout,HostType>& current_partition, const std::vector<double>& manual_times){
        for(int i=0; i<regions.size(); i++){
            regions[i].update_model(current_partition, manual_times[i]);
        }
    }

    void initialize_model(){
        for(int i=0; i<regions.size(); i++){
            regions[i].initialize_model();
        }
    }

    bool model_is_initialized(){
        bool is_initialized=true;
        for(int i=0; i<regions.size(); i++){
            is_initialized = (is_initialized && regions[i].get_model_is_initialized());
        }
        return is_initialized;
    }

    void propose_new_partition(const Kokkos::View<double*,Kokkos::LayoutRight,HostType>& ptl_count, WeightingAlgorithm weighting_algorithm){
        auto constraint1 = ptl_count; // Constrain based on particle count (memory)
        if(weighting_algorithm == WeightingAlgorithm::SingleRegionBalance){
            if(regions.size()!=1) exit_XGC("Error: Load balancing is currently single-region only\n");
            one_weight_balance(regions[0].get_estimated_time_per_vertex(), constraint1);
        }else if(weighting_algorithm == WeightingAlgorithm::ParticleBalance){
            one_weight_balance(ptl_count, constraint1);
        }

        if(verbose){
            printf("\nNewly PROPOSED partition:\n");
            print_new_partition();
        }
    }

    void set_new_partition(const Simulation<DeviceType>& sml, const Grid<DeviceType>& grid, const MagneticField<DeviceType>& magnetic_field,
                   const VelocityGrid& vgrid, Plasma& plasma, DomainDecomposition<DeviceType>& pol_decomp, WeightingAlgorithm weighting_algorithm){
        if(weighting_algorithm==WeightingAlgorithm::Fortran){
            // In the fortran algorithm the new distribution is calculated and set here
            // Count number of particles belonging to each node
            Kokkos::View<double**,Kokkos::LayoutRight,HostType> ptl_count = count_ptl_per_node_elec_main_ion(grid, magnetic_field, plasma, sml.electron_on);

            if(pol_decomp.gvid0_pid_h.size()==2){
                // Skip set_weights if there is only one rank per plane
                pol_decomp.gvid0_pid_h(0) = 1;
                pol_decomp.gvid0_pid_h(1) = pol_decomp.nnodes_on_plane + 1;
            }else{
                TIMER("SET_WEIGHTS",
                    set_weights(pol_decomp.gvid0_pid_h.data(), ptl_count.data(), plasma.f0_node_cost.data()) );
            }
        }else{
#ifdef USE_MPI
            // Broadcast proposal to all ranks
            MPI_Bcast(proposed_partition.data(), proposed_partition.size(), MPI_INT, 0, comm);
#endif

            // Copy proposal to pol_decomp
            Kokkos::deep_copy(pol_decomp.gvid0_pid_h, proposed_partition);

            if(is_rank_zero()){
                printf("\nNEW PARTITION:\n");
                print_new_partition();
            }
        }

        // Update bounds and device copy of the decomposition array
        pol_decomp.update_pol_decomp();
    }

    void redistribute_load(const Simulation<DeviceType>& sml, const Grid<DeviceType>& grid, const MagneticField<DeviceType>& magnetic_field,
                   const VelocityGrid& vgrid, Plasma& plasma, DomainDecomposition<DeviceType>& pol_decomp,
                   const View<int*,CLayout,HostType>& old_partition){
        if (sml.f0_grid){
            // Move f information to correct MPI rank after decomposition update
            TIMER("F0_REDISTRIBUTE",
                f0_redistribute(plasma, pol_decomp, grid, magnetic_field, vgrid, old_partition) );
        }

        // Shift particles to correct MPI rank after decomposition update
        TIMER("SHIFT_R",
            shift_all_species(plasma, grid, magnetic_field, pol_decomp, Shift::NoPhase0) );
    }

    bool will_rebalance(ReweightOption reweight_option, WeightingAlgorithm weighting_algorithm, double f0_cost){
        if(reweight_option==ReweightOption::Always){
            // No need to check if model is an improvement
            return true;
        }else{
            if(weighting_algorithm==WeightingAlgorithm::Fortran){
                // Calculate load imbalance
                calculate_load_imbalance(f0_cost);

                // Assess if load imbalance merits a rebalancing
                int should_rebalance_int = assess_whether_to_rebalance_load();

                // Reset fortran cost trackers
                reset_cost_trackers();

                return (should_rebalance_int==1);
            }else{
                // Evaluate the proposed partition on rank 0
                bool should_rebalance;
                if(weighting_algorithm==WeightingAlgorithm::ParticleBalance){
                    // Since ParticleBalance doesn't contain an internal model, cant evaluate the partition
                    // so just always approve, for now.
                    // Probably better to have the model updatable by something other than time so it can still
                    // be assessed in this mode
                    return true;
                }else{
                    if(is_rank_zero()) should_rebalance = recommend_proposed_partition();
                    // Convert to int because MPI_C_BOOL documentation is confusing
                    int should_rebalance_int = (should_rebalance ? 1 : 0);
#ifdef USE_MPI
                    MPI_Bcast(&should_rebalance_int, 1, MPI_INT, 0, comm);
#endif
                    return (should_rebalance_int==1);
                }
            }
        }
    }

    public:

    LoadBalance(NLReader::NamelistReader& nlr, const DomainDecomposition<DeviceType>& pol_decomp, bool sync_planes = true)
      : proposed_partition(NoInit("proposed_partition"), pol_decomp.gvid0_pid_h.layout())
#ifdef USE_MPI
        , comm(sync_planes ? pol_decomp.mpi.comm : pol_decomp.mpi.plane_comm)
#endif
    {
        double max_mem_redist_gb = 10.0;
        std::string weighting_algorithm_str = "Fortran";
        std::string update_method_str = "NoHistory";
        if(nlr.namelist_present("load_balance_param")){
            nlr.use_namelist("load_balance_param");
            max_mem_redist_gb = nlr.get<double>("max_mem_redist_gb", 10.0);
            weighting_algorithm_str = nlr.get<std::string>("weighting_algorithm", "Fortran");
            threshold_to_rebalance = nlr.get<double>("threshold_to_rebalance", 0.02);
            verbose = nlr.get<bool>("verbose", false);
            update_method_str = nlr.get<std::string>("update_method", "NoHistory");
        }else{
            threshold_to_rebalance = 0.02;
            verbose = false;
        }

        // Simple model for memory usage from load balance - just particles
        double max_n_ptl_on_rank = max_mem_redist_gb*1024*1024*1024/80.0; // ~ 80 B per particle

        if(weighting_algorithm_str=="Fortran"){
            default_weighting_algorithm = WeightingAlgorithm::Fortran;
        }else if(weighting_algorithm_str=="ParticleBalance"){
            default_weighting_algorithm = WeightingAlgorithm::ParticleBalance;
        }else if(weighting_algorithm_str=="SingleRegionBalance"){
            default_weighting_algorithm = WeightingAlgorithm::SingleRegionBalance;
        }else{
            exit_XGC("\nError: weighting_algorithm input not valid.");
        }

#ifdef USE_MPI
        // If planes are sync'd, use intpl_comm to coordinate them; otherwise, use MPI_COMM_SELF
        MPI_Comm inter_period_comm = sync_planes ? pol_decomp.mpi.intpl_comm : MPI_COMM_SELF;

        // Set up load balance regions
        if(default_weighting_algorithm == WeightingAlgorithm::SingleRegionBalance){
            LoadRegion::UpdateMethod update_method;
            if(update_method_str=="NoHistory"){
                update_method = LoadRegion::UpdateMethod::NoHistory;
            }else if(update_method_str=="ExpHistory"){
                update_method = LoadRegion::UpdateMethod::ExpHistory;
            }else{
                exit_XGC("\nError: update_method input not valid.");
            }

            regions.push_back(LoadRegion(pol_decomp.nnodes_on_plane, inter_period_comm, pol_decomp.mpi.plane_comm, update_method, verbose,
                                         "Push",
                                         {"ipc1:PUSHE",
                                          "ipc1:PUSHI",
                                          "ipc2:PUSHE",
                                          "ipc2:PUSHI"}));
        }

        // Hard-code for now
        ConstraintOption constraint_option = ConstraintOption::ParticleCount;

        if(constraint_option == ConstraintOption::ParticleCount){
            // Normalize by n_planes since we are using the sum of planes as the constraint
            int n_periods;
            MPI_Comm_size(inter_period_comm, &n_periods);
            constraint1_max = max_n_ptl_on_rank*n_periods;
        }else{
            exit_XGC("\nError: ParticleConstraint is only available ConstraintOption.\n");
        }
#endif
    }

    void rebalance(const Simulation<DeviceType>& sml, const Grid<DeviceType>& grid, const MagneticField<DeviceType>& magnetic_field,
                   const VelocityGrid& vgrid, Plasma& plasma, DomainDecomposition<DeviceType>& pol_decomp, ReweightOption reweight_option,
                   WeightingAlgorithm weighting_algorithm = WeightingAlgorithm::Default){

        // Skip rebalancing if poloidal decomposition is turned off. Normally we should put the function in
        // an if statement rather than using return; kept here in case we might want some balancing for some reason
        // even without poloidal decomposition
        if(pol_decomp.pol_decomp==false) return;

        // Rebalance is trivial if there is only one rank on a plane (generalize this from no-MPI case)
#ifdef USE_MPI

        if(weighting_algorithm == WeightingAlgorithm::Default){
            weighting_algorithm = default_weighting_algorithm;
        }

        if(default_weighting_algorithm==WeightingAlgorithm::Fortran){
            // Override input algorithm; if the default is Fortran, it should always use the fortran load balancer
            weighting_algorithm = default_weighting_algorithm;
        }

        if(is_rank_zero() && verbose){
            if(weighting_algorithm==WeightingAlgorithm::Fortran) printf("\nLoad balance called with Fortran algorithm\n");
            else if(weighting_algorithm==WeightingAlgorithm::ParticleBalance) printf("\nLoad balance called with Ptl algorithm\n");
            else if(weighting_algorithm==WeightingAlgorithm::SingleRegionBalance) printf("\nLoad balance called with SingleRegion algorithm\n");
        }

        GPTLstart("REBALANCE");
        if(weighting_algorithm!=WeightingAlgorithm::Fortran){
            GPTLstart("PTL_COUNT");
            plasma.for_all_nonadiabatic_species([&](Species<DeviceType>& species){
                long long int total_n_ptl = species.get_total_n_ptl();
                int max_n_ptl = species.get_max_n_ptl();
                double avg_n_ptl = (double)(total_n_ptl)/pol_decomp.mpi.nranks;
                double max_ratio = max_n_ptl/avg_n_ptl;
                if(is_rank_zero()) printf("  Species %d: (max/avg ptl per rank = %1.2f); total n_ptl = %lld\n", species.idx, max_ratio, total_n_ptl);
            }, Plasma::NoDevicePtl);
            GPTLstop("PTL_COUNT");

            GPTLstart("PTL_COUNT_PER_NODE");
            auto ptl_count = count_all_ptl_per_node(grid, magnetic_field, plasma);
            GPTLstop("PTL_COUNT_PER_NODE");

            if(weighting_algorithm!=WeightingAlgorithm::ParticleBalance){
                if(model_is_initialized()){
                    // Update the model with the latest timing
                    TIMER("LOAD_BAL_UPDATE",
                        update_model(pol_decomp.gvid0_pid_h) );
                }else{
                    if(is_rank_zero() && verbose) printf("Initializing timing model, no partition proposal yet.");
                    initialize_model();
                    GPTLstop("REBALANCE");
                    return;
                }
            }

            // With the updated model, proposed a new partition
            GPTLstart("LOAD_BAL_NEW_PART");
            if (is_rank_zero()) propose_new_partition(ptl_count, weighting_algorithm);
            GPTLstop("LOAD_BAL_NEW_PART");
        }

        GPTLstart("LOAD_BAL_REBAL");
        double f0_cost = sml.f0_grid ? sum_view(plasma.f0_node_cost) : 0.0;
        if(will_rebalance(reweight_option, weighting_algorithm, f0_cost)){
            // Save the old partition, it is used to save some communication in the f0 redistribution
            View<int*,CLayout,HostType> old_partition(NoInit("old_partition"), pol_decomp.gvid0_pid_h.layout());
            Kokkos::deep_copy(old_partition, pol_decomp.gvid0_pid_h);

            // Update pol_decomp
            TIMER("LOAD_BAL_SET_NEW",
                set_new_partition(sml, grid, magnetic_field, vgrid, plasma, pol_decomp, weighting_algorithm) );

            // Update particles and f0
            TIMER("LOAD_BAL_REDIST",
                redistribute_load(sml, grid, magnetic_field, vgrid, plasma, pol_decomp, old_partition) );
        }
        GPTLstop("LOAD_BAL_REBAL");
        GPTLstop("REBALANCE");
#endif
    }


    // More generic rebalance
    void rebalance(DomainDecomposition<DeviceType>& pol_decomp, const View<double*,CLayout,HostType>& constraint, const std::vector<double>& timings, double& load_imbalance, View<double*,HostType>& model_belief){
#ifdef USE_MPI
        update_model(pol_decomp.gvid0_pid_h, timings);

        load_imbalance = regions[0].get_observed_load_imbalance();
        model_belief = regions[0].get_estimated_time_per_vertex();

        // Propose new partition
        if (is_rank_zero()) propose_new_partition(constraint, default_weighting_algorithm);

        double f0_cost = 0.0;
        if(will_rebalance(ReweightOption::IfBetter, default_weighting_algorithm, f0_cost)){
            // Broadcast proposal to all ranks
            MPI_Bcast(proposed_partition.data(), proposed_partition.size(), MPI_INT, 0, comm);

            // Copy proposal to pol_decomp
            Kokkos::deep_copy(pol_decomp.gvid0_pid_h, proposed_partition);

            // Update bounds and device copy of the decomposition array
            pol_decomp.update_pol_decomp();

            // If there were a real load, redistribute it here
        }
#endif
    }
};

#endif