Poisson solvers for all types of domains (triply/doubly periodic, channels, boxes) and grids (regular, stretched)

[ali-ramadhan]

Implementing a vertically stretched grid (see PR #543) involves adding new Poisson solvers. This offers us a chance to:

1. refactor the Poisson solvers into a more flexible API,
2. support triply periodic and 3D box domains,
3. unite some of the CPU and GPU Poisson solvers,
4. reuse the Poisson solvers to implicitly time-step diffusion (vertically and possibly in 3D) in the future, and
5. slide in (or dispatch on) a distributed pressure solve for distributed/MPI models without having to modify the time stepping code for regular models in the future.

I am just throwing down some notes and ideas in this issue.

This table summarizes the different boundary conditions and grids we may want to simulate, and the Poisson solver algorithm we would need.

Boundary conditions	Grid	Solver	Alternatives
Triply periodic	Regular	FFT³
	Vertically stretched	FFT²×TRI?
Doubly periodic	Regular	FFT²×TRI	FFT²×DCT
	Vertically stretched	FFT²×TRI
Channel	Regular	FFT×DCT×TRI	FFT²×DCT
	Vertically stretched	FFT×DCT×TRI
Box	Regular	DCT²×TRI	DCT³
	Vertically stretched	DCT²×TRI

Essentially we need an FFT for periodic dimensions and we have the choice of using either a DCT or a TRIdiagonal solve for wall-bounded dimensions, although TRI may only be used once.

Some notes:

1. Multiple wall-bounded dimensions will require the use of a DCT somewhere.
2. Solvers for vertically stretched grids need to use TRI in the vertical.
3. For stretched grids in multiple dimensions I believe that a direct solve is no longer possible and an iterative method such as conjugate gradient must be used. I do not consider this case.
4. We have the option of using a DCT or a TRI. Which one we pick will depend on performance benchmarks between the two solvers. We should pick the faster one. This could depend on the number of vertical levels.
5. Unfortunately I don't think we can get rid of the in-place DCT algorithm that employs permutations on the GPU as it will be needed for channels and boxes at least. It may also turn out to be more efficient for regular grids in some cases.
6. We should reuse solvers as much as possible.
7. While it looks like we have many solvers to implement, I think if we abstract away the steps then each solver just needs to implement something like a forward_x_transform, forward_y_transform, forward_z_transform, backward_z_transform, etc. then we can implement many solvers without repeating code.

For now, I will first just focus on getting Poisson solvers working for vertically stretched grids in doubly periodic domains on the CPU and GPU. This will offer a chance to do some refactoring and benchmarking. We also need to try out faster tridiagonal matrix algorithms.

But I'd like to keep the five initial goals in mind. This issue will probably be open for a while.

[glwagner]

It'd be nice to state in writing the justification for writing a separate CPU solver for certain problems. In general, I think that any algorithm that works on the GPU will also work on the CPU. Thus at least in principle the simplest choice is presumably to use the same solver on both architectures.

For example, benchmarking might show that a GPU-friendly algorithm performs poorly compared to a CPU-specific algorithm, which might justify maintaining separate solvers for the GPU and CPU. Is this the case?

[glwagner]

Also, I'd encourage writing this code into as self-contained a submodule as possible. I think there are other codes in the julia ecosystem (not least FourierFlows.jl!) that would benefit from fast and multi-architecture Poisson solvers. We don't have to break this into a separate package just yet, but we do want to ensure this is easy to do in the future.

[ali-ramadhan]

It'd be nice to state in writing the justification for writing a separate CPU solver for certain problems. In general, I think that any algorithm that works on the GPU will also work on the CPU. Thus at least in principle the simplest choice is presumably to use the same solver on both architectures.

Good point, I've been meaning to set up a script for benchmarking the different pressure solvers. We should use performance benchmarking results to make decisions.

Also, I'd encourage writing this code into as self-contained a submodule as possible. I think there are other codes in the julia ecosystem (not least FourierFlows.jl!) that would benefit from fast and multi-architecture Poisson solvers. We don't have to break this into a separate package just yet, but we do want to ensure this is easy to do in the future.

Another good point. As you pointed out some of these solvers depend on the grid but if we take that out (which would be trivial) then I think the solvers in PR #589 would be pretty reusable by other packages. Although right now they're pretty specific to staggered grids (except for BatchedTridiagonalSolver).

[glwagner]

Almost there...