Improve kernel efficiency of the WENO algorithm

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.3"
 manifest_format = "2.0"
-project_hash = "04d395caf937b0921325a77873167e8baa293a99"
+project_hash = "ebd246712231f728236bc79e507cd0d30883b432"
 
 [[deps.AbstractFFTs]]
 deps = ["LinearAlgebra"]
@@ -403,15 +403,9 @@ version = "1.5.0"
 
 [[deps.JSON3]]
 deps = ["Dates", "Mmap", "Parsers", "PrecompileTools", "StructTypes", "UUIDs"]
-git-tree-sha1 = "eb3edce0ed4fa32f75a0a11217433c31d56bd48b"
+git-tree-sha1 = "95220473901735a0f4df9d1ca5b171b568b2daa3"
 uuid = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
-version = "1.14.0"
-
-    [deps.JSON3.extensions]
-    JSON3ArrowExt = ["ArrowTypes"]
-
-    [deps.JSON3.weakdeps]
-    ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd"
+version = "1.13.2"
 
 [[deps.JuliaNVTXCallbacks_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]

Project.toml

@@ -17,6 +17,7 @@ IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
@@ -32,6 +33,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 

src/Advection/Advection.jl

@@ -40,6 +40,8 @@ import Base: show, summary
 import Oceananigans.Grids: required_halo_size
 import Oceananigans.Architectures: on_architecture
 
+using KernelAbstractions.Extras.LoopInfo: @unroll
+
 abstract type AbstractAdvectionScheme{B, FT} end
 abstract type AbstractCenteredAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end
 abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end
@@ -65,6 +67,9 @@ include("centered_reconstruction.jl")
 include("upwind_biased_reconstruction.jl")
 include("weno_reconstruction.jl")
 include("weno_interpolants.jl")
+include("weno_default_stencils.jl")
+include("weno_function_stencils.jl")
+include("weno_velocity_stencils.jl")
 include("stretched_weno_smoothness.jl")
 include("multi_dimensional_reconstruction.jl")
 include("vector_invariant_upwinding.jl")

src/Advection/tracer_advection_operators.jl

@@ -41,6 +41,7 @@ end
 ##### Tracer advection operator
 #####
 
+
 """
     div_uc(i, j, k, grid, advection, U, c)
 

src/Advection/upwind_biased_advective_fluxes.jl

@@ -100,7 +100,7 @@ end
 ##### Tracer advection operators
 #####
 
-@inline function advective_tracer_flux_x(i, j, k, grid, scheme::UpwindScheme, U, c) 
+@inline function advective_tracer_flux_x(i, j, k, grid, scheme::UpwindScheme, U, c)
 
     @inbounds ũ = U[i, j, k]
     cᴸ =  _left_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, c)

src/Advection/vector_invariant_advection.jl

@@ -108,9 +108,9 @@ Vector Invariant, Dimension-by-dimension reconstruction
  └── smoothness δv²: FunctionStencil f = v_smoothness      
 ```
 """
-function VectorInvariant(; vorticity_scheme = EnstrophyConserving(),
+function VectorInvariant(; vorticity_scheme  = EnstrophyConserving(),
                            vorticity_stencil = VelocityStencil(),
-                           vertical_scheme = EnergyConserving(),
+                           vertical_scheme   = EnergyConserving(),
                            divergence_scheme = vertical_scheme,
                            kinetic_energy_gradient_scheme = divergence_scheme,
                            upwinding  = OnlySelfUpwinding(; cross_scheme = divergence_scheme),
@@ -152,15 +152,33 @@ const VectorInvariantVelocityVerticalUpwinding  = VectorInvariant{<:Any, <:Any,
 Base.summary(a::VectorInvariant)                 = string("Vector Invariant, Dimension-by-dimension reconstruction")
 Base.summary(a::MultiDimensionalVectorInvariant) = string("Vector Invariant, Multidimensional reconstruction")
 
-Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} =
-    print(io, summary(a), " \n",
-              " Vorticity flux scheme: ", "\n",
-              " $(a.vorticity_scheme isa WENO ? "├" : "└")── $(summary(a.vorticity_scheme))",
-              " $(a.vorticity_scheme isa WENO ? "\n └── smoothness ζ: $(a.vorticity_stencil)\n" : "\n")",
-              " Vertical advection / Divergence flux scheme: ", "\n",
-              " $(a.vertical_scheme isa WENO ? "├" : "└")── $(summary(a.vertical_scheme))",
-              "$(a.vertical_scheme isa AbstractUpwindBiasedAdvectionScheme ? 
-              "\n └── upwinding treatment: $(a.upwinding)" : "")")
+function Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} 
+
+    δscheme = a.divergence_scheme
+    vscheme = a.vertical_scheme
+    ζscheme = a.vorticity_scheme
+    kscheme = a.kinetic_energy_gradient_scheme
+
+    msg1 = " Vorticity flux scheme: \n"
+    msg2 = "└── $(summary(ζscheme)) \n"
+    msg3 = " Kinetic energy gradient flux scheme: \n"
+    msg4 = "└── $(summary(kscheme)) \n"
+    msg5 = " Vertical advection scheme: \n"
+    msg6 = "└── $(summary(vscheme)) \n"
+    msg7 = (a.vertical_scheme isa EnergyConserving) ? "" : " Divergence flux scheme: \n"
+    msg8 = isempty(msg7) ? "" : "└── $(summary(a.divergence_scheme)) \n"
+
+    upwinding = (δscheme isa WENO) || (kscheme isa WENO)  || (ζscheme isa WENO)
+
+    msg9  = upwinding ? " WENO smoothness stencils: \n" : ""
+    msg10 = !(ζscheme isa WENO) ? "" :  "└── smoothness ζ: $(a.vorticity_stencil)\n"
+    msg11 = !(δscheme isa WENO) ? "" :  "└── smoothness δx_U : $(a.upwinding.δU_stencil)\n"
+    msg12 = !(δscheme isa WENO) ? "" :  "└── smoothness δy_V : $(a.upwinding.δV_stencil)\n"
+    msg11 = !(kscheme isa WENO) ? "" :  "└── smoothness δx_u² : $(a.upwinding.δu²_stencil)\n"
+    msg12 = !(kscheme isa WENO) ? "" :  "└── smoothness δy_v² : $(a.upwinding.δv²_stencil)\n"
+
+    return print(io, summary(a), "\n", msg1, msg2, msg3, msg4, msg5, msg6, msg7, msg8, msg9, msg10, msg11, msg12)
+end
 
 #####
 ##### Convenience for WENO Vector Invariant
@@ -309,6 +327,10 @@ end
     return 1/Vᶜᶠᶜ(i, j, k, grid) * (Φᵟ + 𝒜ᶻ)
 end
 
+# Fallback for centered advection schemes
+@inline upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme, u, v) = @inbounds u[i, j, k] * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme.divergence_scheme, flux_div_xyᶜᶜᶜ, u, v) 
+@inline upwinded_divergence_flux_Vᶜᶠᶜ(i, j, k, grid, scheme, u, v) = @inbounds v[i, j, k] * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme.divergence_scheme, flux_div_xyᶜᶜᶜ, u, v) 
+
 #####
 ##### Horizontal advection 4 formulations:
 #####  1. Energy conservative         
@@ -346,7 +368,7 @@ end
     return - upwind_biased_product(v̂, ζᴸ, ζᴿ)
 end
 
-@inline function horizontal_advection_V(i, j, k, grid, scheme::VectorInvariantUpwindVorticity, u, v) 
+@inline function horizontal_advection_V(i, j, k, grid, scheme::VectorInvariantUpwindVorticity, u, v)
 
     Sζ = scheme.vorticity_stencil
 

src/Advection/vector_invariant_self_upwinding.jl

@@ -26,7 +26,7 @@
 end
 
 @inline function upwinded_divergence_flux_Vᶜᶠᶜ(i, j, k, grid, scheme::VectorInvariantSelfVerticalUpwinding, u, v)
-    
+
     δV_stencil   = scheme.upwinding.δV_stencil
     cross_scheme = scheme.upwinding.cross_scheme