As written, the halo update code in Pace does not support non-square MPI layouts (e.g. where a tile is decomposed into 3x2 ranks). This is because the abstractions for partitioning of ranks currently assume that for a given rank and a given direction, there is only one rank bordering in that direction. Fortunately, this assumption is not required by anything using the code which has this abstraction.
To fix this, the Partitioner class must be refactored, replacing the boundary(self, boundary_type: int, rank: int) -> Optional[bd.SimpleBoundary]: method with a boundaries(self): -> Iterable[bd.SimpleBoundary] method, and code which currently calls boundary while iterating over directions should instead iterate over the boundaries. Once this is done, the implementation of boundaries can be modified to support non-square layouts, which should make it so that all halo update code works for non-square layouts.
I would strongly recommend test-driven development for this modification, similar to the current partitioner boundaries tests. In developing the current code, it was very, very helpful to be able to try various "guesses" as to what code would work until something passed the (extensive) test cases.
As written, the halo update code in Pace does not support non-square MPI layouts (e.g. where a tile is decomposed into 3x2 ranks). This is because the abstractions for partitioning of ranks currently assume that for a given rank and a given direction, there is only one rank bordering in that direction. Fortunately, this assumption is not required by anything using the code which has this abstraction.
To fix this, the Partitioner class must be refactored, replacing the
boundary(self, boundary_type: int, rank: int) -> Optional[bd.SimpleBoundary]:method with aboundaries(self): -> Iterable[bd.SimpleBoundary]method, and code which currently callsboundarywhile iterating over directions should instead iterate over theboundaries. Once this is done, the implementation ofboundariescan be modified to support non-square layouts, which should make it so that all halo update code works for non-square layouts.I would strongly recommend test-driven development for this modification, similar to the current partitioner boundaries tests. In developing the current code, it was very, very helpful to be able to try various "guesses" as to what code would work until something passed the (extensive) test cases.