Skip to content

[Scheduler][PIC] create pdat as list #1514

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aserhani wants to merge 11 commits into
base: main
Choose a base branch
from
Draft

[Scheduler][PIC] create pdat as list #1514

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
06f3fcf
created a deque container of patch_data_layer_layout
aserhani Dec 31, 2025
0973e56
made the layer_idx variable homogeneous, added shamrock::format where…
aserhani Dec 31, 2025
d81fbd4
created patch_data_list in PatchScheduler
aserhani Dec 31, 2025
2d6f3a6
replaced patch_data by patch_data_list in PatchScheduler
aserhani Dec 31, 2025
7e08f2f
updated .gitignore with lsp file
aserhani Dec 31, 2025
01bf7a8
initialized patch_data_list with default constructor
aserhani Dec 31, 2025
7e33c78
fixed bugs in PatchScheduler
aserhani Dec 31, 2025
9dc2262
renamed patchdata
aserhani Jan 2, 2026
5de280a
made default more hard-coded layer_idx values
aserhani Jan 2, 2026
c502f88
set patch_data as reference to patch_data_list[0]
aserhani Jan 2, 2026
3f3f5a5
reverted "updated .gitignore with lsp file"
aserhani Jan 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 20 additions & 19 deletions src/shampylib/src/pyShamrockCtx.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,26 +286,27 @@ Register_pymod(pyshamrockctxinit) {

py::dict dic_out;

for (auto fname : ctx.pdl->get_field_names()) {
append_to_map<f32>(fname, data, dic_out);
append_to_map<f32_2>(fname, data, dic_out);
append_to_map<f32_3>(fname, data, dic_out);
append_to_map<f32_4>(fname, data, dic_out);
append_to_map<f32_8>(fname, data, dic_out);
append_to_map<f32_16>(fname, data, dic_out);
append_to_map<f64>(fname, data, dic_out);
append_to_map<f64_2>(fname, data, dic_out);
append_to_map<f64_3>(fname, data, dic_out);
append_to_map<f64_4>(fname, data, dic_out);
append_to_map<f64_8>(fname, data, dic_out);
append_to_map<f64_16>(fname, data, dic_out);
append_to_map<u32>(fname, data, dic_out);
append_to_map<u64>(fname, data, dic_out);
append_to_map<u32_3>(fname, data, dic_out);
append_to_map<u64_3>(fname, data, dic_out);
append_to_map<i64_3>(fname, data, dic_out);
for (u32 layer_idx = 0; layer_idx < ctx.pdl_list.size(); layer_idx++) {
for (auto fname : ctx.pdl_list.at(layer_idx)->get_field_names()) {
append_to_map<f32>(fname, data, dic_out);
append_to_map<f32_2>(fname, data, dic_out);
append_to_map<f32_3>(fname, data, dic_out);
append_to_map<f32_4>(fname, data, dic_out);
append_to_map<f32_8>(fname, data, dic_out);
append_to_map<f32_16>(fname, data, dic_out);
append_to_map<f64>(fname, data, dic_out);
append_to_map<f64_2>(fname, data, dic_out);
append_to_map<f64_3>(fname, data, dic_out);
append_to_map<f64_4>(fname, data, dic_out);
append_to_map<f64_8>(fname, data, dic_out);
append_to_map<f64_16>(fname, data, dic_out);
append_to_map<u32>(fname, data, dic_out);
append_to_map<u64>(fname, data, dic_out);
append_to_map<u32_3>(fname, data, dic_out);
append_to_map<u64_3>(fname, data, dic_out);
append_to_map<i64_3>(fname, data, dic_out);
}
}
Comment on lines +289 to 309
Copy link

@gemini-code-assist gemini-code-assist bot Jan 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The current implementation for collecting data across multiple layers is incorrect. ctx.allgather_data() is called only once before the loops, and it only gathers data for layer 0 due to the current implementation of ShamrockCtx::allgather_data. The loop then iterates through field names of all layers, but attempts to find them in the data of layer 0, which is wrong. Additionally, if different layers have fields with the same name, their data will be overwritten in the output dictionary.

To fix this, you should:

  1. Modify ShamrockCtx::gather_data and ShamrockCtx::allgather_data in ShamrockCtx.hpp to accept a layer_idx.
  2. In this file, move the allgather_data call inside the loop to fetch data for each layer.
  3. Consider how to handle fields with the same name across different layers. Returning a list of dictionaries (one per layer) or prefixing field names with the layer index would be a robust solution.


return dic_out;
})
.def("get_patch_list_global", &ShamrockCtx::get_patch_list_global);
Expand Down
98 changes: 52 additions & 46 deletions src/shamrock/include/shamrock/scheduler/PatchScheduler.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include "shamrock/solvergraph/NodeSetEdge.hpp"
#include "shamrock/solvergraph/PatchDataLayerRefs.hpp"
#include <nlohmann/json.hpp>
#include <deque>
#include <unordered_set>
#include <fstream>
#include <functional>
Expand Down Expand Up @@ -64,23 +65,25 @@ class PatchScheduler {
using PatchTree = shamrock::scheduler::PatchTree;
using SchedulerPatchData = shamrock::scheduler::SchedulerPatchData;

std::shared_ptr<shamrock::patch::PatchDataLayerLayout> pdl_ptr;
std::deque<std::shared_ptr<shamrock::patch::PatchDataLayerLayout>> pdl_ptr_list;
std::deque<SchedulerPatchData> patch_data_list; ///< handle the data of the patches of the scheduler for every layer
SchedulerPatchData & patch_data; ///< first layer patch data (legacy) patch_data_list.at(0)

u64 crit_patch_split; ///< splitting limit (if load value > crit_patch_split => patch split)
u64 crit_patch_merge; ///< merging limit (if load value < crit_patch_merge => patch merge)

SchedulerPatchList patch_list; ///< handle the list of the patches of the scheduler
SchedulerPatchData patch_data; ///< handle the data of the patches of the scheduler

PatchTree patch_tree; ///< handle the tree structure of the patches

// using unordered set is not an issue since we use the find command after
std::unordered_set<u64> owned_patch_id; ///< list of owned patch ids updated with
///< (owned_patch_id = patch_list.build_local())

inline shamrock::patch::PatchDataLayerLayout &pdl() { return shambase::get_check_ref(pdl_ptr); }
inline shamrock::patch::PatchDataLayerLayout &pdl(u32 layer_idx = 0) { return shambase::get_check_ref(pdl_ptr_list.at(layer_idx)); }

inline std::shared_ptr<shamrock::patch::PatchDataLayerLayout> get_layout_ptr() const {
return pdl_ptr;
inline std::shared_ptr<shamrock::patch::PatchDataLayerLayout> get_layout_ptr(u32 layer_idx = 0) const {
return pdl_ptr_list.at(layer_idx);
}

/**
Expand All @@ -96,7 +99,7 @@ class PatchScheduler {
void free_mpi_required_types();

PatchScheduler(
const std::shared_ptr<shamrock::patch::PatchDataLayerLayout> &pdl_ptr,
const std::deque<std::shared_ptr<shamrock::patch::PatchDataLayerLayout>> &pdl_ptr_list,
u64 crit_split,
u64 crit_merge);

Expand All @@ -113,10 +116,10 @@ class PatchScheduler {
}

template<class vectype>
std::tuple<vectype, vectype> get_box_tranform();
std::tuple<vectype, vectype> get_box_tranform(u32 layer_idx = 0);

template<class vectype>
std::tuple<vectype, vectype> get_box_volume();
std::tuple<vectype, vectype> get_box_volume(u32 layer_idx = 0);

bool should_resize_box(bool node_in);

Expand All @@ -130,15 +133,18 @@ class PatchScheduler {
template<class vectype>
void set_coord_domain_bound(vectype bmin, vectype bmax) {

if (!pdl().check_main_field_type<vectype>()) {
std::invalid_argument(
std::string("the main field is not of the correct type to call this function\n")
+ "fct called : " + __PRETTY_FUNCTION__
+ "current patch data layout : " + pdl().get_description_str());
for (u32 layer_idx = 0; layer_idx < pdl_ptr_list.size(); layer_idx++) {
if (!pdl(layer_idx).check_main_field_type<vectype>()) {
std::invalid_argument(
std::string("the main field is not of the correct type to call this function\n")
+ "fct called : " + __PRETTY_FUNCTION__
+ shambase::format("current patch data layout of index {} : ", layer_idx)
+ pdl(layer_idx).get_description_str()
);
Comment on lines +138 to +143
Copy link

@gemini-code-assist gemini-code-assist bot Jan 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The std::invalid_argument exception is constructed but not thrown. This makes the type check ineffective, and the program will proceed even if there's a type mismatch, potentially leading to undefined behavior.

                throw std::invalid_argument(
                    std::string("the main field is not of the correct type to call this function\n")
                    + "fct called : " + __PRETTY_FUNCTION__
                    + shambase::format("current patch data layout of index {} : ", layer_idx)
                    + pdl(layer_idx).get_description_str()
                );

}
patch_data_list.at(layer_idx).sim_box.set_bounding_box<vectype>({bmin, bmax});
}

patch_data.sim_box.set_bounding_box<vectype>({bmin, bmax});

shamlog_debug_ln("PatchScheduler", "box resized to :", bmin, bmax);
}

Expand All @@ -165,14 +171,14 @@ class PatchScheduler {
set_coord_domain_bound(a, b);
}

std::string format_patch_coord(shamrock::patch::Patch p);
std::string format_patch_coord(shamrock::patch::Patch p, u32 layer_idx = 0);

void check_patchdata_locality_corectness();

[[deprecated]]
void dump_local_patches(std::string filename);

std::vector<std::unique_ptr<shamrock::patch::PatchDataLayer>> gather_data(u32 rank);
std::vector<std::unique_ptr<shamrock::patch::PatchDataLayer>> gather_data(u32 rank, u32 layer_idx = 0);

/**
* @brief add patch to the scheduler
Expand All @@ -197,7 +203,7 @@ class PatchScheduler {
void add_root_patch();

[[deprecated]]
void sync_build_LB(bool global_patch_sync, bool balance_load);
void sync_build_LB(bool global_patch_sync, bool balance_load, u32 layer_idx = 0);

template<class vec>
inline shamrock::patch::PatchCoordTransform<vec> get_patch_transform() {
Expand All @@ -224,9 +230,9 @@ class PatchScheduler {
* @param fct
*/
template<class Function>
inline void for_each_patch_data(Function &&fct) {
inline void for_each_patch_data(Function &&fct, u32 layer_idx = 0) {

patch_data.for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
patch_data_list.at(layer_idx).for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
shamrock::patch::Patch &cur_p
= patch_list.global[patch_list.id_patch_to_global_idx[patch_id]];

Expand All @@ -237,9 +243,9 @@ class PatchScheduler {
}

template<class Function>
inline void for_each_patch(Function &&fct) {
inline void for_each_patch(Function &&fct, u32 layer_idx = 0) {

patch_data.for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
patch_data_list.at(layer_idx).for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
shamrock::patch::Patch &cur_p
= patch_list.global[patch_list.id_patch_to_global_idx[patch_id]];

Expand Down Expand Up @@ -267,17 +273,17 @@ class PatchScheduler {
}

inline void for_each_local_patchdata(
std::function<void(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &)> fct) {
std::function<void(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &)> fct, u32 layer_idx = 0) {
for (shamrock::patch::Patch p : patch_list.local) {
if (!p.is_err_mode()) {
fct(p, patch_data.get_pdat(p.id_patch));
fct(p, patch_data_list.at(layer_idx).get_pdat(p.id_patch));
}
}
}

inline void for_each_local_patch_nonempty(
std::function<void(const shamrock::patch::Patch &)> fct) {
patch_data.for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
std::function<void(const shamrock::patch::Patch &)> fct, u32 layer_idx = 0) {
patch_data_list.at(layer_idx).for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
shamrock::patch::Patch &cur_p
= patch_list.global[patch_list.id_patch_to_global_idx.at(patch_id)];

Expand All @@ -294,8 +300,8 @@ class PatchScheduler {
}

inline void for_each_patchdata_nonempty(
std::function<void(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &)> fct) {
patch_data.for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
std::function<void(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &)> fct, u32 layer_idx = 0) {
patch_data_list.at(layer_idx).for_each_patchdata([&](u64 patch_id, shamrock::patch::PatchDataLayer &pdat) {
shamrock::patch::Patch &cur_p
= patch_list.global[patch_list.id_patch_to_global_idx.at(patch_id)];

Expand All @@ -307,13 +313,13 @@ class PatchScheduler {

template<class T>
inline shambase::DistributedData<T> map_owned_patchdata(
std::function<T(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &pdat)> fct) {
std::function<T(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &pdat)> fct, u32 layer_idx = 0) {
shambase::DistributedData<T> ret;

using namespace shamrock::patch;
for_each_patch_data([&](u64 id_patch, Patch cur_p, PatchDataLayer &pdat) {
ret.add_obj(id_patch, fct(cur_p, pdat));
});
}, layer_idx);

return ret;
}
Expand Down Expand Up @@ -344,26 +350,26 @@ class PatchScheduler {

template<class T>
inline shambase::DistributedData<T> map_owned_patchdata_fetch_simple(
std::function<T(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &pdat)> fct) {
std::function<T(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &pdat)> fct, u32 layer_idx = 0) {
shambase::DistributedData<T> ret;

using namespace shamrock::patch;
for_each_patch_data([&](u64 id_patch, Patch cur_p, PatchDataLayer &pdat) {
ret.add_obj(id_patch, fct(cur_p, pdat));
});
}, layer_idx);

return distrib_data_local_to_all_simple(ret);
}

template<class T>
inline shambase::DistributedData<T> map_owned_patchdata_fetch_load_store(
std::function<T(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &pdat)> fct) {
std::function<T(const shamrock::patch::Patch, shamrock::patch::PatchDataLayer &pdat)> fct, u32 layer_idx = 0) {
shambase::DistributedData<T> ret;

using namespace shamrock::patch;
for_each_patch_data([&](u64 id_patch, Patch cur_p, PatchDataLayer &pdat) {
ret.add_obj(id_patch, fct(cur_p, pdat));
});
}, layer_idx);

return distrib_data_local_to_all_load_store(ret);
}
Expand All @@ -380,13 +386,13 @@ class PatchScheduler {
return shamrock::patch::PatchField<T>(map_owned_patchdata_fetch_load_store(fct));
}

inline u64 get_rank_count() {
inline u64 get_rank_count(u32 layer_idx = 0) {
StackEntry stack_loc{};
using namespace shamrock::patch;
u64 num_obj = 0; // TODO get_rank_count() in scheduler
for_each_patch_data([&](u64 id_patch, Patch cur_p, PatchDataLayer &pdat) {
num_obj += pdat.get_obj_cnt();
});
}, layer_idx);

return num_obj;
}
Expand All @@ -398,7 +404,7 @@ class PatchScheduler {
}

template<class T>
inline std::unique_ptr<sycl::buffer<T>> rankgather_field(u32 field_idx) {
inline std::unique_ptr<sycl::buffer<T>> rankgather_field(u32 field_idx, u32 layer_idx = 0) {
StackEntry stack_loc{};
std::unique_ptr<sycl::buffer<T>> ret;

Expand Down Expand Up @@ -427,7 +433,7 @@ class PatchScheduler {

ptr += pdat.get_obj_cnt() * nvar;
}
});
}, layer_idx);
}

return ret;
Expand Down Expand Up @@ -456,7 +462,7 @@ class PatchScheduler {
// }

template<class Function, class Pfield>
inline void compute_patch_field(Pfield &field, MPI_Datatype &dtype, Function &&lambda) {
inline void compute_patch_field(Pfield &field, MPI_Datatype &dtype, Function &&lambda, u32 layer_idx = 0) {
field.local_nodes_value.resize(patch_list.local.size());

for (u64 idx = 0; idx < patch_list.local.size(); idx++) {
Expand All @@ -467,21 +473,21 @@ class PatchScheduler {
field.local_nodes_value[idx] = lambda(
shamsys::instance::get_compute_queue(),
cur_p,
patch_data.owned_data.get(cur_p.id_patch));
patch_data_list.at(layer_idx).owned_data.get(cur_p.id_patch));
}
}

field.build_global(dtype);
}

inline auto get_node_set_edge_patchdata_layer_refs() {
inline auto get_node_set_edge_patchdata_layer_refs(u32 layer_idx = 0) {
shamrock::solvergraph::NodeSetEdge<shamrock::solvergraph::PatchDataLayerRefs> node_set_edge(
[&](shamrock::solvergraph::PatchDataLayerRefs &edge) {
edge.free_alloc();
using namespace shamrock::patch;
for_each_patchdata_nonempty([&](Patch cur_p, PatchDataLayer &pdat) {
edge.patchdatas.add_obj(cur_p.id_patch, std::ref(pdat));
});
}, layer_idx);
});

return std::make_shared<decltype(node_set_edge)>(std::move(node_set_edge));
Expand All @@ -493,15 +499,15 @@ class PatchScheduler {
* @param coords coordinates of the patch
* @return u64 the id of the made patch
*/
std::vector<u64> add_root_patches(std::vector<shamrock::patch::PatchCoord<3>> coords);
std::vector<u64> add_root_patches(std::vector<shamrock::patch::PatchCoord<3>> coords, u32 layer_idx = 0);

shamrock::patch::SimulationBoxInfo &get_sim_box() { return patch_data.sim_box; }
shamrock::patch::SimulationBoxInfo &get_sim_box() { return patch_data_list.at(0).sim_box; }

nlohmann::json serialize_patch_metadata();

private:
void split_patches(std::unordered_set<u64> split_rq);
void merge_patches(std::unordered_set<u64> merge_rq);
void split_patches(std::unordered_set<u64> split_rq, u32 layer_idx = 0);
void merge_patches(std::unordered_set<u64> merge_rq, u32 layer_idx = 0);

void set_patch_pack_values(std::unordered_set<u64> merge_rq);
};
Loading
Loading