From ad2f3db0fd8c14d14cd6f0d909f5ee623134e0f8 Mon Sep 17 00:00:00 2001 From: LudwigBoess Date: Fri, 7 Aug 2026 15:40:43 -0500 Subject: [PATCH 1/2] bugfix for loadbalancing when compiled with nvcc --- src/framework/domain/metadomain.h | 3 +- src/framework/domain/metadomain_loadbal.cpp | 8 +- src/kernels/pushers/sr_policies.h | 87 ++++++++++++++------- 3 files changed, 64 insertions(+), 34 deletions(-) diff --git a/src/framework/domain/metadomain.h b/src/framework/domain/metadomain.h index f090dade1..ed35a7140 100644 --- a/src/framework/domain/metadomain.h +++ b/src/framework/domain/metadomain.h @@ -146,8 +146,7 @@ namespace ntt { */ void Rebalance(unsigned int dim_mask, real_t tolerance, - ncells_t max_shift_cells) - requires(MetricClass); + ncells_t max_shift_cells); /* output-related ------------------------------------------------------- */ #if defined(OUTPUT_ENABLED) diff --git a/src/framework/domain/metadomain_loadbal.cpp b/src/framework/domain/metadomain_loadbal.cpp index 0768acad6..fdf167b67 100644 --- a/src/framework/domain/metadomain_loadbal.cpp +++ b/src/framework/domain/metadomain_loadbal.cpp @@ -110,9 +110,7 @@ namespace ntt { template void Metadomain::Rebalance(unsigned int dim_mask, real_t tolerance, - ncells_t max_shift_cells) - requires(MetricClass) - { + ncells_t max_shift_cells) { #if !defined(MPI_ENABLED) (void)dim_mask; (void)tolerance; @@ -377,6 +375,10 @@ namespace ntt { if (tag(p) != ParticleTag::alive) { return; } + // nvcc: force capture of all vars before any constexpr-if branch + (void)i1; (void)i1p; (void)i2; (void)i2p; (void)i3; (void)i3p; + (void)dx1; (void)dx2; (void)dx3; + (void)new_n1; (void)new_n2; (void)new_n3; if constexpr (M::Dim == Dim::_1D or M::Dim == Dim::_2D or M::Dim == Dim::_3D) { i1(p) += dx1; diff --git a/src/kernels/pushers/sr_policies.h b/src/kernels/pushers/sr_policies.h index b247b27b7..195fa380a 100644 --- a/src/kernels/pushers/sr_policies.h +++ b/src/kernels/pushers/sr_policies.h @@ -102,6 +102,57 @@ namespace kernel::sr { } } + // nvcc/nvc++ workaround: if constexpr inside generic lambdas does not reliably + // discard dead branches on some NVC++ versions. Factor the three optional-pgen + // dispatches into regular template functions so the discard happens at template + // instantiation, not inside a lambda operator(). + + template + void sr_dispatch_custom_emission(const PGen& pgen, + simtime_t time, + spidx_t sp, + DOM& dom, + Next&& next) { + if constexpr (::traits::pgen::HasEmissionPolicy) { + next(pgen.EmissionPolicy(time, sp, dom)); + } else { + raise::Error("Custom emission policy flag is set but problem " + "generator does not define an emission policy", + HERE); + } + } + + template + void sr_dispatch_cpu_policy(const PGen& pgen, + simtime_t time, + spidx_t sp, + DOM& dom, + Next&& next) { + if constexpr (::traits::pgen::HasCustomPrtlUpdate) { + next(pgen.CustomParticleUpdate(time, sp, dom)); + } else { + next(::traits::custom_prtl_update::NoPolicy_t {}); + } + } + + template + void sr_dispatch_extfields(const PGen& pgen, + simtime_t time, + spidx_t sp, + DOM& dom, + Next&& next) { + if constexpr (::traits::pgen::HasExternalFields) { + const auto [apply_extfields, external_fields] = pgen.ExternalFields(time, sp, dom); + if (apply_extfields) { + next(external_fields); + } else { + next(::traits::extfields::NoPolicy_t {}); + } + } else { + next(::traits::extfields::NoPolicy_t {}); + } + } + template void MakePusherPolicy(const PGen& pgen, DOM& domain, @@ -125,15 +176,11 @@ namespace kernel::sr { pusher_ctx)); break; case ntt::EmissionType::CUSTOM: - if constexpr (::traits::pgen::HasEmissionPolicy) { - next(pgen.EmissionPolicy(pusher_ctx.time, - pusher_ctx.species_index, - domain)); - } else { - raise::Error("Custom emission policy flag is set but problem " - "generator does not define an emission policy", - HERE); - } + sr_dispatch_custom_emission(pgen, + pusher_ctx.time, + pusher_ctx.species_index, + domain, + next); break; case ntt::EmissionType::NONE: default: @@ -143,29 +190,11 @@ namespace kernel::sr { }; auto with_custom_prtl_upd = [&](auto next) { - if constexpr (::traits::pgen::HasCustomPrtlUpdate) { - next(pgen.CustomParticleUpdate(pusher_ctx.time, - pusher_ctx.species_index, - domain)); - } else { - next(::traits::custom_prtl_update::NoPolicy_t {}); - } + sr_dispatch_cpu_policy(pgen, pusher_ctx.time, pusher_ctx.species_index, domain, next); }; auto with_ext_fields = [&](auto next) { - if constexpr (::traits::pgen::HasExternalFields) { - const auto [apply_extfields, external_fields] = pgen.ExternalFields( - pusher_ctx.time, - pusher_ctx.species_index, - domain); - if (apply_extfields) { - next(external_fields); - } else { - next(::traits::extfields::NoPolicy_t {}); - } - } else { - next(::traits::extfields::NoPolicy_t {}); - } + sr_dispatch_extfields(pgen, pusher_ctx.time, pusher_ctx.species_index, domain, next); }; with_emission([&](auto ep) { From 765784eb2410fbf768a853e396f6e0ce4efc6f8d Mon Sep 17 00:00:00 2001 From: LudwigBoess Date: Tue, 11 Aug 2026 21:23:00 +0000 Subject: [PATCH 2/2] replace the Rebalance particle-shift lambda with a ShiftPrtlIndices_kernel functor --- src/framework/domain/metadomain_loadbal.cpp | 135 +++++++++++++------- 1 file changed, 90 insertions(+), 45 deletions(-) diff --git a/src/framework/domain/metadomain_loadbal.cpp b/src/framework/domain/metadomain_loadbal.cpp index fdf167b67..a7bf97873 100644 --- a/src/framework/domain/metadomain_loadbal.cpp +++ b/src/framework/domain/metadomain_loadbal.cpp @@ -105,6 +105,81 @@ namespace ntt { } Kokkos::deep_copy(dst_dev, dst_h); } + + // Shift particle cell indices by dx_d after the local active-cell offset + // moved, and re-tag every particle whose new index falls outside the new + // active range [0, new_n_d) for the corresponding neighbor. + template + class ShiftPrtlIndices_kernel { + array_t i1, i1_prev, i2, i2_prev, i3, i3_prev; + array_t tag; + + const int dx1, dx2, dx3; + const int new_n1, new_n2, new_n3; + + public: + ShiftPrtlIndices_kernel(array_t& i1, + array_t& i1_prev, + array_t& i2, + array_t& i2_prev, + array_t& i3, + array_t& i3_prev, + array_t& tag, + int dx1, + int dx2, + int dx3, + int new_n1, + int new_n2, + int new_n3) + : i1 { i1 } + , i1_prev { i1_prev } + , i2 { i2 } + , i2_prev { i2_prev } + , i3 { i3 } + , i3_prev { i3_prev } + , tag { tag } + , dx1 { dx1 } + , dx2 { dx2 } + , dx3 { dx3 } + , new_n1 { new_n1 } + , new_n2 { new_n2 } + , new_n3 { new_n3 } {} + + Inline void operator()(prtlidx_t p) const { + if (tag(p) != ParticleTag::alive) { + return; + } + if constexpr (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { + i1(p) += dx1; + i1_prev(p) += dx1; + } + if constexpr (D == Dim::_2D or D == Dim::_3D) { + i2(p) += dx2; + i2_prev(p) += dx2; + } + if constexpr (D == Dim::_3D) { + i3(p) += dx3; + i3_prev(p) += dx3; + } + if constexpr (D == Dim::_1D) { + tag(p) = mpi::SendTag(tag(p), i1(p) < 0, i1(p) >= new_n1); + } else if constexpr (D == 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 (D == 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); + } + } + }; #endif // MPI_ENABLED template @@ -351,9 +426,6 @@ namespace ntt { 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; @@ -368,48 +440,21 @@ namespace ntt { 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; - } - // nvcc: force capture of all vars before any constexpr-if branch - (void)i1; (void)i1p; (void)i2; (void)i2p; (void)i3; (void)i3p; - (void)dx1; (void)dx2; (void)dx3; - (void)new_n1; (void)new_n2; (void)new_n3; - 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); - } - }); + Kokkos::parallel_for("RebalanceShiftPrtls", + sp.rangeActiveParticles(), + ShiftPrtlIndices_kernel { sp.i1, + sp.i1_prev, + sp.i2, + sp.i2_prev, + sp.i3, + sp.i3_prev, + sp.tag, + dx1, + dx2, + dx3, + new_n1, + new_n2, + new_n3 }); sp.set_unsorted(); }