A rework for the advection module that (should) improve WENO performance

src/Advection/Advection.jl

@@ -44,6 +44,8 @@ import Base: show, summary, Callable
 import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z
 import Oceananigans.Architectures: on_architecture
 
+using Oceananigans.Grids: XFlatGrid, YFlatGrid, ZFlatGrid
+
 abstract type AbstractAdvectionScheme{B, FT} end
 abstract type AbstractCenteredAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end
 abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end
@@ -52,11 +54,11 @@ abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionSch
 # are constructed via metaprogramming. (The `advection_buffer` is the width of
 # the halo region required for an advection scheme on a non-immersed-boundary grid.)
 # An upper limit of `advection_buffer = 6` means we can build advection schemes up to
-# `Centered(order=12`) and `UpwindBiased(order=11)`. The list can be extended in order to
+# `Centered(order=10`) and `UpwindBiased(order=9)`. The list can be extended in order to
 # compile schemes with higher orders; for example `advection_buffers = [1, 2, 3, 4, 5, 6, 8]`
 # will compile schemes for `advection_buffer=8` and thus `Centered(order=16)` and `UpwindBiased(order=15)`.
 # Note that it is not possible to compile schemes for `advection_buffer = 41` or higher.
-const advection_buffers = [1, 2, 3, 4, 5, 6]
+const advection_buffers = [1, 2, 3, 4, 5]
 
 @inline Base.eltype(::AbstractAdvectionScheme{<:Any, FT}) where FT = FT
 

src/Advection/adapt_advection_order.jl

@@ -76,21 +76,22 @@ function adapt_advection_order(advection::Centered{B}, N::Int, grid::AbstractGri
     if N >= B
         return advection
     else
-        return Centered(; order=2N)
+        return Centered(grid; order=2N)
     end
 end
 
 function adapt_advection_order(advection::UpwindBiased{B}, N::Int, grid::AbstractGrid) where B
     if N >= B
         return advection
     else
-        return UpwindBiased(; order=2N-1)
+        return UpwindBiased(grid; order=2N-1)
     end
 end
+
 function adapt_advection_order(advection::WENO{B}, N::Int, grid::AbstractGrid) where B
     if N >= B
         return advection
     else
-        return WENO(order=2N-1)
+        return WENO(grid; order=2N-1)
     end
 end

src/Advection/centered_reconstruction.jl

@@ -7,64 +7,78 @@
 
 Centered reconstruction scheme.
 """
-struct Centered{N, FT, CA} <: AbstractCenteredAdvectionScheme{N, FT}
-    "advection scheme used near boundaries"
-    buffer_scheme :: CA
-
-    Centered{N, FT}(buffer_scheme::CA) where {N, FT, CA} = new{N, FT, CA}(buffer_scheme)
-end
+struct Centered{N, FT} <: AbstractCenteredAdvectionScheme{N, FT} end
 
 function Centered(FT::DataType=Oceananigans.defaults.FloatType; grid = nothing, order = 2)
 
-    if !(grid isa Nothing)
-        FT = eltype(grid)
-    end
-
+    # Enforce the grid type if a grid is provided
+    FT = grid isa Nothing ? FT : eltype(grid) 
+
     mod(order, 2) != 0 && throw(ArgumentError("Centered reconstruction scheme is defined only for even orders"))
 
-    N  = Int(order ÷ 2)
-    if N > 1
-        buffer_scheme = Centered(FT; grid, order = order - 2)
-    else
-        buffer_scheme = nothing
-    end
-    return Centered{N, FT}(buffer_scheme)
+    N = Int(order ÷ 2)
+    return Centered{N, FT}()
 end
 
 Base.summary(a::Centered{N}) where N = string("Centered(order=", 2N, ")")
-
-Base.show(io::IO, a::Centered{N, FT}) where {N, FT} =
-    print(io, summary(a), " \n",
-              " Boundary scheme: ", "\n",
-              "    └── ", summary(a.buffer_scheme))
-
-
-Adapt.adapt_structure(to, scheme::Centered{N, FT}) where {N, FT} = Centered{N, FT}(Adapt.adapt(to, scheme.buffer_scheme))
-
-on_architecture(to, scheme::Centered{N, FT}) where {N, FT} = Centered{N, FT}(on_architecture(to, scheme.buffer_scheme))
+Base.show(io::IO, a::Centered{N, FT}) where {N, FT} = summary(a)
 
 # Useful aliases
 Centered(grid, FT::DataType=Float64; kwargs...) = Centered(FT; grid, kwargs...)
 
 const ACAS = AbstractCenteredAdvectionScheme
 
 # left and right biased for Centered reconstruction are just symmetric!
-@inline biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, args...)
-
-# uniform centered reconstruction
-for buffer in advection_buffers, FT in fully_supported_float_types
-    @eval begin
-        @inline symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x, false))
-        @inline symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y, false))
-        @inline symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z, false))
-
-        @inline symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ::Callable, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x,  true))
-        @inline symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ::Callable, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y,  true))
-        @inline symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ::Callable, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z,  true))
+@inline _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, c, args...)
+
+for (side, dir) in zip((:ᶠᵃᵃ, :ᵃᶠᵃ, :ᵃᵃᶠ), (:x, :y, :z))
+    for (F, bool) in zip((:Any, :(Base.Callable)), (false, true))
+        for FT in fully_supported_float_types
+            interp = Symbol(:symmetric_interpolate_, dir, side)
+            @eval begin
+                @inline $interp(i, j, k, grid, ::Centered{1, $FT}, red_order::Int, ψ::$F, args...) = @muladd $(stencil_reconstruction(FT, 1, :symmetric, dir, bool))
+
+                @inline function $interp(i, j, k, grid, ::Centered{2, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1, 
+                           $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                           $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)))
+                end
+
+                @inline function $interp(i, j, k, grid, ::Centered{3, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1,
+                           $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                    ifelse(red_order==2,
+                           $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)),
+                           $(stencil_reconstruction(FT, 3, :symmetric, dir, bool))))
+                end
+
+                @inline function $interp(i, j, k, grid, ::Centered{4, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1,
+                        $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                    ifelse(red_order==2,
+                        $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)),
+                    ifelse(red_order==3,
+                        $(stencil_reconstruction(FT, 3, :symmetric, dir, bool)),
+                        $(stencil_reconstruction(FT, 4, :symmetric, dir, bool)))))
+                end
+
+                @inline function $interp(i, j, k, grid, ::Centered{5, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1,
+                           $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                    ifelse(red_order==2,
+                           $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)),
+                    ifelse(red_order==3,
+                           $(stencil_reconstruction(FT, 3, :symmetric, dir, bool)),
+                    ifelse(red_order==4,
+                           $(stencil_reconstruction(FT, 4, :symmetric, dir, bool)),
+                           $(stencil_reconstruction(FT, 5, :symmetric, dir, bool))))))
+                end
+            end
+        end
     end
 end

src/Advection/flat_advective_fluxes.jl

@@ -34,19 +34,15 @@ FlatGrids = [:XFG, :YFG, :ZFG, :XFG, :YFG, :ZFG]
 for (dir, GridType) in zip((:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ, :xᶜᵃᵃ, :yᵃᶜᵃ, :zᵃᵃᶜ), FlatGrids)
     alt_symm_interp   = Symbol(:_symmetric_interpolate_, dir)
     alt_biased_interp = Symbol(:_biased_interpolate_, dir)
-    @eval begin
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::HOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::LOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::HOADV, ψ::Callable, args...) = @inbounds ψ(i, j, k, grid, args...)
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::LOADV, ψ::Callable, args...) = @inbounds ψ(i, j, k, grid, args...)
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::HOADV, ψ::Callable, ::AS, args...) = @inbounds ψ(i, j, k, grid, args...)
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::LOADV, ψ::Callable, ::AS, args...) = @inbounds ψ(i, j, k, grid, args...)
-
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::HOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::LOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::HOADV, bias, ψ::Callable, args...) = ψ(i, j, k, grid, args...)
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::LOADV, bias, ψ::Callable, args...) = ψ(i, j, k, grid, args...)
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::HOADV, bias, ψ::Callable, ::AS, args...) = ψ(i, j, k, grid, args...)
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::LOADV, bias, ψ::Callable, ::AS, args...) = ψ(i, j, k, grid, args...)
+    for Scheme in [:Centered, :UpwindBiased, :WENO]
+        @eval begin
+            @inline $alt_symm_interp(i, j, k, grid::$GridType, ::$Scheme, ψ, args...) = @inbounds ψ[i, j, k]
+            @inline $alt_symm_interp(i, j, k, grid::$GridType, ::$Scheme, ψ::Callable, args...) = @inbounds ψ(i, j, k, grid, args...)
+            @inline $alt_symm_interp(i, j, k, grid::$GridType, ::$Scheme, ψ::Callable, ::AS, args...) = @inbounds ψ(i, j, k, grid, args...)
+
+            @inline $alt_biased_interp(i, j, k, grid::$GridType, ::$Scheme, ψ, args...) = @inbounds ψ[i, j, k]
+            @inline $alt_biased_interp(i, j, k, grid::$GridType, ::$Scheme, bias, ψ::Callable, args...) = ψ(i, j, k, grid, args...)
+            @inline $alt_biased_interp(i, j, k, grid::$GridType, ::$Scheme, bias, ψ::Callable, ::AS, args...) = ψ(i, j, k, grid, args...)
+        end
     end
-end
+end