Add Enzyme test for differentiating a single column model with CATKEVerticalDiffusivity

ext/OceananigansEnzymeExt.jl

@@ -7,7 +7,6 @@ using Oceananigans.Utils: contiguousrange
 
 using KernelAbstractions
 
-#isdefined(Base, :get_extension) ? (import EnzymeCore) : (import ..EnzymeCore)
 isdefined(Base, :get_extension) ? (import Enzyme) : (import ..Enzyme)
 
 using Enzyme: EnzymeCore
@@ -17,320 +16,14 @@ EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Base.:(==)), ::Oceananigans.Abstr
 EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.validate_grid), x...) = nothing
 EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.metric_function), x...) = nothing
 EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.flatten_reduced_dimensions), x...) = nothing
+EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.prettytime), x...) = nothing
 EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.Grids.total_size), x...) = nothing
 EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.parent_size_and_offset), x...) = nothing
 @inline EnzymeCore.EnzymeRules.inactive_type(v::Type{Oceananigans.Utils.KernelParameters}) = true
 
 @inline batch(::Val{1}, ::Type{T}) where T = T
 @inline batch(::Val{N}, ::Type{T}) where {T, N} = NTuple{N, T}
 
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  func::EnzymeCore.Const{Type{Field}},
-#                                                  ::Type{<:EnzymeCore.Annotation{RT}},
-#                                                  loc::Union{EnzymeCore.Const{<:Tuple},
-#                                                  EnzymeCore.Duplicated{<:Tuple}},
-#                                                  grid::EnzymeCore.Annotation{<: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
-# 
-#     P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing
-#     B = batch(Val(EnzymeCore.EnzymeRules.width(config)), RT)
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, Nothing}(primal, shadow, nothing)
-# end
-# 
-# #####
-# ##### Field
-# #####
-# 
-# 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)
-# 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.Active{<:Oceananigans.Grids.AbstractGrid},
-#                                         T::EnzymeCore.Const{<:DataType}; kw...) where RT
-#     return (nothing, EnzymeCore.make_ero(grid), nothing)
-# end
-
-#####
-##### FunctionField
-#####
-
-# @inline FunctionField(L::Tuple, func, grid) = FunctionField{L[1], L[2], L[3]}(func, grid)
-
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  enzyme_func::Union{EnzymeCore.Const{<:Type{<:FunctionField}}, EnzymeCore.Const{Type{FT2}}},
-#                                                  ::Type{<:EnzymeCore.Annotation{RT}},
-#                                                  function_field_func,
-#                                                  grid;
-#                                                  clock = nothing,
-#                                                  parameters = nothing) where {RT, FT2 <: FunctionField}
-# 
-#     FunctionFieldType = enzyme_func.val     
-# 
-#     primal = if EnzymeCore.EnzymeRules.needs_primal(config)
-#         FunctionFieldType(function_field_func.val, grid.val; clock, parameters)
-#     else
-#         nothing
-#     end
-# 
-#     # function_field_func can be Active, Const (inactive), Duplicated (active but mutable)
-#     function_field_is_active = function_field_func isa Active
-#     # @show function_field_func
-# 
-#     # Support batched differentiation!
-#     config_width = EnzymeCore.EnzymeRules.width(config)
-# 
-#     dactives = if function_field_is_active
-#         if config_width == 1
-#             Ref(EnzymeCore.make_zero(function_field_func.val))
-#         else
-#             ntuple(Val(config_width)) do i
-#                 Base.@_inline_meta
-#                 Ref(EnzymeCore.make_zero(function_field_func.val))
-#             end
-#         end
-#     else
-#         nothing
-#     end
-# 
-#     shadow = if config_width == 1
-#         dfunction_field_func = if function_field_is_active
-#             dactives[]
-#         else
-#             function_field_func.dval
-#         end
-# 
-#         FunctionFieldType(dfunction_field_func, grid.val; clock, parameters)
-#     else
-#   	    ntuple(Val(config_width)) do i
-#   		    Base.@_inline_meta
-# 
-#             dfunction_field_func = if function_field_is_active
-#                 dactives[i][]
-#             else
-#                 function_field_func.dval[i]
-#             end
-# 
-#             FunctionFieldType(dfunction_field_func, grid.val; clock, parameters)
-#   	    end
-#     end
-# 
-#     P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing
-#     B = batch(Val(EnzymeCore.EnzymeRules.width(config)), RT)
-#     D = typeof(dactives)
-# 
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, D}(primal, shadow, dactives)
-# end
-# 
-# function EnzymeCore.EnzymeRules.reverse(config,
-#                                         enzyme_func::Union{EnzymeCore.Const{<:Type{<:FunctionField}}, EnzymeCore.Const{Type{FT2}}},
-#                                         ::RT,
-#                                         tape,
-#                                         function_field_func,
-#                                         grid;
-#                                         clock = nothing,
-#                                         parameters = nothing) where {RT, FT2 <: FunctionField}
-# 
-#     dactives = if function_field_func isa Active
-#         if EnzymeCore.EnzymeRules.width(config) == 1
-#             tape[]
-#         else
-#             ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i
-#                 Base.@_inline_meta
-#                 tape[i][]
-#             end
-#         end
-#     else
-#         nothing
-#     end
-# 
-#     # return (dactives, grid (nothing))
-#     return (dactives, nothing)
-# end
-
-#####
-##### launch!
-#####
-
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  func::EnzymeCore.Const{typeof(Oceananigans.Models.flattened_unique_values)},
-#                                                  ::Type{<:EnzymeCore.Annotation{RT}},
-#                                                  a) where RT
-# 
-#     sprimal = if EnzymeCore.EnzymeRules.needs_primal(config) || EnzymeCore.EnzymeRules.needs_shadow(config)
-#         func.val(a.val)
-#     else
-#         nothing
-#     end
-# 
-#     shadow = if EnzymeCore.EnzymeRules.needs_shadow(config)
-#         if EnzymeCore.EnzymeRules.width(config) == 1
-#             (typeof(a) <: Const) ? EnzymeCore.make_zero(sprimal)::RT : func.val(a.dval)
-#         else
-#             ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i
-#                 Base.@_inline_meta
-#                 (typeof(a) <: Const) ? EnzymeCore.make_zero(sprimal)::RT : func.val(a.dval[i])
-#             end
-#         end
-#     else
-#         nothing
-#     end
-# 
-#     primal = if EnzymeCore.EnzymeRules.needs_primal(config)
-#         sprimal
-#     else
-#         nothing
-#     end
-# 
-#     P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing
-#     B = EnzymeCore.EnzymeRules.needs_primal(config) ? batch(Val(EnzymeCore.EnzymeRules.width(config)), RT) : Nothing
-# 
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, Nothing}(primal, shadow, nothing)
-# end
-# 
-# function EnzymeCore.EnzymeRules.reverse(config,
-#                                         func::EnzymeCore.Const{typeof(Oceananigans.Models.flattened_unique_values)},
-#                                          ::Type{<:EnzymeCore.Annotation{RT}},
-#                                          tape,
-#                                          a) where RT
-# 
-#   return (nothing,)
-# end
-
-#####
-##### launch!
-#####
-
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  func::EnzymeCore.Const{typeof(Oceananigans.Utils.launch!)},
-#                                                  ::Type{EnzymeCore.Const{Nothing}},
-#                                                  arch,
-#                                                  grid,
-#                                                  workspec,
-#                                                  kernel!,
-#                                                  kernel_args::Vararg{Any,N};
-#                                                  include_right_boundaries = false,
-#                                                  reduced_dimensions = (),
-#                                                  location = nothing,
-#                                                  active_cells_map = nothing,
-#                                                  kwargs...) where N
-# 
-# 
-#     workgroup, worksize = Oceananigans.Utils.work_layout(grid.val, workspec.val;
-#                                                          include_right_boundaries,
-#                                                          reduced_dimensions,
-#                                                          location)
-# 
-#     offset = Oceananigans.Utils.offsets(workspec.val)
-# 
-#     if !isnothing(active_cells_map) 
-#         workgroup, worksize = Oceananigans.Utils.active_cells_work_layout(workgroup, worksize, active_cells_map, grid.val) 
-#         offset = nothing
-#     end
-# 
-#     if worksize != 0
-#       
-#       # 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
-# 
-#       @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!)
-# 
-#       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)
-#     else
-#       tape = nothing
-#     end
-# 
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{Nothing, Nothing, Any}(nothing, nothing, tape)
-# end
-# 
-# @inline arg_elem_type(::Type{T}, ::Val{i}) where {T<:Tuple, i} = eltype(T.parameters[i])
-# 
-# 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::Vararg{Any,N};
-#                                                  include_right_boundaries = false,
-#                                                  reduced_dimensions = (),
-#                                                  location = nothing,
-#                                                  active_cells_map = nothing,
-#                                                  kwargs...) where N
-# 
-#   subrets = if tape !== nothing
-#     duploop, subtape = tape
-#     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
-#     res2 = ntuple(Val(N)) do i
-#       Base.@_inline_meta
-#       if kernel_args[i] isa Active
-#         EnzymeCore.make_zero(kernel_args[i].val)
-#       else
-#         nothing
-#       end
-#     end
-#   end
-# 
-#   subrets2 =  ntuple(Val(N)) do i
-#       Base.@_inline_meta
-#       if kernel_args[i] isa Active
-#         subrets[i]::arg_elem_type(typeof(kernel_args), Val(i))
-#       else
-#         nothing
-#       end
-#     end
-# 
-#   return (nothing, nothing, nothing, nothing, subrets2...)
-# 
-# end
-
 #####
 ##### update_model_field_time_series!
 #####

src/BuoyancyModels/seawater_buoyancy.jl

@@ -9,15 +9,15 @@ temperature and salinity are active, or of type `FT` if temperature
 or salinity are constant, respectively.
 """
 struct SeawaterBuoyancy{FT, EOS, T, S} <: AbstractBuoyancyModel{EOS}
-             equation_of_state :: EOS
+    equation_of_state :: EOS
     gravitational_acceleration :: FT
-          constant_temperature :: T
-             constant_salinity :: S
+    constant_temperature :: T
+    constant_salinity :: S
 end
 
 required_tracers(::SeawaterBuoyancy) = (:T, :S)
-required_tracers(::SeawaterBuoyancy{FT, EOS, <:Nothing, <:Number}) where {FT, EOS} = (:T,) # active temperature only
-required_tracers(::SeawaterBuoyancy{FT, EOS, <:Number, <:Nothing}) where {FT, EOS} = (:S,) # active salinity only
+required_tracers(::SeawaterBuoyancy{FT, EOS, <:Nothing, <:Number}) where {FT, EOS} = tuple(:T) # active temperature only
+required_tracers(::SeawaterBuoyancy{FT, EOS, <:Number, <:Nothing}) where {FT, EOS} = tuple(:S) # active salinity only
 
 Base.nameof(::Type{SeawaterBuoyancy}) = "SeawaterBuoyancy"
 Base.summary(b::SeawaterBuoyancy) = string(nameof(typeof(b)), " with g=", prettysummary(b.gravitational_acceleration),

src/Fields/field.jl

@@ -404,8 +404,8 @@ Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...)
 @propagate_inbounds Base.lastindex(f::Field) = lastindex(f.data)
 @propagate_inbounds Base.lastindex(f::Field, dim) = lastindex(f.data, dim)
 
-Base.fill!(f::Field, val) = fill!(parent(f), val)
-Base.parent(f::Field) = parent(f.data)
+@inline Base.fill!(f::Field, val) = fill!(parent(f), val)
+@inline Base.parent(f::Field) = parent(f.data)
 Adapt.adapt_structure(to, f::Field) = Adapt.adapt(to, f.data)
 
 total_size(f::Field) = total_size(f.grid, location(f), f.indices)