(0.88.0) MPI communication and computation overlap in the HydrostaticFreeSurfaceModel and NonhydrostaticModel

Manifest.toml

@@ -1,8 +1,8 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.8.5"
+julia_version = "1.8.0"
 manifest_format = "2.0"
-project_hash = "efa1b3c0f878c0fe77caba96f053e1145d126cd0"
+project_hash = "b61348c5ba4009d3da1a3d8c47bdeb84513faa2c"
 
 [[deps.AbstractFFTs]]
 deps = ["ChainRulesCore", "LinearAlgebra"]
@@ -106,7 +106,7 @@ version = "4.6.1"
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
-version = "1.0.1+0"
+version = "0.5.2+0"
 
 [[deps.Crayons]]
 git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15"
@@ -277,7 +277,9 @@ version = "1.12.0"
 
 [[deps.KernelAbstractions]]
 deps = ["Adapt", "Atomix", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"]
-git-tree-sha1 = "47be64f040a7ece575c2b5f53ca6da7b548d69f4"
+git-tree-sha1 = "bbb7ac4a3194c0d1561b9dea2a20e8f1ab68f709"
+repo-rev = "main"
+repo-url = "https://github.com/simone-silvestri/KernelAbstractions.jl"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 version = "0.9.4"
 
@@ -668,7 +670,7 @@ version = "1.10.1"
 [[deps.Tar]]
 deps = ["ArgTools", "SHA"]
 uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
-version = "1.10.1"
+version = "1.10.0"
 
 [[deps.TaylorSeries]]
 deps = ["LinearAlgebra", "Markdown", "Requires", "SparseArrays"]

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.82.0"
+version = "0.83.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -45,7 +45,6 @@ Glob = "1.3"
 IncompleteLU = "0.2"
 IterativeSolvers = "0.9"
 JLD2 = "^0.4"
-KernelAbstractions = "0.9"
 MPI = "0.16, 0.17, 0.18, 0.19, 0.20"
 NCDatasets = "0.12.10"
 OffsetArrays = "1.4"

src/Advection/centered_reconstruction.jl

@@ -100,8 +100,8 @@ for (dir, ξ, val) in zip((:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ), (:x, :y, :z),
 
     for buffer in [1, 2, 3, 4, 5, 6]
         @eval begin
-            @inline $stencil(i, j, k, grid, scheme::Centered{$buffer}, ψ, idx, loc, args...)           = @inbounds sum($(reconstruction_stencil(buffer, :symm, ξ, false)) .* retrieve_coeff(scheme, Val($val), idx, loc))
-            @inline $stencil(i, j, k, grid, scheme::Centered{$buffer}, ψ::Function, idx, loc, args...) = @inbounds sum($(reconstruction_stencil(buffer, :symm, ξ,  true)) .* retrieve_coeff(scheme, Val($val), idx, loc))
+            @inline $stencil(i, j, k, grid, scheme::Centered{$buffer, FT}, ψ, idx, loc, args...)           where FT = @inbounds sum($(reconstruction_stencil(buffer, :symm, ξ, false)) .* retrieve_coeff(scheme, Val($val), idx, loc))
+            @inline $stencil(i, j, k, grid, scheme::Centered{$buffer, FT}, ψ::Function, idx, loc, args...) where FT = @inbounds sum($(reconstruction_stencil(buffer, :symm, ξ,  true)) .* retrieve_coeff(scheme, Val($val), idx, loc))
         end
     end
 end

src/Advection/reconstruction_coefficients.jl

@@ -123,19 +123,19 @@ Examples
 julia> using Oceananigans.Advection: calc_reconstruction_stencil
 
 julia> calc_reconstruction_stencil(1, :right, :x)
-:(+(coeff1_right[1] * ψ[i + 0, j, k]))
+:(+(FT(coeff1_right[1]) * ψ[i + 0, j, k]))
 
 julia> calc_reconstruction_stencil(1, :left, :x)
-:(+(coeff1_left[1] * ψ[i + -1, j, k]))
+:(+(FT(coeff1_left[1]) * ψ[i + -1, j, k]))
 
 julia> calc_reconstruction_stencil(1, :symm, :x)
-:(coeff2_symm[2] * ψ[i + -1, j, k] + coeff2_symm[1] * ψ[i + 0, j, k])
+:(FT(coeff2_symm[2]) * ψ[i + -1, j, k] + FT(coeff2_symm[1]) * ψ[i + 0, j, k])
 
 julia> calc_reconstruction_stencil(2, :symm, :x)
-:(coeff4_symm[4] * ψ[i + -2, j, k] + coeff4_symm[3] * ψ[i + -1, j, k] + coeff4_symm[2] * ψ[i + 0, j, k] + coeff4_symm[1] * ψ[i + 1, j, k])
+:(FT(coeff4_symm[4]) * ψ[i + -2, j, k] + FT(coeff4_symm[3]) * ψ[i + -1, j, k] + FT(coeff4_symm[2]) * ψ[i + 0, j, k] + FT(coeff4_symm[1]) * ψ[i + 1, j, k])
 
 julia> calc_reconstruction_stencil(3, :left, :x)
-:(coeff5_left[5] * ψ[i + -3, j, k] + coeff5_left[4] * ψ[i + -2, j, k] + coeff5_left[3] * ψ[i + -1, j, k] + coeff5_left[2] * ψ[i + 0, j, k] + coeff5_left[1] * ψ[i + 1, j, k])
+:(FT(coeff5_left[5]) * ψ[i + -3, j, k] + FT(coeff5_left[4]) * ψ[i + -2, j, k] + FT(coeff5_left[3]) * ψ[i + -1, j, k] + FT(coeff5_left[2]) * ψ[i + 0, j, k] + FT(coeff5_left[1]) * ψ[i + 1, j, k])
 ```
 """
 @inline function calc_reconstruction_stencil(buffer, shift, dir, func::Bool = false)
@@ -154,16 +154,16 @@ julia> calc_reconstruction_stencil(3, :left, :x)
         c = n - buffer - 1
         if func
             stencil_full[idx] = dir == :x ? 
-                                :($coeff[$(order - idx + 1)] * ψ(i + $c, j, k, grid, args...)) :
+                                :(FT($coeff[$(order - idx + 1)]) * ψ(i + $c, j, k, grid, args...)) :
                                 dir == :y ?
-                                :($coeff[$(order - idx + 1)] * ψ(i, j + $c, k, grid, args...)) :
-                                :($coeff[$(order - idx + 1)] * ψ(i, j, k + $c, grid, args...))
+                                :(FT($coeff[$(order - idx + 1)]) * ψ(i, j + $c, k, grid, args...)) :
+                                :(FT($coeff[$(order - idx + 1)]) * ψ(i, j, k + $c, grid, args...))
         else
             stencil_full[idx] =  dir == :x ? 
-                                :($coeff[$(order - idx + 1)] * ψ[i + $c, j, k]) :
+                                :(FT($coeff[$(order - idx + 1)]) * ψ[i + $c, j, k]) :
                                 dir == :y ?
-                                :($coeff[$(order - idx + 1)] * ψ[i, j + $c, k]) :
-                                :($coeff[$(order - idx + 1)] * ψ[i, j, k + $c])
+                                :(FT($coeff[$(order - idx + 1)]) * ψ[i, j + $c, k]) :
+                                :(FT($coeff[$(order - idx + 1)]) * ψ[i, j, k + $c])
         end
     end
     return Expr(:call, :+, stencil_full...)

src/Advection/upwind_biased_reconstruction.jl

@@ -125,8 +125,8 @@ for (sd, side) in enumerate((:left, :right)), (dir, ξ, val) in zip((:xᶠᵃᵃ
 
     for buffer in [1, 2, 3, 4, 5, 6]
         @eval begin
-            @inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer}, ψ, idx, loc, args...)           = @inbounds sum($(reconstruction_stencil(buffer, side, ξ, false)) .* retrieve_coeff(scheme, Val($sd), Val($val), idx, loc))
-            @inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer}, ψ::Function, idx, loc, args...) = @inbounds sum($(reconstruction_stencil(buffer, side, ξ,  true)) .* retrieve_coeff(scheme, Val($sd), Val($val), idx, loc))
+            @inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer, FT}, ψ, idx, loc, args...)           where FT = @inbounds sum($(reconstruction_stencil(buffer, side, ξ, false)) .* retrieve_coeff(scheme, Val($sd), Val($val), idx, loc))
+            @inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer, FT}, ψ::Function, idx, loc, args...) where FT = @inbounds sum($(reconstruction_stencil(buffer, side, ξ,  true)) .* retrieve_coeff(scheme, Val($sd), Val($val), idx, loc))
         end
     end
 end

src/Advection/weno_reconstruction.jl

@@ -114,7 +114,7 @@ function WENO(FT::DataType=Float64;
 
     if order < 3
         # WENO(order = 1) is equivalent to UpwindBiased(order = 1)
-        return UpwindBiased(order = 1)
+        return UpwindBiased(FT; order = 1)
     else
         N  = Int((order + 1) ÷ 2)
 

src/Architectures.jl

@@ -86,7 +86,7 @@ function unified_array(::GPU, arr::AbstractArray)
     return vec
 end
 
-## Only for contiguous data!! (i.e. only if the offset for pointer(dst::CuArrat, offset::Int) is 1)
+## Only for contiguous data!! (i.e. only if the offset for pointer(dst::CuArray, offset::Int) is 1)
 @inline function device_copy_to!(dst::CuArray, src::CuArray; async::Bool = false) 
     n = length(src)
     context!(context(src)) do

src/BoundaryConditions/BoundaryConditions.jl

@@ -14,7 +14,7 @@ using CUDA
 using KernelAbstractions: @index, @kernel
 
 using Oceananigans.Architectures: CPU, GPU, device
-using Oceananigans.Utils: work_layout, launch!
+using Oceananigans.Utils: work_layout, launch!, KernelParameters
 using Oceananigans.Operators: Ax, Ay, Az, volume
 using Oceananigans.Grids
 

src/BoundaryConditions/apply_flux_bcs.jl

@@ -9,8 +9,8 @@ using Oceananigans.Grids: AbstractGrid
 #####
 
 # Unpack
-apply_x_bcs!(Gc, c, args...) = apply_x_bcs!(Gc, Gc.grid, c, c.boundary_conditions.west, c.boundary_conditions.east, args...)
-apply_y_bcs!(Gc, c, args...) = apply_y_bcs!(Gc, Gc.grid, c, c.boundary_conditions.south, c.boundary_conditions.north, args...)
+apply_x_bcs!(Gc, c, args...) = apply_x_bcs!(Gc, Gc.grid, c, c.boundary_conditions.west,   c.boundary_conditions.east, args...)
+apply_y_bcs!(Gc, c, args...) = apply_y_bcs!(Gc, Gc.grid, c, c.boundary_conditions.south,  c.boundary_conditions.north, args...)
 apply_z_bcs!(Gc, c, args...) = apply_z_bcs!(Gc, Gc.grid, c, c.boundary_conditions.bottom, c.boundary_conditions.top, args...)
 
 # Shortcuts for...

src/BoundaryConditions/fill_halo_regions.jl

@@ -9,8 +9,8 @@ import Base
 ##### General halo filling functions
 #####
 
-fill_halo_regions!(::Nothing, args...) = nothing
-fill_halo_regions!(::NamedTuple{(), Tuple{}}, args...) = nothing
+fill_halo_regions!(::Nothing, args...; kwargs...) = nothing
+fill_halo_regions!(::NamedTuple{(), Tuple{}}, args...; kwargs...) = nothing
 
 """
     fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...)
@@ -163,28 +163,22 @@ fill_first(bc1, bc2)               = true
 ##### General fill_halo! kernels
 #####
 
-@kernel function _fill_west_and_east_halo!(c, west_bc, east_bc, offset, loc, grid, args) 
+@kernel function _fill_west_and_east_halo!(c, west_bc, east_bc, loc, grid, args) 
     j, k = @index(Global, NTuple)
-    j′ = j + offset[1]
-    k′ = k + offset[2]
-    _fill_west_halo!(j′, k′, grid, c, west_bc, loc, args...)
-    _fill_east_halo!(j′, k′, grid, c, east_bc, loc, args...)
+    _fill_west_halo!(j, k, grid, c, west_bc, loc, args...)
+    _fill_east_halo!(j, k, grid, c, east_bc, loc, args...)
 end
 
-@kernel function _fill_south_and_north_halo!(c, south_bc, north_bc, offset, loc, grid, args)
+@kernel function _fill_south_and_north_halo!(c, south_bc, north_bc, loc, grid, args)
     i, k = @index(Global, NTuple)
-    i′ = i + offset[1]
-    k′ = k + offset[2]
-    _fill_south_halo!(i′, k′, grid, c, south_bc, loc, args...)
-    _fill_north_halo!(i′, k′, grid, c, north_bc, loc, args...)
+    _fill_south_halo!(i, k, grid, c, south_bc, loc, args...)
+    _fill_north_halo!(i, k, grid, c, north_bc, loc, args...)
 end
 
-@kernel function _fill_bottom_and_top_halo!(c, bottom_bc, top_bc, offset, loc, grid, args)
+@kernel function _fill_bottom_and_top_halo!(c, bottom_bc, top_bc, loc, grid, args)
     i, j = @index(Global, NTuple)
-    i′ = i + offset[1]
-    j′ = j + offset[2]
-    _fill_bottom_halo!(i′, j′, grid, c, bottom_bc, loc, args...)
-       _fill_top_halo!(i′, j′, grid, c, top_bc,    loc, args...)
+    _fill_bottom_halo!(i, j, grid, c, bottom_bc, loc, args...)
+       _fill_top_halo!(i, j, grid, c, top_bc,    loc, args...)
 end
 
 #####
@@ -193,7 +187,7 @@ end
 
 import Oceananigans.Utils: @constprop
 
-@kernel function _fill_west_and_east_halo!(c::NTuple, west_bc, east_bc, offset, loc, grid, args)
+@kernel function _fill_west_and_east_halo!(c::NTuple, west_bc, east_bc, loc, grid, args)
     j, k = @index(Global, NTuple)
     ntuple(Val(length(west_bc))) do n
         Base.@_inline_meta
@@ -205,7 +199,7 @@ import Oceananigans.Utils: @constprop
     end
 end
 
-@kernel function _fill_south_and_north_halo!(c::NTuple, south_bc, north_bc, offset, loc, grid, args) 
+@kernel function _fill_south_and_north_halo!(c::NTuple, south_bc, north_bc, loc, grid, args) 
     i, k = @index(Global, NTuple)
     ntuple(Val(length(south_bc))) do n
         Base.@_inline_meta
@@ -217,7 +211,7 @@ end
     end
 end
 
-@kernel function _fill_bottom_and_top_halo!(c::NTuple, bottom_bc, top_bc, offset, loc, grid, args) 
+@kernel function _fill_bottom_and_top_halo!(c::NTuple, bottom_bc, top_bc, loc, grid, args) 
     i, j = @index(Global, NTuple)
     ntuple(Val(length(bottom_bc))) do n
         Base.@_inline_meta
@@ -230,13 +224,13 @@ end
 end
 
 fill_west_and_east_halo!(c, west_bc, east_bc, size, offset, loc, arch, grid, args...; kwargs...) =
-    launch!(arch, grid, size, _fill_west_and_east_halo!, c, west_bc, east_bc, offset, loc, grid, Tuple(args); kwargs...)
+    launch!(arch, grid, KernelParameters(size, offset), _fill_west_and_east_halo!, c, west_bc, east_bc, loc, grid, Tuple(args); kwargs...)
 
 fill_south_and_north_halo!(c, south_bc, north_bc, size, offset, loc, arch, grid, args...; kwargs...) =
-    launch!(arch, grid, size, _fill_south_and_north_halo!, c, south_bc, north_bc, offset, loc, grid, Tuple(args); kwargs...)
+    launch!(arch, grid, KernelParameters(size, offset), _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...)
 
 fill_bottom_and_top_halo!(c, bottom_bc, top_bc, size, offset, loc, arch, grid, args...; kwargs...) =
-    launch!(arch, grid, size, _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, offset, loc, grid, Tuple(args); kwargs...)
+    launch!(arch, grid, KernelParameters(size, offset), _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, loc, grid, Tuple(args); kwargs...)
 
 #####
 ##### Calculate kernel size and offset for Windowed and Sliced Fields

src/BoundaryConditions/fill_halo_regions_flux.jl

@@ -38,62 +38,50 @@ using KernelAbstractions.Extras.LoopInfo: @unroll
 ##### Single halo filling kernels
 #####
 
-@kernel function fill_flux_west_halo!(c, offset, grid)
+@kernel function fill_flux_west_halo!(c, grid)
     j, k = @index(Global, NTuple)
-    j′ = j + offset[1]
-    k′ = k + offset[2]
-    _fill_flux_west_halo!(1, j′, k′, grid, c)
+    _fill_flux_west_halo!(1, j, k, grid, c)
 end
 
-@kernel function fill_flux_south_halo!(c, offset, grid)
+@kernel function fill_flux_south_halo!(c, grid)
     i, k = @index(Global, NTuple)
-    i′ = i + offset[1]
-    k′ = k + offset[2]
-    _fill_flux_south_halo!(i′, 1, k′, grid, c)
+    _fill_flux_south_halo!(i, 1, k, grid, c)
 end
 
-@kernel function fill_flux_bottom_halo!(c, offset, grid)
+@kernel function fill_flux_bottom_halo!(c, grid)
     i, j = @index(Global, NTuple)
-    i′ = i + offset[1]
-    j′ = j + offset[2]
-    _fill_flux_bottom_halo!(i′, j′, 1, grid, c)
+    _fill_flux_bottom_halo!(i, j, 1, grid, c)
 end
 
-@kernel function fill_flux_east_halo!(c, offset, grid)
+@kernel function fill_flux_east_halo!(c, grid)
     j, k = @index(Global, NTuple)
-    j′ = j + offset[1]
-    k′ = k + offset[2]
-    _fill_flux_east_halo!(1, j′, k′, grid, c)
+    _fill_flux_east_halo!(1, j, k, grid, c)
 end
 
-@kernel function fill_flux_north_halo!(c, offset, grid)
+@kernel function fill_flux_north_halo!(c, grid)
     i, k = @index(Global, NTuple)
-    i′ = i + offset[1]
-    k′ = k + offset[2]
-    _fill_flux_north_halo!(i′, 1, k′, grid, c)
+    _fill_flux_north_halo!(i, 1, k, grid, c)
 end
 
-@kernel function fill_flux_top_halo!(c, offset, grid)
+@kernel function fill_flux_top_halo!(c, grid)
     i, j = @index(Global, NTuple)
-    j′ = j + offset[1]
-    k′ = k + offset[2]
-    _fill_flux_top_halo!(i′, j′, 1, grid, c)
+    _fill_flux_top_halo!(i, j, 1, grid, c)
 end
 
 #####
 ##### Kernel launchers for flux boundary conditions
 #####
 
 fill_west_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, kernel_size, fill_flux_west_halo!, c, offset,grid; kwargs...)
+            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_west_halo!, c,grid; kwargs...)
 fill_east_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, kernel_size, fill_flux_east_halo!, c, offset,grid; kwargs...)
+            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_east_halo!, c, grid; kwargs...)
 fill_south_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, kernel_size, fill_flux_south_halo!, c, offset,grid; kwargs...)
+            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_south_halo!, c, grid; kwargs...)
 fill_north_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, kernel_size, fill_flux_north_halo!, c, offset,grid; kwargs...)
+            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_north_halo!, c, grid; kwargs...)
 fill_bottom_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, kernel_size, fill_flux_bottom_halo!, c, offset,grid; kwargs...)
+            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_bottom_halo!, c, grid; kwargs...)
 fill_top_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, kernel_size, fill_flux_top_halo!, c, offset, grid; kwargs...)
+            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_top_halo!, c, grid; kwargs...)