Oceananigans.jl/src/DistributedComputations/DistributedComputations.jl at 329632d52525c145167124ec2138b60f97450446 · CliMA/Oceananigans.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
module DistributedComputations

export
    Distributed, Partition, Equal, Fractional,
    child_architecture, reconstruct_global_grid, partition,
    inject_halo_communication_boundary_conditions,
    DistributedFFTBasedPoissonSolver, mpi_initialized, mpi_rank,
    mpi_size, global_barrier, global_communicator,
    @root, @onrank, @distribute, @handshake

using MPI

using Oceananigans
using Oceananigans.Utils
using Oceananigans.Grids
using OffsetArrays
using Oceananigans.Grids: XYZRegularRG
using Oceananigans.Solvers: GridWithFourierTridiagonalSolver

import Oceananigans.Solvers: fft_poisson_solver

include("distributed_macros.jl")
include("distributed_architectures.jl")
include("partition_assemble.jl")
include("distributed_grids.jl")
include("distributed_immersed_boundaries.jl")
include("distributed_on_architecture.jl")
include("distributed_kernel_launching.jl")
include("halo_communication_bcs.jl")
include("communication_buffers.jl")
include("distributed_fields.jl")
include("halo_communication.jl")
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")

fft_poisson_solver(grid::DistributedRectilinearGrid) = fft_poisson_solver(grid, reconstruct_global_grid(grid))

fft_poisson_solver(local_grid::DistributedRectilinearGrid, global_grid::XYZRegularRG) =
    DistributedFFTBasedPoissonSolver(global_grid, local_grid)

fft_poisson_solver(local_grid::DistributedRectilinearGrid, global_grid::GridWithFourierTridiagonalSolver) =
    DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid)

import Oceananigans.Solvers: compute_preconditioner_rhs!, precondition!

# But we need to define the precondition! methods here
function precondition!(p, preconditioner::DistributedFFTBasedPoissonSolver, r, args...)
    compute_preconditioner_rhs!(preconditioner, r)
    solve!(p, preconditioner)
    return p
end

function precondition!(p, preconditioner::DistributedFourierTridiagonalPoissonSolver, r, args...)
    compute_preconditioner_rhs!(preconditioner, r)
    solve!(p, preconditioner)
    return p
end

# Correctly pass architecture to determine the default weno_weight_computation
Oceananigans.Advection.default_weno_weight_computation(arch::Distributed) =
    Oceananigans.Advection.default_weno_weight_computation(child_architecture(arch))


end # module