Skip to content

Change ORB to support arbitrary decomposition axes #560

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

Open
aliemen wants to merge 5 commits into
base: master
Choose a base branch
from
Open

Change ORB to support arbitrary decomposition axes #560

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2ae70be
Change ORB to support arbitrary decomposition axes and improve domain…
aliemen Jul 1, 2026
9e8162f
Add ORB unit test for Dim>1 that tests arbitrary cut axis
aliemen Jul 1, 2026
f812516
Collective clang-format
aliemen Jul 1, 2026
eac1088
Merge branch 'master' into 559-orb-does-not-allow-arbitrary-decomposi…
aliemen Jul 10, 2026
8869c39
Merge branch 'master' into 559-orb-does-not-allow-arbitrary-decomposi…
aliemen Jul 11, 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
44 changes: 43 additions & 1 deletion src/Decomposition/OrthogonalRecursiveBisection.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,11 @@
#ifndef IPPL_ORTHOGONAL_RECURSIVE_BISECTION_H
#define IPPL_ORTHOGONAL_RECURSIVE_BISECTION_H

#include <algorithm>
#include <array>
#include <numeric>
#include <vector>

#include "FieldLayout/FieldLayout.h"
#include "Index/Index.h"
#include "Index/NDIndex.h"
Expand Down Expand Up @@ -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
Expand All @@ -55,12 +61,48 @@ namespace ippl {
bool binaryRepartition(const Attrib& R, FieldLayout<Dim>& 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 <typename Attrib>
bool binaryRepartition(const Attrib& R, FieldLayout<Dim>& fl,
const bool& isFirstRepartition,
const std::array<bool, Dim>& allowedAxes);

@srikrrish srikrrish Jul 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Why do we need to have these overloads? Can't we just use this signature alone with a default argument of allowedAxes=true for all dimensions to preserve the legacy behaviour? Or is it just to avoid default arguments as it may not be a good practice (although I don't immediately see the disadvantages here)?

@aliemen aliemen Jul 2, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I think having a default would be nicer. In 3D one could simply write std::array<bool, Dim> = {true, true, true}. But since it's dimension independent, we would either need a helper or a lambda invokation. The latter would look like this (for your second example):

int findCutAxis(
    const NDIndex<Dim>& dom,
    const std::array<bool, Dim>& allowedAxes = [] {
        std::array<bool, Dim> axes{};
        axes.fill(true);
        return axes;
    }());

I'm honestly not sure what's cleaner. But the overload is perhaps easier to read?

If we really want a simple default argument, we could change the semantics to "forbiddenAxes". Then false would be the default ($\equiv$ all axes allowed) and one could write std::array<bool, Dim> allowedAxes = {}. Note that this is now by value, not by reference.

@aliemen aliemen Jul 2, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Just let me know what you prefer, I'm fine with all three versions.


/*!
* Find cutting axis as the longest axis of the field layout.
* @param dom Domain to reduce
*/
int findCutAxis(NDIndex<Dim>& 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<Dim>& dom, const std::array<bool, Dim>& allowedAxes);

@srikrrish srikrrish Jul 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Similar comment as above can't we have only one signature with a default argument?


/*!
* Check whether two domains overlap.
*/
bool domainsOverlap(const NDIndex<Dim>& lhs, const NDIndex<Dim>& rhs) const;

/*!
* Check whether proposed domains tile the global domain without cutting serial axes.
*/
bool domainsTileAllowedDecomposition(const std::vector<NDIndex<Dim>>& domains,
const NDIndex<Dim>& globalDomain,
const std::array<bool, Dim>& allowedAxes) const;

/*!
* Performs reduction on local field in all dimension except that determined
* by cutAxis, stores result in res
Expand Down
114 changes: 98 additions & 16 deletions src/Decomposition/OrthogonalRecursiveBisection.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,16 @@ namespace ippl {
template <typename Attrib>
bool OrthogonalRecursiveBisection<Field, Tp>::binaryRepartition(
const Attrib& R, FieldLayout<Dim>& fl, const bool& isFirstRepartition) {
std::array<bool, Dim> allowedAxes;
allowedAxes.fill(true);
return binaryRepartition(R, fl, isFirstRepartition, allowedAxes);
}

template <class Field, class Tp>
template <typename Attrib>
bool OrthogonalRecursiveBisection<Field, Tp>::binaryRepartition(
const Attrib& R, FieldLayout<Dim>& fl, const bool& isFirstRepartition,
const std::array<bool, Dim>& allowedAxes) {
// Timings
static IpplTimings::TimerRef tbasicOp = IpplTimings::getTimer("basicOperations");
static IpplTimings::TimerRef tperpReduction = IpplTimings::getTimer("perpReduction");
Expand Down Expand Up @@ -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<int>(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);
Expand All @@ -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
Expand All @@ -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<Dim>& globalDomain = fl.getDomain();
if (domains.size() != static_cast<size_t>(nprocs)) {
IpplTimings::stopTimer(tbasicOp);
return false;
}
if (!domainsTileAllowedDecomposition(domains, globalDomain, allowedAxes)) {
IpplTimings::stopTimer(tbasicOp);
return false;
}

// Update FieldLayout with new indices
Expand All @@ -123,11 +143,73 @@ namespace ippl {

template <class Field, class Tp>
int OrthogonalRecursiveBisection<Field, Tp>::findCutAxis(NDIndex<Dim>& 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<bool, Dim> allowedAxes;
allowedAxes.fill(true);
return findCutAxis(dom, allowedAxes);
}

template <class Field, class Tp>
int OrthogonalRecursiveBisection<Field, Tp>::findCutAxis(
const NDIndex<Dim>& dom, const std::array<bool, Dim>& 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 <class Field, class Tp>
bool OrthogonalRecursiveBisection<Field, Tp>::domainsOverlap(const NDIndex<Dim>& lhs,
const NDIndex<Dim>& 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 <class Field, class Tp>
bool OrthogonalRecursiveBisection<Field, Tp>::domainsTileAllowedDecomposition(
const std::vector<NDIndex<Dim>>& domains, const NDIndex<Dim>& globalDomain,
const std::array<bool, Dim>& 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 <class Field, class Tp>
Expand Down
75 changes: 74 additions & 1 deletion unit_tests/PIC/ORB.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ TYPED_TEST(ORBTest, Charge) {
typename TestFixture::value_type tol = tolerance<typename TestFixture::value_type>;

double charge = 0.5;
bunch->Q = charge/this->nParticles;
bunch->Q = charge / this->nParticles;

bunch->update();

Expand All @@ -164,7 +164,80 @@ TYPED_TEST(ORBTest, Charge) {
double totalCharge = field->sum();

ASSERT_NEAR((charge - totalCharge), 0., tol);
}

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<bool, Dim> 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<bool, Dim> 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[]) {
Expand Down
Loading