From 4e9e9314eacb9919854ea4587f650505ae54be90 Mon Sep 17 00:00:00 2001 From: LudwigBoess Date: Sat, 27 Jun 2026 16:28:11 +0000 Subject: [PATCH] port of load balancing from my fork --- input.example.toml | 29 ++ src/engines/engine.hpp | 33 ++ src/engines/grpic/grpic.hpp | 6 +- src/engines/srpic/srpic.hpp | 7 +- src/framework/CMakeLists.txt | 2 + src/framework/containers/fields.h | 5 +- src/framework/containers/fields_io.cpp | 27 +- src/framework/domain/domain.h | 4 + src/framework/domain/mesh.h | 9 + src/framework/domain/metadomain.h | 17 + src/framework/domain/metadomain_chckpt.cpp | 13 +- src/framework/domain/metadomain_io.cpp | 3 + src/framework/domain/metadomain_loadbal.cpp | 428 ++++++++++++++++++++ src/framework/parameters/parameters.cpp | 46 +++ src/output/utils/writers.cpp | 18 +- src/output/utils/writers.h | 3 +- src/output/writer.cpp | 83 +++- src/output/writer.h | 11 + 18 files changed, 710 insertions(+), 34 deletions(-) create mode 100644 src/framework/domain/metadomain_loadbal.cpp diff --git a/input.example.toml b/input.example.toml index 87ca48f3c..8d6f25117 100644 --- a/input.example.toml +++ b/input.example.toml @@ -28,6 +28,35 @@ # @example: [2, 2, 2] (total of 8 domains) decomposition = "" + # Diffusion-style dynamic load balancing (Cartesian metrics only). + # Domain boundaries between MPI neighbors are nudged to equalize the + # active-particle count per rank. All inter-rank traffic uses only the + # existing nearest-neighbor field/particle communication paths. + [simulation.domain.load_balance] + # Enable dynamic load balancing + # @type: bool + # @default: false + enable = "" + # Run the rebalancer every `interval` timesteps (0 disables) + # @type: int + # @default: 0 + interval = "" + # Dimensions along which load is redistributed (1 = x1, 2 = x2, 3 = x3) + # @type: array of int, subset of [1, 2, 3] + # @default: [1] + dimensions = "" + # Skip rebalancing along a dim when (max - min) / mean of the per-slice + # particle count is below this fraction + # @type: float + # @default: 0.1 + tolerance = "" + # Maximum cell-shift per interior boundary per event; clamped at compile + # time to N_GHOSTS so the migrating field strip is already cached in the + # rank's ghost zone. + # @type: int + # @default: N_GHOSTS + max_shift = "" + [grid] # Spatial resolution of the grid # @required diff --git a/src/engines/engine.hpp b/src/engines/engine.hpp index 057bfcb5a..161339a9d 100644 --- a/src/engines/engine.hpp +++ b/src/engines/engine.hpp @@ -253,6 +253,7 @@ namespace ntt { "ParticlePusher", "FieldBoundaries", "ParticleBoundaries", "Communications", "Injector", "Custom", + "LoadBalance", "ParticleSort", "Output", "Checkpoint" }, []() { @@ -267,6 +268,17 @@ namespace ntt { const auto clear_interval = m_params.template get( "particles.clear_interval"); + const auto lb_enable = m_params.template get( + "simulation.domain.load_balance.enable"); + const auto lb_interval = m_params.template get( + "simulation.domain.load_balance.interval"); + const auto lb_dim_mask = m_params.template get( + "simulation.domain.load_balance.dim_mask"); + const auto lb_tolerance = m_params.template get( + "simulation.domain.load_balance.tolerance"); + const auto lb_max_shift = m_params.template get( + "simulation.domain.load_balance.max_shift"); + // main algorithm loop while (step < max_steps) { // run the engine-dependent algorithm step @@ -282,6 +294,27 @@ namespace ntt { }); timers.stop("Custom"); } + if constexpr (MetricClass) { + if (lb_enable and lb_interval > 0 and (step + 1) % lb_interval == 0 and + lb_dim_mask != 0u) { + timers.start("LoadBalance"); + m_metadomain.Rebalance(lb_dim_mask, + lb_tolerance, + static_cast(lb_max_shift)); + timers.stop("LoadBalance"); + } + } + // Sort particles last — after step_forward, CustomPostStep, and + // LoadBalance — so the tile layout (and dead-particle compaction) the + // next step's deposit relies on reflects every particle change made this + // step: moving-window shift/injection and in-place dead-tagging done in + // CustomPostStep, plus any rebalance migration. + timers.start("ParticleSort"); + m_metadomain.runOnLocalDomains([this](auto& dom) { + m_metadomain.SortParticles(time, step, m_params, dom); + }); + timers.stop("ParticleSort"); + auto print_prtl_clear = (clear_interval > 0 and step % clear_interval == 0 and step > 0); diff --git a/src/engines/grpic/grpic.hpp b/src/engines/grpic/grpic.hpp index 621592955..3487d2867 100644 --- a/src/engines/grpic/grpic.hpp +++ b/src/engines/grpic/grpic.hpp @@ -613,9 +613,9 @@ namespace ntt { timers.stop("FieldBoundaries"); } - timers.start("ParticleSort"); - m_metadomain.SortParticles(time, step, m_params, dom); - timers.stop("ParticleSort"); + // NOTE: particle sorting is intentionally NOT done here. It runs once per + // step in the engine loop (Engine::run) after CustomPostStep and + // LoadBalance — see the SRPIC engine for the rationale. /** * Finally: em0::B at n-1/2 diff --git a/src/engines/srpic/srpic.hpp b/src/engines/srpic/srpic.hpp index 94a8949c1..d0352318e 100644 --- a/src/engines/srpic/srpic.hpp +++ b/src/engines/srpic/srpic.hpp @@ -182,9 +182,10 @@ namespace ntt { timers.stop("Injector"); } - timers.start("ParticleSort"); - m_metadomain.SortParticles(time, step, m_params, dom); - timers.stop("ParticleSort"); + // NOTE: particle sorting is intentionally NOT done here. It runs once per + // step in the engine loop (Engine::run) after CustomPostStep and + // LoadBalance, so the layout the next deposit uses reflects window + // shifts/injection and dead-tagging performed in CustomPostStep. } }; diff --git a/src/framework/CMakeLists.txt b/src/framework/CMakeLists.txt index 5343f3d8a..377c1f98a 100644 --- a/src/framework/CMakeLists.txt +++ b/src/framework/CMakeLists.txt @@ -19,6 +19,7 @@ # * domain/metadomain_stats.cpp # * domain/metadomain_io.cpp # * domain/metadomain_reshape.cpp +# * domain/metadomain_loadbal.cpp # * containers/particles.cpp # * containers/particles_comm.cpp # * containers/particles_io.cpp @@ -59,6 +60,7 @@ set(SOURCES ${SRC_DIR}/domain/metadomain_sort.cpp ${SRC_DIR}/domain/metadomain_stats.cpp ${SRC_DIR}/domain/metadomain_reshape.cpp + ${SRC_DIR}/domain/metadomain_loadbal.cpp ${SRC_DIR}/containers/particles.cpp ${SRC_DIR}/containers/particles_sort.cpp ${SRC_DIR}/containers/fields.cpp) diff --git a/src/framework/containers/fields.h b/src/framework/containers/fields.h index 4acabf7b4..2dcfd0a7e 100644 --- a/src/framework/containers/fields.h +++ b/src/framework/containers/fields.h @@ -175,7 +175,10 @@ namespace ntt { void CheckpointRead(adios2::IO&, adios2::Engine&, const adios2::Box&); - void CheckpointWrite(adios2::IO&, adios2::Engine&) const; + void CheckpointWrite(adios2::IO&, + adios2::Engine&, + const std::vector&, + const std::vector&) const; #endif }; diff --git a/src/framework/containers/fields_io.cpp b/src/framework/containers/fields_io.cpp index 683519493..a17ab1430 100644 --- a/src/framework/containers/fields_io.cpp +++ b/src/framework/containers/fields_io.cpp @@ -62,13 +62,27 @@ namespace ntt { } template - void Fields::CheckpointWrite(adios2::IO& io, adios2::Engine& writer) const { + void Fields::CheckpointWrite( + adios2::IO& io, + adios2::Engine& writer, + const std::vector& local_shape, + const std::vector& local_offset) const { logger::Checkpoint("Writing fields checkpoint", HERE); - out::WriteNDField(io, writer, "em", em); + // Per-rank slab: re-set the variable selection to track the (possibly + // rebalanced) local layout. The component axis is always full. + auto build_range = [&](unsigned short ncomp) { + auto start = adios2::Dims(local_offset.begin(), local_offset.end()); + auto count = adios2::Dims(local_shape.begin(), local_shape.end()); + start.push_back(0); + count.push_back(ncomp); + return adios2::Box(start, count); + }; + + out::WriteNDField(io, writer, "em", em, build_range(6)); if (S == ntt::SimEngine::GRPIC) { - out::WriteNDField(io, writer, "em0", em0); - out::WriteNDField(io, writer, "cur", cur); + out::WriteNDField(io, writer, "em0", em0, build_range(6)); + out::WriteNDField(io, writer, "cur", cur, build_range(3)); } } @@ -81,7 +95,10 @@ namespace ntt { template void Fields::CheckpointRead(adios2::IO&, \ adios2::Engine&, \ const adios2::Box&); \ - template void Fields::CheckpointWrite(adios2::IO&, adios2::Engine&) const; + template void Fields::CheckpointWrite(adios2::IO&, \ + adios2::Engine&, \ + const std::vector&, \ + const std::vector&) const; FIELDS_CHECKPOINTS(Dim::_1D, SimEngine::SRPIC) FIELDS_CHECKPOINTS(Dim::_2D, SimEngine::SRPIC) diff --git a/src/framework/domain/domain.h b/src/framework/domain/domain.h index 747ab4576..0b708e555 100644 --- a/src/framework/domain/domain.h +++ b/src/framework/domain/domain.h @@ -165,6 +165,10 @@ namespace ntt { m_neighbor_idx[dir] = idx; } + void set_offset_ncells(const std::vector& off) { + m_offset_ncells = off; + } + /* printer overload ----------------------------------------------------- */ auto Report() const -> std::string { std::string report; diff --git a/src/framework/domain/mesh.h b/src/framework/domain/mesh.h index aa2eea726..50449ae49 100644 --- a/src/framework/domain/mesh.h +++ b/src/framework/domain/mesh.h @@ -65,6 +65,15 @@ namespace ntt { new (&metric) M { this->m_resolution, new_extent, m_metric_params_raw }; } + void set_resolution_and_extent(const std::vector& new_res, + const boundaries_t& new_extent) { + raise::ErrorIf(new_res.size() != D, "invalid resolution dim", HERE); + this->m_resolution = new_res; + m_extent = new_extent; + metric.~M(); + new (&metric) M { this->m_resolution, m_extent, m_metric_params_raw }; + } + /** * @brief Get the intersection of the mesh with a box * @param box physical extent diff --git a/src/framework/domain/metadomain.h b/src/framework/domain/metadomain.h index 8af9e6e0b..f090dade1 100644 --- a/src/framework/domain/metadomain.h +++ b/src/framework/domain/metadomain.h @@ -132,6 +132,23 @@ namespace ntt { /* domain update-related ------------------------------------------------ */ void ShiftByCells(int, in = in::x1); + /** + * @brief Rebalance the load (active particles) across MPI domains by + * shifting interior domain boundaries between neighbors. + * @param dim_mask bitmask: bit d (0,1,2) set => balance along dim x1/x2/x3 + * @param tolerance skip if (max-min)/mean of the per-slice load is below + * this fraction + * @param max_shift_cells per-event cap for any single boundary movement, + * additionally clamped to N_GHOSTS so the field strip we need is already + * present in the local ghost zone + * @note Only neighbor communication is used (CommunicateFields ghosts + + * CommunicateParticles). + */ + void Rebalance(unsigned int dim_mask, + real_t tolerance, + ncells_t max_shift_cells) + requires(MetricClass); + /* output-related ------------------------------------------------------- */ #if defined(OUTPUT_ENABLED) void InitWriter(adios2::ADIOS*, const SimulationParams&); diff --git a/src/framework/domain/metadomain_chckpt.cpp b/src/framework/domain/metadomain_chckpt.cpp index 1116df85e..fbf0f11cd 100644 --- a/src/framework/domain/metadomain_chckpt.cpp +++ b/src/framework/domain/metadomain_chckpt.cpp @@ -125,8 +125,19 @@ namespace ntt { } params.saveTOML(g_checkpoint_writer.written().back().second, current_time); + // Recompute the local with-ghosts shape/offset every step so the + // ADIOS variable selection tracks any rebalance that has happened + // since InitCheckpointWriter. + std::vector loc_off_with_ghosts; + for (auto d { 0u }; d < M::Dim; ++d) { + loc_off_with_ghosts.push_back( + local_domain->offset_ncells()[d] + + 2 * N_GHOSTS * local_domain->offset_ndomains()[d]); + } local_domain->fields.CheckpointWrite(g_checkpoint_writer.io(), - g_checkpoint_writer.writer()); + g_checkpoint_writer.writer(), + local_domain->mesh.n_all(), + loc_off_with_ghosts); #if !defined(MPI_ENABLED) const std::size_t dom_tot = 1, dom_offset = 0; #else diff --git a/src/framework/domain/metadomain_io.cpp b/src/framework/domain/metadomain_io.cpp index 3fc0f5f57..d53dc49a1 100644 --- a/src/framework/domain/metadomain_io.cpp +++ b/src/framework/domain/metadomain_io.cpp @@ -409,6 +409,9 @@ namespace ntt { } } } + // Refresh the writer's cached per-rank slab so that field/mesh writes + // pick up the (possibly rebalanced) current local layout. + g_writer.setLocalLayout(off_ncells_with_ghosts, loc_shape_with_ghosts); for (auto dim { 0u }; dim < M::Dim; ++dim) { const auto l_size = local_domain->mesh.n_active()[dim]; const auto l_offset = local_domain->offset_ncells()[dim]; diff --git a/src/framework/domain/metadomain_loadbal.cpp b/src/framework/domain/metadomain_loadbal.cpp new file mode 100644 index 000000000..0768acad6 --- /dev/null +++ b/src/framework/domain/metadomain_loadbal.cpp @@ -0,0 +1,428 @@ +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "traits/metric.h" +#include "utils/error.h" +#include "utils/log.h" +#include "utils/numeric.h" + +#include "framework/containers/fields.h" +#include "framework/domain/metadomain.h" +#include "framework/specialization_registry.h" + +#if defined(MPI_ENABLED) + #include "arch/mpi_aliases.h" + #include "arch/mpi_tags.h" + + #include +#endif + +#include + +#include +#include +#include +#include +#include +#include + +namespace ntt { + +#if defined(MPI_ENABLED) + // Copy old field into new (already-zero) field, applying a shift in the + // active-cell offset along each dimension. + // For new (with-ghost) index i, old (with-ghost) index = i + delta[d]. + // Cells whose old index is out of bounds are left at 0 and will be filled + // by CommunicateFields() afterwards. + template + void CopyShifted(const ndfield_mirror_t& src_h, + ndfield_t& dst_dev, + const std::vector& delta) { + auto dst_h = Kokkos::create_mirror_view(dst_dev); + Kokkos::deep_copy(dst_h, ZERO); + + if constexpr (D == Dim::_1D) { + const int new_n1 = static_cast(dst_h.extent(0)); + const int old_n1 = static_cast(src_h.extent(0)); + for (int i1 = 0; i1 < new_n1; ++i1) { + const int oi1 = i1 + delta[0]; + if (oi1 < 0 or oi1 >= old_n1) { + continue; + } + for (auto c { 0u }; c < NC; ++c) { + dst_h(i1, c) = src_h(oi1, c); + } + } + } else if constexpr (D == Dim::_2D) { + const int new_n1 = static_cast(dst_h.extent(0)); + const int new_n2 = static_cast(dst_h.extent(1)); + const int old_n1 = static_cast(src_h.extent(0)); + const int old_n2 = static_cast(src_h.extent(1)); + for (int i1 = 0; i1 < new_n1; ++i1) { + const int oi1 = i1 + delta[0]; + if (oi1 < 0 or oi1 >= old_n1) { + continue; + } + for (int i2 = 0; i2 < new_n2; ++i2) { + const int oi2 = i2 + delta[1]; + if (oi2 < 0 or oi2 >= old_n2) { + continue; + } + for (auto c { 0u }; c < NC; ++c) { + dst_h(i1, i2, c) = src_h(oi1, oi2, c); + } + } + } + } else if constexpr (D == Dim::_3D) { + const int new_n1 = static_cast(dst_h.extent(0)); + const int new_n2 = static_cast(dst_h.extent(1)); + const int new_n3 = static_cast(dst_h.extent(2)); + const int old_n1 = static_cast(src_h.extent(0)); + const int old_n2 = static_cast(src_h.extent(1)); + const int old_n3 = static_cast(src_h.extent(2)); + for (int i1 = 0; i1 < new_n1; ++i1) { + const int oi1 = i1 + delta[0]; + if (oi1 < 0 or oi1 >= old_n1) { + continue; + } + for (int i2 = 0; i2 < new_n2; ++i2) { + const int oi2 = i2 + delta[1]; + if (oi2 < 0 or oi2 >= old_n2) { + continue; + } + for (int i3 = 0; i3 < new_n3; ++i3) { + const int oi3 = i3 + delta[2]; + if (oi3 < 0 or oi3 >= old_n3) { + continue; + } + for (auto c { 0u }; c < NC; ++c) { + dst_h(i1, i2, i3, c) = src_h(oi1, oi2, oi3, c); + } + } + } + } + } + Kokkos::deep_copy(dst_dev, dst_h); + } +#endif // MPI_ENABLED + + template + void Metadomain::Rebalance(unsigned int dim_mask, + real_t tolerance, + ncells_t max_shift_cells) + requires(MetricClass) + { +#if !defined(MPI_ENABLED) + (void)dim_mask; + (void)tolerance; + (void)max_shift_cells; + return; +#else + raise::ErrorIf(l_subdomain_indices().size() != 1, + "Rebalance assumes one local subdomain per rank", + HERE); + // The theta dimension (idx 1) is bounded by the polar axis for any + // non-Cartesian metric: moving an interior boundary in theta is fine + // in principle, but the safety net here forbids it pending validation. + if constexpr (M::CoordType != ntt::Coord::Cartesian) { + raise::ErrorIf((dim_mask & (1u << 1)) != 0u, + "Rebalance along the polar axis is not supported", + HERE); + } + // strip-width is constrained by N_GHOSTS so that any new active cells + // are already present in the rank's old ghost zone (after the most + // recent ghost-cell exchange). + if (max_shift_cells > N_GHOSTS) { + max_shift_cells = static_cast(N_GHOSTS); + } + if (max_shift_cells == 0 or dim_mask == 0) { + return; + } + + const auto local_idx = l_subdomain_indices()[0]; + + /* --- 1. Allgather active particle counts per domain ------------------ */ + npart_t local_npart { 0 }; + for (const auto& sp : g_subdomains[local_idx].species) { + local_npart += sp.npart(); + } + std::vector npart_per_dom(g_ndomains, 0); + MPI_Allgather(&local_npart, + 1, + mpi::get_type(), + npart_per_dom.data(), + 1, + mpi::get_type(), + MPI_COMM_WORLD); + + /* --- 2. Project ncells / load onto each balanced dim ----------------- */ + std::vector> ncells_per_pos(M::Dim); + std::vector> load_per_pos(M::Dim); + for (auto d { 0u }; d < M::Dim; ++d) { + ncells_per_pos[d].assign(g_ndomains_per_dim[d], 0); + load_per_pos[d].assign(g_ndomains_per_dim[d], 0.0); + } + for (unsigned int idx { 0 }; idx < g_ndomains; ++idx) { + const auto& off = g_domain_offsets[idx]; + const auto ncells = g_subdomains[idx].mesh.n_active(); + for (auto d { 0u }; d < M::Dim; ++d) { + ncells_per_pos[d][off[d]] = ncells[d]; + load_per_pos[d][off[d]] += static_cast(npart_per_dom[idx]); + } + } + + /* --- 3. Diffusion-style boundary shifts per balanced dim ------------- */ + std::vector> new_ncells_per_pos = ncells_per_pos; + const auto MIN_NCELLS = static_cast(2 * N_GHOSTS + 4); + bool any_shift = false; + + for (auto d { 0u }; d < M::Dim; ++d) { + if ((dim_mask & (1u << d)) == 0u) { + continue; + } + const auto N = g_ndomains_per_dim[d]; + if (N < 2) { + continue; + } + + double total_load = 0.0; + double max_load = 0.0; + double min_load = std::numeric_limits::infinity(); + for (auto p { 0u }; p < N; ++p) { + total_load += load_per_pos[d][p]; + max_load = std::max(max_load, load_per_pos[d][p]); + min_load = std::min(min_load, load_per_pos[d][p]); + } + if (total_load <= 0.0) { + continue; + } + const auto mean = total_load / static_cast(N); + if (((max_load - min_load) / mean) < + static_cast(tolerance)) { + continue; + } + + // bnd_shift[k] is the number of cells transferred from position k-1 to + // position k by moving the (interior) boundary k. bnd_shift[0] and + // bnd_shift[N] are exterior boundaries that are pinned at 0. + std::vector bnd_shift(N + 1, 0); + const int cap = static_cast(max_shift_cells); + for (auto k { 1u }; k < N; ++k) { + const auto l_load = load_per_pos[d][k - 1]; + const auto r_load = load_per_pos[d][k]; + const auto l_density = (ncells_per_pos[d][k - 1] > 0) + ? l_load / static_cast( + ncells_per_pos[d][k - 1]) + : 0.0; + const auto r_density = (ncells_per_pos[d][k] > 0) + ? r_load / static_cast( + ncells_per_pos[d][k]) + : 0.0; + const auto avg_density = std::max(0.5 * (l_density + r_density), 1.0); + // Move boundary towards the lighter side. Halve the gradient so that + // a single sweep does roughly one diffusion step. + int shift = static_cast( + std::round(0.5 * (l_load - r_load) / avg_density)); + shift = std::clamp(shift, -cap, cap); + // Don't shrink either side below MIN_NCELLS. + const int max_pos = static_cast(ncells_per_pos[d][k - 1]) - + static_cast(MIN_NCELLS); + const int max_neg = static_cast(ncells_per_pos[d][k]) - + static_cast(MIN_NCELLS); + shift = std::clamp(shift, + -std::max(max_neg, 0), + std::max(max_pos, 0)); + bnd_shift[k] = shift; + if (shift != 0) { + any_shift = true; + } + } + // new_ncells[p] = ncells[p] + bnd_shift[p] - bnd_shift[p+1] + for (auto p { 0u }; p < N; ++p) { + new_ncells_per_pos[d][p] = static_cast( + static_cast(ncells_per_pos[d][p]) + bnd_shift[p] - + bnd_shift[p + 1]); + } + } + + if (not any_shift) { + return; + } + + /* --- 4. Per-position prefix sums (offset and physical extent) -------- */ + // Face positions are queried from g_mesh.metric in the global code-coordinate + // system, so curvilinear stretches (log-r, eta-stretching) are honored. + std::vector> new_offset_per_pos(M::Dim); + std::vector>> extent_per_pos(M::Dim); + auto face_phys = [this](unsigned int d, real_t x_code) -> real_t { + if (d == 0u) { + return g_mesh.metric.template convert<1, Crd::Cd, Crd::Ph>(x_code); + } + if constexpr (M::Dim == Dim::_2D or M::Dim == Dim::_3D) { + if (d == 1u) { + return g_mesh.metric.template convert<2, Crd::Cd, Crd::Ph>(x_code); + } + } + if constexpr (M::Dim == Dim::_3D) { + if (d == 2u) { + return g_mesh.metric.template convert<3, Crd::Cd, Crd::Ph>(x_code); + } + } + raise::Error("Invalid dimension index in Rebalance face_phys", HERE); + return ZERO; + }; + for (auto d { 0u }; d < M::Dim; ++d) { + const auto N = g_ndomains_per_dim[d]; + new_offset_per_pos[d].assign(N, 0); + extent_per_pos[d].resize(N); + ncells_t running { 0 }; + for (auto p { 0u }; p < N; ++p) { + new_offset_per_pos[d][p] = running; + const auto x_lo = face_phys(d, static_cast(running)); + running += new_ncells_per_pos[d][p]; + const auto x_hi = face_phys(d, static_cast(running)); + extent_per_pos[d][p] = { x_lo, x_hi }; + } + } + + /* --- 5. Save local em (and em0/cur0 for GRPIC) to host --------------- */ + auto& local_dom = g_subdomains[local_idx]; + const auto old_offset_ncells = local_dom.offset_ncells(); + + auto em_old_h = Kokkos::create_mirror_view(local_dom.fields.em); + Kokkos::deep_copy(em_old_h, local_dom.fields.em); + auto em0_old_h = decltype(Kokkos::create_mirror_view(local_dom.fields.em0)) {}; + auto cur0_old_h = decltype(Kokkos::create_mirror_view(local_dom.fields.cur0)) {}; + if constexpr (S == SimEngine::GRPIC) { + em0_old_h = Kokkos::create_mirror_view(local_dom.fields.em0); + cur0_old_h = Kokkos::create_mirror_view(local_dom.fields.cur0); + Kokkos::deep_copy(em0_old_h, local_dom.fields.em0); + Kokkos::deep_copy(cur0_old_h, local_dom.fields.cur0); + } + + /* --- 6. Update bookkeeping for every g_subdomain --------------------- */ + std::vector new_local_ncells(M::Dim); + std::vector new_local_offset(M::Dim); + for (unsigned int idx { 0 }; idx < g_ndomains; ++idx) { + auto& sub = g_subdomains[idx]; + const auto& off_ndoms = g_domain_offsets[idx]; + std::vector ncells_d(M::Dim); + std::vector offset_d(M::Dim); + boundaries_t ext_d; + for (auto d { 0u }; d < M::Dim; ++d) { + ncells_d[d] = new_ncells_per_pos[d][off_ndoms[d]]; + offset_d[d] = new_offset_per_pos[d][off_ndoms[d]]; + ext_d.push_back(extent_per_pos[d][off_ndoms[d]]); + } + sub.mesh.set_resolution_and_extent(ncells_d, ext_d); + sub.set_offset_ncells(offset_d); + if (idx == local_idx) { + new_local_ncells = ncells_d; + new_local_offset = offset_d; + } + } + // Boundary conditions and neighbor topology are unchanged. + + /* --- 7. Reallocate local fields, copy from saved buffer with shift --- */ + // delta[d] = new_offset[d] - old_offset[d] (in the with-ghost field + // coordinate system the same delta applies). + std::vector delta(M::Dim); + for (auto d { 0u }; d < M::Dim; ++d) { + delta[d] = static_cast(new_local_offset[d]) - + static_cast(old_offset_ncells[d]); + } + + local_dom.fields = Fields { new_local_ncells }; + CopyShifted(em_old_h, local_dom.fields.em, delta); + if constexpr (S == SimEngine::GRPIC) { + CopyShifted(em0_old_h, local_dom.fields.em0, delta); + CopyShifted(cur0_old_h, local_dom.fields.cur0, delta); + } + + /* --- 8. Refill ghost zones from neighbors ---------------------------- */ + CommunicateFields(local_dom, Comm::E | Comm::B); + + /* --- 9. Shift particle indices, retag, and migrate ------------------- */ + // Particle indices i_d are in active-cell coordinates: i_d in [0, n_active) + // for an in-domain particle. After the offset moves by delta, particles + // get i_d_new = i_d_old - delta[d]. Particles whose new index falls + // outside [0, n_active) are tagged for the appropriate neighbor and + // sent by CommunicateParticles(). + for (auto& sp : local_dom.species) { + if (sp.npart() == 0) { + continue; + } + auto i1 = sp.i1, i2 = sp.i2, i3 = sp.i3; + auto i1p = sp.i1_prev, i2p = sp.i2_prev, i3p = sp.i3_prev; + auto tag = sp.tag; + const int dx1 = -delta[0]; + int dx2 = 0; + int dx3 = 0; + int new_n1 = static_cast(new_local_ncells[0]); + int new_n2 = 1; + int new_n3 = 1; + if constexpr (M::Dim == Dim::_2D or M::Dim == Dim::_3D) { + dx2 = -delta[1]; + new_n2 = static_cast(new_local_ncells[1]); + } + if constexpr (M::Dim == Dim::_3D) { + dx3 = -delta[2]; + new_n3 = static_cast(new_local_ncells[2]); + } + Kokkos::parallel_for( + "RebalanceShiftPrtls", + sp.rangeActiveParticles(), + Lambda(prtlidx_t p) { + if (tag(p) != ParticleTag::alive) { + return; + } + if constexpr (M::Dim == Dim::_1D or M::Dim == Dim::_2D or + M::Dim == Dim::_3D) { + i1(p) += dx1; + i1p(p) += dx1; + } + if constexpr (M::Dim == Dim::_2D or M::Dim == Dim::_3D) { + i2(p) += dx2; + i2p(p) += dx2; + } + if constexpr (M::Dim == Dim::_3D) { + i3(p) += dx3; + i3p(p) += dx3; + } + if constexpr (M::Dim == Dim::_1D) { + tag(p) = mpi::SendTag(tag(p), i1(p) < 0, i1(p) >= new_n1); + } else if constexpr (M::Dim == Dim::_2D) { + tag(p) = mpi::SendTag(tag(p), + i1(p) < 0, + i1(p) >= new_n1, + i2(p) < 0, + i2(p) >= new_n2); + } else if constexpr (M::Dim == Dim::_3D) { + tag(p) = mpi::SendTag(tag(p), + i1(p) < 0, + i1(p) >= new_n1, + i2(p) < 0, + i2(p) >= new_n2, + i3(p) < 0, + i3(p) >= new_n3); + } + }); + sp.set_unsorted(); + } + + CommunicateParticles(local_dom); + logger::Checkpoint("Rebalance: domains shifted, fields and particles redistributed", + HERE); +#endif // MPI_ENABLED + } + + // NOLINTBEGIN(bugprone-macro-parentheses) +#define METADOMAIN_REBAL(S, M, D) \ + template void Metadomain>::Rebalance(unsigned int, real_t, ncells_t); + + NTT_FOREACH_SPECIALIZATION(METADOMAIN_REBAL) +#undef METADOMAIN_REBAL + // NOLINTEND(bugprone-macro-parentheses) + +} // namespace ntt diff --git a/src/framework/parameters/parameters.cpp b/src/framework/parameters/parameters.cpp index 7372da510..8b525d49a 100644 --- a/src/framework/parameters/parameters.cpp +++ b/src/framework/parameters/parameters.cpp @@ -239,6 +239,52 @@ namespace ntt { alg_params.read(get("scales.dx0"), alg_extra_flags, toml_data); alg_params.setParams(alg_extra_flags, this); + /* [simulation.domain.load_balance] ------------------------------------- */ + set("simulation.domain.load_balance.enable", + toml::find_or(toml_data, "simulation", "domain", "load_balance", "enable", false)); + set("simulation.domain.load_balance.interval", + toml::find_or(toml_data, + "simulation", + "domain", + "load_balance", + "interval", + 0u)); + set("simulation.domain.load_balance.tolerance", + toml::find_or(toml_data, + "simulation", + "domain", + "load_balance", + "tolerance", + static_cast(0.1))); + set("simulation.domain.load_balance.max_shift", + toml::find_or(toml_data, + "simulation", + "domain", + "load_balance", + "max_shift", + static_cast(N_GHOSTS))); + { + // dimensions: list of 1/2/3 mapped to a bitmask + const auto dim_ints = toml::find_or>( + toml_data, + "simulation", + "domain", + "load_balance", + "dimensions", + std::vector { 1 }); + unsigned int mask = 0u; + for (const auto& d : dim_ints) { + if (d == 1 or d == 2 or d == 3) { + mask |= 1u << (d - 1); + } else { + raise::Error( + "simulation.domain.load_balance.dimensions: unknown dim, expected 1/2/3", + HERE); + } + } + set("simulation.domain.load_balance.dim_mask", mask); + } + /* extra physics ------------------------------------------------------ */ params::Extra extra_params {}; const std::map extra_extra_flags = { diff --git a/src/output/utils/writers.cpp b/src/output/utils/writers.cpp index a02f298e4..d89cd8cab 100644 --- a/src/output/utils/writers.cpp +++ b/src/output/utils/writers.cpp @@ -76,13 +76,18 @@ namespace out { } template - void WriteNDField(adios2::IO& io, - adios2::Engine& writer, - const std::string& name, - const ndfield_t& data) { + void WriteNDField(adios2::IO& io, + adios2::Engine& writer, + const std::string& name, + const ndfield_t& data, + const adios2::Box& range) { + auto var = io.InquireVariable(name); + if (not range.first.empty()) { + var.SetSelection(range); + } auto data_h = Kokkos::create_mirror_view(data); Kokkos::deep_copy(data_h, data); - writer.Put(io.InquireVariable(name), data_h.data(), adios2::Mode::Sync); + writer.Put(var, data_h.data(), adios2::Mode::Sync); } // NOLINTBEGIN(bugprone-macro-parentheses) @@ -122,7 +127,8 @@ namespace out { template void WriteNDField(adios2::IO&, \ adios2::Engine&, \ const std::string&, \ - const ndfield_t&); + const ndfield_t&, \ + const adios2::Box&); NDFIELD_WRITERS(Dim::_1D, 3) NDFIELD_WRITERS(Dim::_1D, 6) NDFIELD_WRITERS(Dim::_2D, 3) diff --git a/src/output/utils/writers.h b/src/output/utils/writers.h index cbd089452..4b8909a19 100644 --- a/src/output/utils/writers.h +++ b/src/output/utils/writers.h @@ -78,7 +78,8 @@ namespace out { void WriteNDField(adios2::IO&, adios2::Engine&, const std::string&, - const ndfield_t&); + const ndfield_t&, + const adios2::Box& = {}); } // namespace out diff --git a/src/output/writer.cpp b/src/output/writer.cpp index a84745c5b..5b9ae89e2 100644 --- a/src/output/writer.cpp +++ b/src/output/writer.cpp @@ -73,6 +73,33 @@ namespace out { m_mode = mode; } + void Writer::setLocalLayout(const std::vector& loc_corner, + const std::vector& loc_shape) { + raise::ErrorIf(loc_corner.size() != m_flds_l_corner.size() or + loc_shape.size() != m_flds_l_shape.size(), + "setLocalLayout dim mismatch with the original layout", + HERE); + m_flds_l_corner = loc_corner; + m_flds_l_shape = loc_shape; + m_flds_l_corner_dwn.clear(); + m_flds_l_shape_dwn.clear(); + m_flds_l_first.clear(); + for (auto i { 0u }; i < m_flds_g_shape.size(); ++i) { + const double d = static_cast(m_dwn[i]); + const double l = static_cast(loc_corner[i]); + const double n = static_cast(loc_shape[i]); + const double f = math::ceil(l / d) * d - l; + m_flds_l_corner_dwn.push_back(static_cast(math::ceil(l / d))); + m_flds_l_first.push_back(static_cast(f)); + m_flds_l_shape_dwn.push_back(static_cast(math::ceil((n - f) / d))); + } + if constexpr (not std::is_same::array_layout, + Kokkos::LayoutRight>::value) { + std::reverse(m_flds_l_corner_dwn.begin(), m_flds_l_corner_dwn.end()); + std::reverse(m_flds_l_shape_dwn.begin(), m_flds_l_shape_dwn.end()); + } + } + void Writer::defineMeshLayout( const std::vector& glob_shape, const std::vector& loc_corner, @@ -111,25 +138,24 @@ namespace out { for (auto i { 0u }; i < m_flds_g_shape.size(); ++i) { // cell-centers + // ConstantDims is intentionally NOT set: per-rank slab shape can change + // when dynamic load balancing shifts domain boundaries between writes. m_io.DefineVariable("X" + std::to_string(i + 1), { m_flds_g_shape_dwn[i] }, { m_flds_l_corner_dwn[i] }, - { m_flds_l_shape_dwn[i] }, - adios2::ConstantDims); + { m_flds_l_shape_dwn[i] }); // cell-edges const auto is_last = (m_flds_l_corner[i] + m_flds_l_shape[i] == m_flds_g_shape[i]); m_io.DefineVariable("X" + std::to_string(i + 1) + "e", { m_flds_g_shape_dwn[i] + 1 }, { m_flds_l_corner_dwn[i] }, - { m_flds_l_shape_dwn[i] + (is_last ? 1 : 0) }, - adios2::ConstantDims); + { m_flds_l_shape_dwn[i] + (is_last ? 1 : 0) }); m_io.DefineVariable( "N" + std::to_string(i + 1) + "l", { static_cast(2 * domain_idx.second) }, { static_cast(2 * domain_idx.first) }, - { static_cast(2) }, - adios2::ConstantDims); + { static_cast(2) }); } if constexpr (std::is_same::array_layout, @@ -154,21 +180,21 @@ namespace out { m_flds_writers.emplace_back(S, fld); } for (const auto& fld : m_flds_writers) { + // ConstantDims is intentionally NOT set: per-rank slab shape can change + // when dynamic load balancing shifts domain boundaries between writes. if (fld.comp.empty()) { // scalar m_io.DefineVariable(fld.name(), m_flds_g_shape_dwn, m_flds_l_corner_dwn, - m_flds_l_shape_dwn, - adios2::ConstantDims); + m_flds_l_shape_dwn); } else { // vector or tensor for (auto i { 0u }; i < fld.comp.size(); ++i) { m_io.DefineVariable(fld.name(i), m_flds_g_shape_dwn, m_flds_l_corner_dwn, - m_flds_l_shape_dwn, - adios2::ConstantDims); + m_flds_l_shape_dwn); } } } @@ -198,9 +224,14 @@ namespace out { std::size_t comp, std::vector dwn, std::vector first_cell, - bool ghosts) { + bool ghosts, + const adios2::Dims& loc_corner_dwn, + const adios2::Dims& loc_shape_dwn) { // when dwn != 1 in any direction, it is assumed that ghosts == false - auto var = io.InquireVariable(varname); + auto var = io.InquireVariable(varname); + // Refresh the per-step (start, count) so the slab tracks any rebalance + // that happened since the variable was declared. + var.SetSelection(adios2::Box(loc_corner_dwn, loc_shape_dwn)); const auto gh_zones = ghosts ? 0 : N_GHOSTS; ndarray_t output_field {}; @@ -332,7 +363,9 @@ namespace out { addresses[i], m_dwn, m_flds_l_first, - m_flds_ghosts); + m_flds_ghosts, + m_flds_l_corner_dwn, + m_flds_l_shape_dwn); } } @@ -410,8 +443,28 @@ namespace out { const array_t& xc, const array_t& xe, const std::vector& loc_off_sz) { + // Per-step (start, count) for the per-rank slab; tracks the (possibly + // rebalanced) layout cached by setLocalLayout(). auto varc = m_io.InquireVariable("X" + std::to_string(dim + 1)); auto vare = m_io.InquireVariable("X" + std::to_string(dim + 1) + "e"); + // m_flds_l_corner_dwn / m_flds_l_shape_dwn are reversed for non-LayoutRight + // (see defineMeshLayout / setLocalLayout); m_flds_l_corner / m_flds_l_shape + // / m_flds_g_shape are not. Map the dim-order index to the dwn-array index. + constexpr bool layout_right = std::is_same< + typename ndfield_t::array_layout, + Kokkos::LayoutRight>::value; + const auto i_dwn = layout_right + ? static_cast(dim) + : (m_flds_g_shape.size() - 1u - + static_cast(dim)); + const auto is_last = (m_flds_l_corner[dim] + m_flds_l_shape[dim] == + m_flds_g_shape[dim]); + varc.SetSelection(adios2::Box( + { m_flds_l_corner_dwn[i_dwn] }, + { m_flds_l_shape_dwn[i_dwn] })); + vare.SetSelection(adios2::Box( + { m_flds_l_corner_dwn[i_dwn] }, + { m_flds_l_shape_dwn[i_dwn] + (is_last ? 1ul : 0ul) })); auto xc_h = Kokkos::create_mirror_view(xc); auto xe_h = Kokkos::create_mirror_view(xe); Kokkos::deep_copy(xc_h, xc); @@ -506,7 +559,9 @@ namespace out { std::size_t, \ std::vector, \ std::vector, \ - bool); + bool, \ + const adios2::Dims&, \ + const adios2::Dims&); WRITE_FIELD(Dim::_1D, 3) WRITE_FIELD(Dim::_1D, 6) WRITE_FIELD(Dim::_2D, 3) diff --git a/src/output/writer.h b/src/output/writer.h index 055cb215e..7231c738b 100644 --- a/src/output/writer.h +++ b/src/output/writer.h @@ -109,6 +109,17 @@ namespace out { bool, Coord); + /** + * @brief Refresh the cached per-rank slab corner and shape. + * + * Must be called before any per-step write whenever the local domain has + * been resized (e.g. by `Metadomain::Rebalance`); the cached values are + * passed to `var.SetSelection()` inside `writeMesh` and `WriteField`. + * The global shape is preserved (load balancing conserves total cells). + */ + void setLocalLayout(const std::vector& loc_corner, + const std::vector& loc_shape); + void defineFieldOutputs(const SimEngine&, const std::vector&); void defineSpectraOutputs(const std::vector&);