Distributed tridiagonal Fourier solver

[simone-silvestri]

This PR introduces a distributed pressure solver for grids that are stretched in one direction. The algorithm implemented here is the same as the one described in #2538 but it uses the FFT and transposes built in #3279

[glwagner]

Just a question about implementation:
It looks like #2538 implemented just one solver, with an optional tridiagonal component, is that right. Why does this PR take a different approach? Wouldn't using a single solver result in less code / duplication of transform logic?

Here, the approach follows the serial implementation quite closely, where we have an FFT solver and a FourierTridiagonal solver. Additionally, the fields of the solvers are tailored to their task; for example, there is no tridiagonal solver in the pure FFT solver. You might argue that we just put nothing if we don't need it, but then we have more "ambiguous" fields like eigenvalues that are not needed in the tridiagonal solver and source_term that is not needed in the FFT. I am not convinced that a single solver would lead to less (or cleaner) code. Mostly because the underlying code (constructor and solve! functions) is compact and slender enough to justify writing individual functions for different grids (the number of lines in the docstring for the DistributedFourierTridiagonalPoissonSolver is comparable to the code related to the solver). I think this improves the interpretability of the algorithm. Maybe an improvement would be writing a unified constructor that would spit out the different solvers.

2. It might make sense to split the constructor into two parts, so that we can build a tridiagonal solver even when all three directions are regular. That could be useful for testing, for example.

I have added this capability by passing the stretched_direction kwarg.

Also note that in terms of operation count the tridiagonal solve is cheaper than FFT...

I think, all things considered, the mixed FFT / tridiagonal solve will have basically the same computational cost as the pure FFT solve only for a stretched x direction. The additional transposes required for a y or stretched z direction will completely dominate the cost of the actual operations.

As an example, this is a slab decomposition with a fairly big grid (512 x 256^2) split on 2 GPUs on Tartarus The AlltoAllv is the dominant cost, while the FFT (in between the two transposes) is quite irrelevant. In the near future, I'll perform scaling tests on Perlmutter, which has a much better network, so it might be that (even if I think it's unlikely) the cost is not all communication after all.

Just to clarify --- the mixed tridiagonal + FFT solver also needs eigenvalues, doesn't it?

[simone-silvestri]

Just to clarify --- the mixed tridiagonal + FFT solver also needs eigenvalues, doesn't it?

yep, it does but they are "embedded" in the diagonal terms of the tridiagonal solver. So to clarify, there is no need for an additional eigenvalue field because these are already included in bathced_tridiagonal_solver

[tomchor]

I meant to ask this on the other PR but forgot: how are the scalings for this algorithm? Any recommendations when using GPU partitioning on nonhydrostatic models?

[simone-silvestri]

A long time ago, I did some scaling tests of the pure FFT algorithm. These were the results

I will probably redo the scaling test later on after the summer.
In general, always use slab partitioning if you can because you avoid one transposition.
This might lead to larger halo to domain ratio, but the fill halo, contrarily to the transpose, is hidden so it should still be better to have a slab partitioning then a pencil one.

[tomchor]

Thanks that's super useful info!