Distributed tridiagonal Fourier solver

src/DistributedComputations/DistributedComputations.jl

@@ -23,5 +23,6 @@ include("transposable_field.jl")
 include("distributed_transpose.jl")
 include("plan_distributed_transforms.jl")
 include("distributed_fft_based_poisson_solver.jl")
+include("distributed_fft_tridiagonal_solver.jl")
 
 end # module

src/DistributedComputations/distributed_architectures.jl

@@ -24,7 +24,7 @@ end
 Return `Partition` representing the division of a domain in
 the `x` (first), `y` (second) and `z` (third) dimension
 
-Keyword arguments: 
+Keyword arguments:
 ==================
 
 - `x`: partitioning of the first dimension 

src/DistributedComputations/distributed_fft_based_poisson_solver.jl

@@ -2,6 +2,7 @@ import FFTW
 
 using CUDA: @allowscalar
 using Oceananigans.Grids: XYZRegularRG
+using Oceananigans.Solvers: stretched_dimensions, stretched_direction
 
 import Oceananigans.Solvers: poisson_eigenvalues, solve!
 import Oceananigans.Architectures: architecture
@@ -60,9 +61,9 @@ In the algorithm below, the first dimension is always the local dimension. In ou
 
 1. `storage.zfield`, partitioned over ``(x, y)`` is initialized with the `rhs` that is ``b``.
 2. Transform along ``z``.
-3  Transpose + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``.
+3  Transpose `storage.zfield` + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``.
 4. Transform along ``y``.
-5. Transpose + communicate to `storage.xfield` partitioned into `(Rx, Ry)` processes in ``(y, z)``.
+5. Transpose `storage.yfield` + communicate to `storage.xfield` partitioned into `(Rx, Ry)` processes in ``(y, z)``.
 6. Transform in ``x``.
 
 At this point the three in-place forward transforms are complete, and we
@@ -91,8 +92,8 @@ Restrictions
 """
 function DistributedFFTBasedPoissonSolver(global_grid, local_grid, planner_flag=FFTW.PATIENT)
 
-    validate_global_grid(global_grid)
-    validate_configuration(global_grid, local_grid)
+    validate_poisson_solver_distributed_grid(global_grid)
+    validate_poisson_solver_configuration(global_grid, local_grid)
 
     FT = Complex{eltype(local_grid)}
 
@@ -188,38 +189,41 @@ end
 end
 
 # TODO: bring up to speed the PCG to remove this error
-validate_global_grid(global_grid) = 
+validate_poisson_solver_distributed_grid(global_grid) = 
         throw("Grids other than the RectilinearGrid are not supported in the Distributed NonhydrostaticModels")
 
-function validate_global_grid(global_grid::RectilinearGrid) 
+function validate_poisson_solver_distributed_grid(global_grid::RectilinearGrid) 
     TX, TY, TZ = topology(global_grid)
 
     if (TY == Bounded && TZ == Periodic) || (TX == Bounded && TY == Periodic) || (TX == Bounded && TZ == Periodic)
-        throw("NonhydrostaticModels on Distributed grids do not support topology ($TX, $TY, $TZ) at the moment.
-               TZ Periodic requires also TY and TX to be Periodic, while TY Periodic requires also TX to be Periodic.
-               Please rotate the domain to obtain the required topology")
+        throw("Distributed Poisson solvers do not support grids with topology ($TX, $TY, $TZ) at the moment.
+               A Periodic z-direction requires also the y- and and x-directions to be Periodic, while a Periodic y-direction requires also 
+               the x-direction to be Periodic.")
     end
 
-    if !(global_grid isa XYZRegularRG) 
-        throw("Stretched directions are not supported with distributed grids at the moment.")
+    if !(global_grid isa YZRegularRG) && !(global_grid isa XYRegularRG) && !(global_grid isa XZRegularRG) 
+        throw("The provided grid is stretched in directions $(stretched_dimensions(global_grid)). 
+               A distributed Poisson solver supports only RectilinearGrids that have variably-spaced cells in at most one direction.")
     end
 
     return nothing
 end
 
-function validate_configuration(global_grid, local_grid)
+function validate_poisson_solver_configuration(global_grid, local_grid)
 
     # We don't support distributing anything in z.
-    Rz = architecture(local_grid).ranks[3]
-    Rz == 1 || throw("Non-singleton ranks in the vertical are not supported by DistributedFFTBasedPoissonSolver.")
+    Rx, Ry, Rz = architecture(local_grid).ranks
+    Rz == 1 || throw("Non-singleton ranks in the vertical are not supported by distributed Poisson solvers.")
 
     # Limitation of the current implementation (see the docstring)
-    if global_grid.Nz % architecture(local_grid).ranks[2] != 0
-        throw("The number of ranks in the y direction must divide Nz. See the docstring for more information.")
+    if global_grid.Nz % Ry != 0
+        throw("The number of ranks in the y-direction are $(Ry) with Nz = $(global_grid.Nz) cells in the z-direction.
+               The distributed Poisson solver requires that the number of ranks in the y-direction divide Nz.")
     end
 
-    if global_grid.Ny % architecture(local_grid).ranks[1] != 0
-        throw("The number of ranks in the x direction must divide Ny. See the docstring for more information.")
+    if global_grid.Ny % Rx != 0
+        throw("The number of ranks in the y-direction are $(Rx) with Ny = $(global_grid.Ny) cells in the y-direction.
+               The distributed Poisson solver requires that the number of ranks in the x-direction divide Ny.")
     end
 
     return nothing