Some CATKE performance optimizations

Manifest.toml

@@ -1,8 +1,8 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.9.3"
+julia_version = "1.9.4"
 manifest_format = "2.0"
-project_hash = "72ed8b1b7715053c6d7b675f75dd867b9f153685"
+project_hash = "21eb6b02d2870a916430d805acf3d926ca95d5b2"
 
 [[deps.AbstractFFTs]]
 deps = ["LinearAlgebra"]
@@ -420,12 +420,12 @@ uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 [[deps.LibCURL]]
 deps = ["LibCURL_jll", "MozillaCACerts_jll"]
 uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
-version = "0.6.3"
+version = "0.6.4"
 
 [[deps.LibCURL_jll]]
 deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
 uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
-version = "7.84.0+0"
+version = "8.4.0+0"
 
 [[deps.LibGit2]]
 deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
@@ -434,7 +434,7 @@ uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 [[deps.LibSSH2_jll]]
 deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
 uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
-version = "1.10.2+0"
+version = "1.11.0+1"
 
 [[deps.Libdl]]
 uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
@@ -988,7 +988,7 @@ version = "5.8.0+0"
 [[deps.nghttp2_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
-version = "1.48.0+0"
+version = "1.52.0+1"
 
 [[deps.p7zip_jll]]
 deps = ["Artifacts", "Libdl"]

ext/OceananigansEnzymeCoreExt.jl

@@ -9,7 +9,7 @@ EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.flatten_reduce
 EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.Grids.total_size), x...) = nothing
 
 @inline batch(::Val{1}, ::Type{T}) where T = T
-@inline batch(::Val{N}, ::Type{T}) where {T,N} = NTuple{N,T}
+@inline batch(::Val{N}, ::Type{T}) where {N, T} = NTuple{N, T}
 
 function EnzymeCore.EnzymeRules.augmented_primal(config,
                                                  func::EnzymeCore.Const{Type{Field}},
@@ -18,31 +18,32 @@ function EnzymeCore.EnzymeRules.augmented_primal(config,
                                                  EnzymeCore.Duplicated{<:Tuple}},
                                                  grid::EnzymeCore.Const{<:Oceananigans.Grids.AbstractGrid},
                                                  T::EnzymeCore.Const{<:DataType}; kw...) where RT
-  primal = if EnzymeCore.EnzymeRules.needs_primal(config)
-      func.val(loc.val, grid.val, T.val; kw...)
-  else
-      nothing
-  end
-
-  if haskey(kw, :a)
-    # copy zeroing
-    kw[:data] = copy(kw[:data])
-  end
-
-  shadow = if EnzymeCore.EnzymeRules.width(config) == 1
-  	func.val(loc.val, grid.val, T.val; kw...)
-  else
-  	ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i
-  		Base.@_inline_meta
-	  	func.val(loc.val, grid.val, T.val; kw...)
-  	end
-  end
-
-  return EnzymeCore.EnzymeRules.AugmentedReturn{EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing, batch(Val(EnzymeCore.EnzymeRules.width(config)), RT), Nothing}(primal, shadow, nothing)
+
+    primal = if EnzymeCore.EnzymeRules.needs_primal(config)
+                 func.val(loc.val, grid.val, T.val; kw...)
+             else
+                 nothing
+             end
+
+    if haskey(kw, :a)
+        # copy zeroing
+        kw[:data] = copy(kw[:data])
+    end
+
+    shadow = if EnzymeCore.EnzymeRules.width(config) == 1
+                 func.val(loc.val, grid.val, T.val; kw...)
+             else
+                 ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i
+                     Base.@_inline_meta
+                     func.val(loc.val, grid.val, T.val; kw...)
+                 end
+             end
+
+    return EnzymeCore.EnzymeRules.AugmentedReturn{EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing, batch(Val(EnzymeCore.EnzymeRules.width(config)), RT), Nothing}(primal, shadow, nothing)
 end
 
 function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, func::EnzymeCore.Const{Type{Field}}, ::RT, tape, loc::Union{EnzymeCore.Const{<:Tuple}, EnzymeCore.Duplicated{<:Tuple}}, grid::EnzymeCore.Const{<:Oceananigans.Grids.AbstractGrid}, T::EnzymeCore.Const{<:DataType}; kw...) where RT
-  return (nothing, nothing, nothing)
+    return (nothing, nothing, nothing)
 end
 
 
@@ -69,68 +70,67 @@ function EnzymeCore.EnzymeRules.augmented_primal(config,
     offset = Oceananigans.Utils.offsets(workspec.val)
 
     if !isnothing(only_active_cells) 
-        workgroup, worksize = Oceananigans.Utils.active_cells_work_layout(workgroup, worksize, only_active_cells, grid.val) 
+        workgroup, worksize = Oceananigans.Utils.active_cells_work_layout(workgroup, worksize, only_active_cells, grid.val)
         offset = nothing
     end
 
     if worksize != 0
 
-      # We can only launch offset kernels with Static sizes!!!!
+        # We can only launch offset kernels with Static sizes!!!!
 
-      if isnothing(offset)
-          loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), workgroup, worksize)
-          dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.dval(Oceananigans.Architectures.device(arch.val), workgroup, worksize)
-      else
-          loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset))) 
-          dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset)))
-      end
+        if isnothing(offset)
+            loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), workgroup, worksize)
+            dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.dval(Oceananigans.Architectures.device(arch.val), workgroup, worksize)
+        else
+            loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset))) 
+            dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset)))
+        end
 
-      @debug "Launching kernel $kernel! with worksize $worksize and offsets $offset from $workspec.val"
+        @debug "Launching kernel $kernel! with worksize $worksize and offsets $offset from $workspec.val"
 
+        duploop = (typeof(kernel!) <: EnzymeCore.Const) ? EnzymeCore.Const(loop!) : EnzymeCore.Duplicated(loop!, dloop!)
 
-      duploop = (typeof(kernel!) <: EnzymeCore.Const) ? EnzymeCore.Const(loop!) : EnzymeCore.Duplicated(loop!, dloop!)
+        config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}()
+        subtape = EnzymeCore.EnzymeRules.augmented_primal(config2, duploop, EnzymeCore.Const{Nothing}, kernel_args...).tape
 
-      config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}()
-      subtape = EnzymeCore.EnzymeRules.augmented_primal(config2, duploop, EnzymeCore.Const{Nothing}, kernel_args...).tape
-
-      tape = (duploop, subtape)
+        tape = (duploop, subtape)
     else
-      tape = nothing
+        tape = nothing
     end
 
     return EnzymeCore.EnzymeRules.AugmentedReturn{Nothing, Nothing, Any}(nothing, nothing, tape)
 end
 
 function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1},
-                                                func::EnzymeCore.Const{typeof(Oceananigans.Utils.launch!)},
-                                                 ::Type{EnzymeCore.Const{Nothing}},
-                                                 tape,
-                                                 arch,
-                                                 grid,
-                                                 workspec,
-                                                 kernel!,
-                                                 kernel_args...;
-                                                 include_right_boundaries = false,
-                                                 reduced_dimensions = (),
-                                                 location = nothing,
-                                                 only_active_cells = nothing,
-                                                 kwargs...)
+                                        func::EnzymeCore.Const{typeof(Oceananigans.Utils.launch!)},
+                                         ::Type{EnzymeCore.Const{Nothing}},
+                                         tape,
+                                         arch,
+                                         grid,
+                                         workspec,
+                                         kernel!,
+                                         kernel_args...;
+                                         include_right_boundaries = false,
+                                         reduced_dimensions = (),
+                                         location = nothing,
+                                         only_active_cells = nothing,
+                                         kwargs...)
 
   subrets = if tape !== nothing
-    duploop, subtape = tape
+                duploop, subtape = tape
 
-    config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}()
+                config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}()
 
-    EnzymeCore.EnzymeRules.reverse(config2, duploop, EnzymeCore.Const{Nothing}, subtape, kernel_args...)
-  else
-    ntuple(Val(length(kernel_args))) do _
-      Base.@_inline_meta
-      nothing
-    end
-  end
+                EnzymeCore.EnzymeRules.reverse(config2, duploop, EnzymeCore.Const{Nothing}, subtape, kernel_args...)
+            else
+                ntuple(Val(length(kernel_args))) do _
+                    Base.@_inline_meta
+                    nothing
+                end
+            end
 
   return (nothing, nothing, nothing, nothing, subrets...)
 
 end
 
-end
+end # module

src/Advection/tracer_advection_operators.jl

@@ -1,6 +1,34 @@
 using Oceananigans.Operators: Vᶜᶜᶜ
 using Oceananigans.Fields: ZeroField
 
+struct TracerAdvection{N, FT, A, B, C} <: AbstractAdvectionScheme{N, FT}
+    x :: A
+    y :: B
+    z :: C
+
+    TracerAdvection{N, FT}(x::A, y::B, z::C) where {N, FT, A, B, C} = new{N, FT, A, B, C}(x, y, z)
+end
+
+"""
+    function TracerAdvection(; x, y, z)
+
+Builds a `TracerAdvection` type with reconstruction schemes in `x`, `y`, and `z`.
+"""
+function TracerAdvection(; x, y, z)
+    Nx = required_halo_size(x)
+    Ny = required_halo_size(y)
+    Nz = required_halo_size(z)
+
+    FT = eltype(x)
+
+    return TracerAdvection{max(Nx, Ny, Nz), FT}(x, y, z)
+end
+
+Adapt.adapt_structure(to, scheme::TracerAdvection{N, FT}) where {N, FT} =
+    TracerAdvection{N, FT}(Adapt.adapt(to, scheme.x),
+                           Adapt.adapt(to, scheme.y),
+                           Adapt.adapt(to, scheme.z))
+
 @inline _advective_tracer_flux_x(args...) = advective_tracer_flux_x(args...)
 @inline _advective_tracer_flux_y(args...) = advective_tracer_flux_y(args...)
 @inline _advective_tracer_flux_z(args...) = advective_tracer_flux_z(args...)
@@ -32,3 +60,9 @@ which ends up at the location `ccc`.
                                     δyᵃᶜᵃ(i, j, k, grid, _advective_tracer_flux_y, advection, U.v, c) +
                                     δzᵃᵃᶜ(i, j, k, grid, _advective_tracer_flux_z, advection, U.w, c))
 end
+
+@inline function div_Uc(i, j, k, grid, advection::TracerAdvection, U, c)
+    return 1/Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, _advective_tracer_flux_x, advection.x, U.u, c) +
+                                    δyᵃᶜᵃ(i, j, k, grid, _advective_tracer_flux_y, advection.y, U.v, c) +
+                                    δzᵃᵃᶜ(i, j, k, grid, _advective_tracer_flux_z, advection.z, U.w, c))
+end

src/Advection/vector_invariant_advection.jl

@@ -111,9 +111,9 @@ Vector Invariant, Dimension-by-dimension reconstruction
 function VectorInvariant(; vorticity_scheme = EnstrophyConserving(), 
                            vorticity_stencil = VelocityStencil(),
                            vertical_scheme = EnergyConserving(),
-                           kinetic_energy_gradient_scheme = vertical_scheme,
                            divergence_scheme = vertical_scheme,
-                           upwinding  = OnlySelfUpwinding(; cross_scheme = vertical_scheme),
+                           kinetic_energy_gradient_scheme = divergence_scheme,
+                           upwinding  = OnlySelfUpwinding(; cross_scheme = divergence_scheme),
                            multi_dimensional_stencil = false)
 
     N = required_halo_size(vorticity_scheme)
@@ -132,7 +132,6 @@ end
 const MultiDimensionalVectorInvariant           = VectorInvariant{<:Any, <:Any, true}
 
 #                                                 VectorInvariant{N,     FT,    M,     Z (vorticity scheme)
-const MultiDimensionalVectorInvariant           = VectorInvariant{<:Any, <:Any, true}
 const VectorInvariantEnergyConserving           = VectorInvariant{<:Any, <:Any, <:Any, <:EnergyConserving}
 const VectorInvariantEnstrophyConserving        = VectorInvariant{<:Any, <:Any, <:Any, <:EnstrophyConserving}
 const VectorInvariantUpwindVorticity            = VectorInvariant{<:Any, <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme}
@@ -145,10 +144,10 @@ const VectorInvariantKEGradientEnergyConserving = VectorInvariant{<:Any, <:Any,
 const VectorInvariantKineticEnergyUpwinding     = VectorInvariant{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme}
 
 
-#                                                 VectorInvariant{N,     FT,    M,     Z,     ZS,    V,                                     K,     D,     U (upwinding)
-const VectorInvariantCrossVerticalUpwinding     = VectorInvariant{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme, <:Any, <:Any, <:CrossAndSelfUpwinding}
-const VectorInvariantSelfVerticalUpwinding      = VectorInvariant{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme, <:Any, <:Any, <:OnlySelfUpwinding}
-const VectorInvariantVelocityVerticalUpwinding  = VectorInvariant{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme, <:Any, <:Any, <:VelocityUpwinding}
+#                                                 VectorInvariant{N,     FT,    M,     Z,     ZS,     V,     K,     D,                                     U (upwinding)
+const VectorInvariantCrossVerticalUpwinding     = VectorInvariant{<:Any, <:Any, <:Any, <:Any, <:Any,  <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme, <:CrossAndSelfUpwinding}
+const VectorInvariantSelfVerticalUpwinding      = VectorInvariant{<:Any, <:Any, <:Any, <:Any, <:Any,  <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme, <:OnlySelfUpwinding}
+const VectorInvariantVelocityVerticalUpwinding  = VectorInvariant{<:Any, <:Any, <:Any, <:Any, <:Any,  <:Any, <:Any, <:AbstractUpwindBiasedAdvectionScheme, <:VelocityUpwinding}
 
 Base.summary(a::VectorInvariant)                 = string("Vector Invariant, Dimension-by-dimension reconstruction")
 Base.summary(a::MultiDimensionalVectorInvariant) = string("Vector Invariant, Multidimensional reconstruction")
@@ -167,10 +166,7 @@ Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} =
 ##### Convenience for WENO Vector Invariant
 #####
 
-#                           VectorInvariant{N,     FT,    M,     Z (vorticity scheme),      ZS,    V (vertical scheme),      K (kinetic energy gradient scheme)
-const WENOVectorInvariant = VectorInvariant{<:Any, <:Any, <:Any, <:WENO, <:Any, <:WENO, <:WENO}
-
-nothing_to_default(user_value, default) = isnothing(user_value) ? default : user_value
+nothing_to_default(user_value; default) = isnothing(user_value) ? default : user_value
 
 """
     function WENOVectorInvariant(; upwinding = nothing,
@@ -189,23 +185,23 @@ function WENOVectorInvariant(; upwinding = nothing,
                                weno_kw...)
 
     if isnothing(order) # apply global defaults
-        vorticity_order               = nothing_to_default(vorticity_order, default=9)
-        vertical_order                = nothing_to_default(vertical_order, default=5)
-        divergence_order              = nothing_to_default(divergence_order, default=5)
-        kinetic_energy_gradient_order = nothing_to_default(kinetic_energy_gradient_order, default=5)
+        vorticity_order               = nothing_to_default(vorticity_order,  default = 9)
+        vertical_order                = nothing_to_default(vertical_order,   default = 5)
+        divergence_order              = nothing_to_default(divergence_order, default = 5)
+        kinetic_energy_gradient_order = nothing_to_default(kinetic_energy_gradient_order, default = 5)
     else # apply user supplied `order` unless overridden by more specific value
-        vorticity_order               = nothing_to_default(vorticity_order, default=order)
-        vertical_order                = nothing_to_default(vertical_order, default=order)
-        divergence_order              = nothing_to_default(divergence_order, default=order)
-        kinetic_energy_gradient_order = nothing_to_default(kinetic_energy_gradient_order, default=order)
+        vorticity_order               = nothing_to_default(vorticity_order,  default = order)
+        vertical_order                = nothing_to_default(vertical_order,   default = order)
+        divergence_order              = nothing_to_default(divergence_order, default = order)
+        kinetic_energy_gradient_order = nothing_to_default(kinetic_energy_gradient_order, default = order)
     end
 
-    vorticity_scheme               = WENO(; order=vorticity_order, weno_kw...)
-    vertical_scheme                = WENO(; order=vertical_order, weno_kw...)
-    kinetic_energy_gradient_scheme = WENO(; order=kinetic_energy_gradient_order, weno_kw...)
-    divergence_scheme              = WENO(; order=divergence_order, weno_kw...)
+    vorticity_scheme               = WENO(; order = vorticity_order, weno_kw...)
+    vertical_scheme                = WENO(; order = vertical_order, weno_kw...)
+    kinetic_energy_gradient_scheme = WENO(; order = kinetic_energy_gradient_order, weno_kw...)
+    divergence_scheme              = WENO(; order = divergence_order, weno_kw...)
 
-    default_upwinding = OnlySelfUpwinding(cross_scheme=divergence_scheme)
+    default_upwinding = OnlySelfUpwinding(cross_scheme = divergence_scheme)
     upwinding = nothing_to_default(upwinding; default = default_upwinding)
 
     schemes = (vorticity_scheme, vertical_scheme, kinetic_energy_gradient_scheme, divergence_scheme)

src/Architectures.jl

@@ -112,4 +112,11 @@ end
 @inline unsafe_free!(a::CuArray) = CUDA.unsafe_free!(a)
 @inline unsafe_free!(a)          = nothing
 
+# Convert arguments to GPU-compatible types
+
+@inline convert_args(::CPU, args) = args
+@inline convert_args(::GPU, args) = CUDA.cudaconvert(args)
+@inline convert_args(::GPU, args::Tuple) = map(CUDA.cudaconvert, args)
+
+
 end # module

src/BoundaryConditions/fill_halo_regions.jl

@@ -50,7 +50,7 @@ function fill_halo_regions!(c::MaybeTupledData, boundary_conditions, indices, lo
 
     arch = architecture(grid)
 
-    fill_halos!, bcs  = permute_boundary_conditions(boundary_conditions)
+    fill_halos!, bcs = permute_boundary_conditions(boundary_conditions)
     number_of_tasks = length(fill_halos!)
 
     # Fill halo in the three permuted directions (1, 2, and 3), making sure dependencies are fulfilled