From 2ae70bea46d486a88b3830b89b29a08385702eac Mon Sep 17 00:00:00 2001 From: Alexander Liemen <47395086+aliemen@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:13:06 +0200 Subject: [PATCH 1/3] Change ORB to support arbitrary decomposition axes and improve domain validation. Conserve legacy behavior. --- .../OrthogonalRecursiveBisection.h | 44 ++++++- .../OrthogonalRecursiveBisection.hpp | 114 +++++++++++++++--- 2 files changed, 141 insertions(+), 17 deletions(-) diff --git a/src/Decomposition/OrthogonalRecursiveBisection.h b/src/Decomposition/OrthogonalRecursiveBisection.h index cb8767fe0..305c27324 100644 --- a/src/Decomposition/OrthogonalRecursiveBisection.h +++ b/src/Decomposition/OrthogonalRecursiveBisection.h @@ -10,6 +10,11 @@ #ifndef IPPL_ORTHOGONAL_RECURSIVE_BISECTION_H #define IPPL_ORTHOGONAL_RECURSIVE_BISECTION_H +#include +#include +#include +#include + #include "FieldLayout/FieldLayout.h" #include "Index/Index.h" #include "Index/NDIndex.h" @@ -44,7 +49,8 @@ namespace ippl { /*! * Performs scatter operation of particle positions in field (weights) and - * repartitions FieldLayout's global domain + * repartitions FieldLayout's global domain. This overload preserves the + * legacy ORB behavior by allowing cuts along all axes. * @tparam Attrib the particle attribute type (memory space must be accessible to field * memory) * @param R Weights to scatter @@ -55,12 +61,48 @@ namespace ippl { bool binaryRepartition(const Attrib& R, FieldLayout& fl, const bool& isFirstRepartition); + /*! + * Performs scatter operation of particle positions in field (weights) and + * repartitions FieldLayout's global domain using only the enabled axes. + * The FieldLayout and the ORB weight field are updated only after the proposed + * domains pass validation. + * @tparam Attrib the particle attribute type (memory space must be accessible to field + * memory) + * @param R Weights to scatter + * @param fl FieldLayout + * @param isFirstRepartition boolean which tells whether to scatter or not + * @param allowedAxes true for axes ORB is allowed to cut + */ + template + bool binaryRepartition(const Attrib& R, FieldLayout& fl, + const bool& isFirstRepartition, + const std::array& allowedAxes); + /*! * Find cutting axis as the longest axis of the field layout. * @param dom Domain to reduce */ int findCutAxis(NDIndex& dom); + /*! + * Find cutting axis as the longest enabled axis of the field layout. + * @param dom Domain to reduce + * @param allowedAxes true for axes ORB is allowed to cut + */ + int findCutAxis(const NDIndex& dom, const std::array& allowedAxes); + + /*! + * Check whether two domains overlap. + */ + bool domainsOverlap(const NDIndex& lhs, const NDIndex& rhs) const; + + /*! + * Check whether proposed domains tile the global domain without cutting serial axes. + */ + bool domainsTileAllowedDecomposition(const std::vector>& domains, + const NDIndex& globalDomain, + const std::array& allowedAxes) const; + /*! * Performs reduction on local field in all dimension except that determined * by cutAxis, stores result in res diff --git a/src/Decomposition/OrthogonalRecursiveBisection.hpp b/src/Decomposition/OrthogonalRecursiveBisection.hpp index 0ada6cfca..a30a94acb 100644 --- a/src/Decomposition/OrthogonalRecursiveBisection.hpp +++ b/src/Decomposition/OrthogonalRecursiveBisection.hpp @@ -13,6 +13,16 @@ namespace ippl { template bool OrthogonalRecursiveBisection::binaryRepartition( const Attrib& R, FieldLayout& fl, const bool& isFirstRepartition) { + std::array allowedAxes; + allowedAxes.fill(true); + return binaryRepartition(R, fl, isFirstRepartition, allowedAxes); + } + + template + template + bool OrthogonalRecursiveBisection::binaryRepartition( + const Attrib& R, FieldLayout& fl, const bool& isFirstRepartition, + const std::array& allowedAxes) { // Timings static IpplTimings::TimerRef tbasicOp = IpplTimings::getTimer("basicOperations"); static IpplTimings::TimerRef tperpReduction = IpplTimings::getTimer("perpReduction"); @@ -53,13 +63,22 @@ namespace ippl { while (maxprocs > 1) { // Find cut axis IpplTimings::startTimer(tbasicOp); - int cutAxis = findCutAxis(domains[it]); + int cutAxis = findCutAxis(domains[it], allowedAxes); + if (cutAxis < 0) { + IpplTimings::stopTimer(tbasicOp); + return false; + } + int cutAxisLength = static_cast(domains[it][cutAxis].length()); + if (cutAxisLength < 4) { + IpplTimings::stopTimer(tbasicOp); + return false; + } IpplTimings::stopTimer(tbasicOp); // Reserve space IpplTimings::startTimer(tperpReduction); - reduced.resize(domains[it][cutAxis].length()); - reducedRank.resize(domains[it][cutAxis].length()); + reduced.resize(cutAxisLength); + reducedRank.resize(cutAxisLength); std::fill(reducedRank.begin(), reducedRank.end(), 0.0); std::fill(reduced.begin(), reduced.end(), 0.0); @@ -77,7 +96,7 @@ namespace ippl { IpplTimings::startTimer(tbasicOp); // Initialize median to some value (1 is lower bound value) int median = 1; - median = findMedian(reduced); + median = std::clamp(findMedian(reduced), 1, cutAxisLength - 3); IpplTimings::stopTimer(tbasicOp); // Cut domains and procs @@ -101,14 +120,15 @@ namespace ippl { IpplTimings::stopTimer(tperpReduction); } - // Check that no plane was obtained in the repartition IpplTimings::startTimer(tbasicOp); - for (const auto& domain : domains) { - for (const auto& axis : domain) { - if (axis.length() == 1) { - return false; - } - } + const NDIndex& globalDomain = fl.getDomain(); + if (domains.size() != static_cast(nprocs)) { + IpplTimings::stopTimer(tbasicOp); + return false; + } + if (!domainsTileAllowedDecomposition(domains, globalDomain, allowedAxes)) { + IpplTimings::stopTimer(tbasicOp); + return false; } // Update FieldLayout with new indices @@ -123,11 +143,73 @@ namespace ippl { template int OrthogonalRecursiveBisection::findCutAxis(NDIndex& dom) { - // Find longest domain size - return std::distance(dom.begin(), std::max_element(dom.begin(), dom.end(), - [&](const Index& a, const Index& b) { - return a.length() < b.length(); - })); + std::array allowedAxes; + allowedAxes.fill(true); + return findCutAxis(dom, allowedAxes); + } + + template + int OrthogonalRecursiveBisection::findCutAxis( + const NDIndex& dom, const std::array& allowedAxes) { + int cutAxis = -1; + for (unsigned d = 0; d < Dim; ++d) { + if (!allowedAxes[d]) { + continue; + } + if (cutAxis < 0 || dom[d].length() > dom[cutAxis].length()) { + cutAxis = d; + } + } + return cutAxis; + } + + template + bool OrthogonalRecursiveBisection::domainsOverlap( + const NDIndex& lhs, const NDIndex& rhs) const { + for (unsigned d = 0; d < Dim; ++d) { + if (lhs[d].last() < rhs[d].first() || rhs[d].last() < lhs[d].first()) { + return false; + } + } + return true; + } + + template + bool OrthogonalRecursiveBisection::domainsTileAllowedDecomposition( + const std::vector>& domains, const NDIndex& globalDomain, + const std::array& allowedAxes) const { + size_t globalCells = 1; + for (unsigned d = 0; d < Dim; ++d) { + globalCells *= globalDomain[d].length(); + } + + size_t localCells = 0; + for (size_t i = 0; i < domains.size(); ++i) { + const auto& domain = domains[i]; + size_t domainCells = 1; + for (unsigned d = 0; d < Dim; ++d) { + if (domain[d].length() < 2) { + return false; + } + if (domain[d].first() < globalDomain[d].first() + || domain[d].last() > globalDomain[d].last()) { + return false; + } + if (!allowedAxes[d] && domain[d] != globalDomain[d]) { + return false; + } + domainCells *= domain[d].length(); + } + + for (size_t j = 0; j < i; ++j) { + if (domainsOverlap(domain, domains[j])) { + return false; + } + } + localCells += domainCells; + } + + return localCells == globalCells; } template From 9e8162f2e7443483825a6fbdcf4a22c5f776d816 Mon Sep 17 00:00:00 2001 From: Alexander Liemen <47395086+aliemen@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:25:17 +0200 Subject: [PATCH 2/3] Add ORB unit test for Dim>1 that tests arbitrary cut axis --- unit_tests/PIC/ORB.cpp | 74 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/unit_tests/PIC/ORB.cpp b/unit_tests/PIC/ORB.cpp index 7debb8cd2..49e553d03 100644 --- a/unit_tests/PIC/ORB.cpp +++ b/unit_tests/PIC/ORB.cpp @@ -167,6 +167,80 @@ TYPED_TEST(ORBTest, Charge) { } +TYPED_TEST(ORBTest, AllowedAxesPreserveDisabledDimensions) { + constexpr unsigned Dim = TestFixture::dim; + + if constexpr (Dim < 2) { + GTEST_SKIP() << "Need at least one disabled dimension."; + } else { + if (ippl::Comm->size() < 2) { + GTEST_SKIP() << "Allowed-axis repartition needs more than one rank."; + } + + auto& bunch = this->bunch; + auto& layout = this->layout; + auto& orb = this->orb; + + std::array allowedAxes; + allowedAxes.fill(false); + allowedAxes[0] = true; + + bool fromAnalyticDensity = true; + ASSERT_TRUE(orb.binaryRepartition(bunch->R, layout, fromAnalyticDensity, allowedAxes)); + + const auto& globalDomain = layout.getDomain(); + auto localDomains = layout.getHostLocalDomains(); + + bool cutAllowedAxis = false; + for (size_t rank = 0; rank < localDomains.size(); ++rank) { + const auto& localDomain = localDomains(rank); + if (localDomain[0].first() != globalDomain[0].first() + || localDomain[0].last() != globalDomain[0].last()) { + cutAllowedAxis = true; + } + + for (unsigned d = 1; d < Dim; ++d) { + EXPECT_EQ(localDomain[d].first(), globalDomain[d].first()); + EXPECT_EQ(localDomain[d].last(), globalDomain[d].last()); + EXPECT_EQ(localDomain[d].stride(), globalDomain[d].stride()); + } + } + + EXPECT_TRUE(cutAllowedAxis); + } +} + +TYPED_TEST(ORBTest, RejectsEmptyAllowedAxisMask) { + constexpr unsigned Dim = TestFixture::dim; + + if (ippl::Comm->size() < 2) { + GTEST_SKIP() << "Empty allowed-axis mask only fails when ORB needs to cut."; + } + + auto& bunch = this->bunch; + auto& layout = this->layout; + auto& orb = this->orb; + + auto originalDomains = layout.getHostLocalDomains(); + + std::array allowedAxes; + allowedAxes.fill(false); + + bool fromAnalyticDensity = true; + EXPECT_FALSE(orb.binaryRepartition(bunch->R, layout, fromAnalyticDensity, allowedAxes)); + + auto currentDomains = layout.getHostLocalDomains(); + ASSERT_EQ(originalDomains.size(), currentDomains.size()); + + for (size_t rank = 0; rank < currentDomains.size(); ++rank) { + for (unsigned d = 0; d < Dim; ++d) { + EXPECT_EQ(currentDomains(rank)[d].first(), originalDomains(rank)[d].first()); + EXPECT_EQ(currentDomains(rank)[d].last(), originalDomains(rank)[d].last()); + EXPECT_EQ(currentDomains(rank)[d].stride(), originalDomains(rank)[d].stride()); + } + } +} + int main(int argc, char* argv[]) { int success = 1; ippl::initialize(argc, argv); From f8125162dea1834ee389e2df2a503debbe5d7a9a Mon Sep 17 00:00:00 2001 From: Alexander Liemen <47395086+aliemen@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:26:05 +0200 Subject: [PATCH 3/3] Collective clang-format --- src/Decomposition/OrthogonalRecursiveBisection.hpp | 4 ++-- unit_tests/PIC/ORB.cpp | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/Decomposition/OrthogonalRecursiveBisection.hpp b/src/Decomposition/OrthogonalRecursiveBisection.hpp index a30a94acb..f1455ee34 100644 --- a/src/Decomposition/OrthogonalRecursiveBisection.hpp +++ b/src/Decomposition/OrthogonalRecursiveBisection.hpp @@ -164,8 +164,8 @@ namespace ippl { } template - bool OrthogonalRecursiveBisection::domainsOverlap( - const NDIndex& lhs, const NDIndex& rhs) const { + bool OrthogonalRecursiveBisection::domainsOverlap(const NDIndex& lhs, + const NDIndex& rhs) const { for (unsigned d = 0; d < Dim; ++d) { if (lhs[d].last() < rhs[d].first() || rhs[d].last() < lhs[d].first()) { return false; diff --git a/unit_tests/PIC/ORB.cpp b/unit_tests/PIC/ORB.cpp index 49e553d03..92c7eebf7 100644 --- a/unit_tests/PIC/ORB.cpp +++ b/unit_tests/PIC/ORB.cpp @@ -151,7 +151,7 @@ TYPED_TEST(ORBTest, Charge) { typename TestFixture::value_type tol = tolerance; double charge = 0.5; - bunch->Q = charge/this->nParticles; + bunch->Q = charge / this->nParticles; bunch->update(); @@ -164,7 +164,6 @@ TYPED_TEST(ORBTest, Charge) { double totalCharge = field->sum(); ASSERT_NEAR((charge - totalCharge), 0., tol); - } TYPED_TEST(ORBTest, AllowedAxesPreserveDisabledDimensions) {