DataInserterUtility.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once

 
#include "shamcomm/worldInfo.hpp"
#include "shamrock/patch/PatchDataLayer.hpp"
#include "shamrock/scheduler/PatchScheduler.hpp"
#include "shamrock/scheduler/ReattributeDataUtility.hpp"
#include "shamrock/scheduler/SerialPatchTree.hpp"
#include <mpi.h>
 
namespace shamrock {

    class DataInserterUtility {
        PatchScheduler &sched; 
 
        public:
        DataInserterUtility(PatchScheduler &sched) : sched(sched) {}
 
        inline void balance_load(std::function<void(void)> load_balance_update) {
            // it seems that we need multiple runs to converge the load balance
            ON_RANK_0(
                logger::info_ln(
                    "DataInserterUtility", "---------------------------------------------"));
 
            u64 npatch_last = sched.patch_list.global.size();
 
            u32 max_runs = 30;
 
            int i = 0;
            for (; i < max_runs; i++) {
                if (shamcomm::world_rank() == 0) {
                    logger::info_ln("DataInserterUtility", "Compute load ...");
                }
 
                load_balance_update();
 
                shamcomm::mpi::Barrier(MPI_COMM_WORLD);
 
                if (shamcomm::world_rank() == 0) {
                    logger::info_ln("DataInserterUtility", "run scheduler step ...");
                }
 
                sched.scheduler_step(false, false);
                sched.scheduler_step(true, true);
 
                u64 npatch_new = sched.patch_list.global.size();
                if (npatch_new != npatch_last) {
                    npatch_last = npatch_new;
                } else {
                    break;
                }
            }
 
            ON_RANK_0(
                logger::info_ln(
                    "DataInserterUtility",
                    "patch count stable after",
                    i + 1,
                    "runs npatch =",
                    npatch_last));
 
            ON_RANK_0(
                logger::info_ln(
                    "DataInserterUtility", "---------------------------------------------"));
        }

        template<class Tvec>
        u64 push_patch_data(
            shamrock::patch::PatchDataLayer &pdat_ins,
            std::string main_field_name,
            u32 split_threshold,
            std::function<void(void)> load_balance_update) {
            using namespace shamrock::patch;
 
            u64 pdat_ob_cnt = pdat_ins.get_obj_cnt();
 
            u64 sum_push = shamalgs::collective::allreduce_sum(pdat_ob_cnt);
            if (shamcomm::world_rank() == 0) {
                logger::info_ln("DataInserterUtility", "pushing data in scheduler, N =", sum_push);
            }
 
            if (pdat_ob_cnt < split_threshold) {
                bool should_insert = true;
                sched.for_each_local_patchdata([&](const Patch &p, PatchDataLayer &pdat) {
                    if (should_insert) {
                        pdat.insert_elements(pdat_ins);
                        should_insert = false; // We insert only in first patch (no duplicates)
                    }
                });
            } else {
                shambase::throw_unimplemented(
                    "Not implemented yet please keep the obj count to be "
                    "inserted below the split_threshold, sorrrrrry ...");
            }
 
            if (shamcomm::world_rank() == 0) {
                logger::info_ln("DataInserterUtility", "reattributing data ...");
            }
 
            shambase::Timer treatrib;
            treatrib.start();
            // move data into the correct patches
            SerialPatchTree<Tvec> sptree = SerialPatchTree<Tvec>::build(sched);
            ReattributeDataUtility reatrib(sched);
            sptree.attach_buf();
            reatrib.reatribute_patch_objects(sptree, main_field_name);
            sched.check_patchdata_locality_correctness();
 
            treatrib.stop();
            if (shamcomm::world_rank() == 0) {
                logger::info_ln(
                    "DataInserterUtility", "reattributing data done in ", treatrib.get_time_str());
            }
            shamcomm::mpi::Barrier(MPI_COMM_WORLD);
 
            balance_load(load_balance_update);
 
            return sum_push;
        }
    };
 
} // namespace shamrock