Change ORB to support arbitrary decomposition axes

[aliemen]

closes #559

This PR generalizes the ORB class to allow decomposing in arbitrary axis. Now you can e.g. pass a std::array<bool, Dim> allowedAxes = {false, false, true}; array to indicate that you only want to decompose in z. The actual algorithm is more or less untouched. The addition is mostly in the helpers like findAxis. Here an example in our SwissFEL-booster-sc input file from the OPALX side with decomposition only in z:

The decomposition works just fine only in z direction.

Other changes:

• Add allowed-axis cut selection:
  • findCutAxis(domain, allowedAxes) chooses the longest enabled axis,
  • returns failure if no enabled cut axis exists.
• Add validation before changing the layout:
  • reject selected cut axes with fewer than 4 cells,
  • clamp weighted median cuts to [1, cutAxisLength - 3],
  • reject one-cell domains,
  • reject boxes outside the global domain,
  • reject cuts along disabled axes,
  • reject overlapping boxes,
  • reject incomplete tiling of the global domain,
  • reject proposals whose domain count does not match the communicator size.
• Update FieldLayout and ORB weight field only after correct decomposition is found.

[aliemen]

cscs-ci run cscs-ci-gh200

[aliemen]

Jülich CI passes just fine. I don't know what cscs/cscs-ci-gh200 is doing, it says "no CUDA-capable device is detected" and doesn't re-start... Looks like it worked now

[aaadelmann]

General comment: before this can go into master we need to make sure it works with several k of GPUs. @srikrrish in ALPINE which of the tests are with ORB == ON?

[srikrrish]

General comment: before this can go into master we need to make sure it works with several k of GPUs. @srikrrish in ALPINE which of the tests are with ORB == ON?

Any of the Alpine example with load balancing threshold < 1 may invoke ORB (depending on the threshold). For example the usual value of 0.01 which means load balancing should be triggered when particle load imbalance in any rank exceeds 1% is a good start. I think Penning trap with more than 8 ranks might be the good stress test.

@aliemen yes the initial parallel/serial decomposition should be preserved during ORB also. If not that was a bug and if this PR fixes it that's a good thing I would say.

[srikrrish]

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]

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]

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

[srikrrish]

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