Oceananigans.jl Add test of tracer advection-diffusion for Enzyme

Added an additional test to the enzyme test set for differentiating a HydrostaticFreeSurfaceModel with tracer advection and diffusion.

Feb 14 '24 21:02 jlk9

btw, the new test seem to fail? https://buildkite.com/clima/oceananigans/builds/14443#018da99b-3a02-41a5-9663-3522fb224813/16-409

Feb 15 '24 12:02 navidcy

Hm haven't seen that error before. In any case this PR should resolve: https://github.com/EnzymeAD/Enzyme.jl/pull/1297

Feb 16 '24 16:02 wsmoses

Just to summarize where we are here:

The new test fails.
We also expect that changes on #3477 will cause the test to fail. A possible solution to the changes that #3477 will incur is to replace the slurped locations with args::Varargs{N, T} @wsmoses
We need to merge #3477 to restore normal CPU performance. Otherwise, Oceananigans can't be used even merely to demonstrate capabilities on the CPU.

Feb 16 '24 16:02 glwagner

Hm haven't seen that error before. In any case this PR should resolve: EnzymeAD/Enzyme.jl#1297

Okay, let's see if we can test whether that fixes it before merging in changes from #3477, just so we can test one thing at a time.

Feb 16 '24 16:02 glwagner

@michel2323 @vchuravy this looks like an error in KA indexing, which I recall you guys resolved previously? [and @glwagner there was that weird offset issue before iirc]


Enzyme on advection and diffusion: Error During Test at /var/lib/buildkite-agent/builds/tartarus-3/clima/oceananigans/test/test_enzyme.jl:93
--
  | Got exception outside of a @test
  | BoundsError: attempt to access 16×16×1 Array{Tuple{Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}}, 3} at index [1, 1, 2]
  | Stacktrace:
  | [1] setindex!
  | @ ./array.jl:971 [inlined]
  | [2] setindex!
  | @ ./multidimensional.jl:670 [inlined]
  | [3] aug_fwd
  | @ /storage5/buildkite-agent/.julia-14477/packages/KernelAbstractions/WoCk1/ext/EnzymeExt.jl:20 [inlined]
  | [4] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{KernelAbstractions.CPU, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, KernelAbstractions.NDIteration.StaticSize{(64, 64, 8)}, typeof(EnzymeExt.aug_fwd)}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{3, KernelAbstractions.NDIteration.StaticSize{(4, 4, 8)}, KernelAbstractions.NDIteration.StaticSize{(16, 16, 1)}, Nothing, Nothing}, args::Tuple{typeof(Oceananigans.Models.HydrostaticFreeSurfaceModels.cpu_compute_hydrostatic_free_surface_Gc!), Val{(false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, true)}, Array{Tuple{Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}}, 3}, Duplicated{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}, Const{RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}}, Const{Nothing}, Const{Val{1}}, Const{Val{:c}}, Const{WENO{3, Float64, Nothing, Nothing, Nothing, true, Nothing, WENO{2, Float64, Nothing, Nothing, Nothing, true, Nothing, UpwindBiased{1, Float64, Nothing, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{2, Float64, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}}}, Active{VerticalScalarDiffusivity{ExplicitTimeDiscretization, Float64, NamedTuple{(:c,), Tuple{Float64}}, Float64}}, Const{BoundaryCondition{Flux, Nothing}}, Const{Nothing}, Const{Nothing}, Duplicated{PrescribedVelocityFields{Field{Face, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Face, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Oceananigans.Fields.ZeroField{Int64, 3}, Nothing}}, Const{Nothing}, Duplicated{NamedTuple{(:c,), Tuple{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}}}, Const{NamedTuple{(:c,), Tuple{BoundaryCondition{Flux, Nothing}}}}, Const{Nothing}, Const{NamedTuple{(), Tuple{}}}, Const{typeof(Oceananigans.Forcings.zeroforcing)}, Duplicated{Clock{Float64}}}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck)
  | @ KernelAbstractions /storage5/buildkite-agent/.julia-14477/packages/KernelAbstractions/WoCk1/src/cpu.jl:115
  | [5] __run(obj::KernelAbstractions.Kernel{KernelAbstractions.CPU, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, KernelAbstractions.NDIteration.StaticSize{(64, 64, 8)}, typeof(EnzymeExt.aug_fwd)}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{3, KernelAbstractions.NDIteration.StaticSize{(4, 4, 8)}, KernelAbstractions.NDIteration.StaticSize{(16, 16, 1)}, Nothing, Nothing}, args::Tuple{typeof(Oceananigans.Models.HydrostaticFreeSurfaceModels.cpu_compute_hydrostatic_free_surface_Gc!), Val{(false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, true)}, Array{Tuple{Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{UInt64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{UInt64, 0}}, 3}, Duplicated{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}, Const{RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}}, Const{Nothing}, Const{Val{1}}, Const{Val{:c}}, Const{WENO{3, Float64, Nothing, Nothing, Nothing, true, Nothing, WENO{2, Float64, Nothing, Nothing, Nothing, true, Nothing, UpwindBiased{1, Float64, Nothing, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{2, Float64, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}}}, Active{VerticalScalarDiffusivity{ExplicitTimeDiscretization, Float64, NamedTuple{(:c,), Tuple{Float64}}, Float64}}, Const{BoundaryCondition{Flux, Nothing}}, Const{Nothing}, Const{Nothing}, Duplicated{PrescribedVelocityFields{Field{Face, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Face, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Oceananigans.Fields.ZeroField{Int64, 3}, Nothing}}, Const{Nothing}, Duplicated{NamedTuple{(:c,), Tuple{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}}}, Const{NamedTuple{(:c,), Tuple{BoundaryCondition{Flux, Nothing}}}}, Const{Nothing}, Const{NamedTuple{(), Tuple{}}}, Const{typeof(Oceananigans.Forcings.zeroforcing)}, Duplicated{Clock{Float64}}}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck, static_threads::Bool)
  | @ KernelAbstractions /storage5/buildkite-agent/.julia-14477/packages/KernelAbstractions/WoCk1/src/cpu.jl:82
  | [6] (::KernelAbstractions.Kernel{KernelAbstractions.CPU, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, KernelAbstractions.NDIteration.StaticSize{(64, 64, 8)}, typeof(EnzymeExt.aug_fwd)})(::Function, ::Vararg{Any}; ndrange::Nothing, workgroupsize::Nothing)
  | @ KernelAbstractions /storage5/buildkite-agent/.julia-14477/packages/KernelAbstractions/WoCk1/src/cpu.jl:44
  | [7] augmented_primal(::EnzymeCore.EnzymeRules.ConfigWidth{1, false, false, (false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, true)}, ::Const{KernelAbstractions.Kernel{KernelAbstractions.CPU, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, KernelAbstractions.NDIteration.StaticSize{(64, 64, 8)}, typeof(Oceananigans.Models.HydrostaticFreeSurfaceModels.cpu_compute_hydrostatic_free_surface_Gc!)}}, ::Type{Const{Nothing}}, ::Duplicated{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}, ::Vararg{Any}; ndrange::Nothing, workgroupsize::Nothing)
  | @ EnzymeExt /storage5/buildkite-agent/.julia-14477/packages/KernelAbstractions/WoCk1/ext/EnzymeExt.jl:62
  | [8] augmented_primal(::EnzymeCore.EnzymeRules.ConfigWidth{1, false, false, (false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, true)}, ::Const{KernelAbstractions.Kernel{KernelAbstractions.CPU, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, KernelAbstractions.NDIteration.StaticSize{(64, 64, 8)}, typeof(Oceananigans.Models.HydrostaticFreeSurfaceModels.cpu_compute_hydrostatic_free_surface_Gc!)}}, ::Type{Const{Nothing}}, ::Duplicated{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArray{Float64, 3, Array{Float64, 3}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}, BoundaryCondition{Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}, ::Const{RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, CPU}}, ::Const{Nothing}, ::Const{Val{1}}, ::Const{Val{:c}}, ::Const{WENO{3, Float64, Nothing, Nothing, Nothing, true, Nothing, WENO{2, Float64, Nothing, Nothing, Nothing, true, Nothing, UpwindBiased{1, Float64, Nothing, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{2, Float64, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}}}, ::Vararg{Any})
  | @ EnzymeExt /storage5/buildkite-agent/.julia-14477/packages/KernelAbstractions/WoCk1/ext/EnzymeExt.jl:39

Feb 20 '24 02:02 wsmoses

So I switched from KernelAbstraction v0.9.16 (what I have on my dev Oceananigans) to v0.9.10 (on the CI test) and replicated the error. Updating to 9.16 fixes this.

Feb 20 '24 03:02 jlk9

we can update KernelAbstractions to v0.9.16!

Feb 20 '24 06:02 navidcy

Enzyme tests pass now: https://buildkite.com/clima/oceananigans/builds/14478#018dc549-6a59-40cd-a6f1-5e4d2cb7ced4

🎉

Feb 20 '24 07:02 navidcy

@navidcy That's great and what I intended, Billy and I think we have a workaround for the splatting, just wanted to get all the checks to pass first to eliminate any other hurdles.

Feb 20 '24 18:02 jlk9

We'll have to do some benchmarking on the CPU too

Feb 26 '24 23:02 glwagner

@glwagner I think we're getting some non-deterministic CI errors again... tests are failing for this commit that didn't on the previous one.

May 23 '24 20:05 jlk9

@wsmoses @glwagner It passed all tests!

May 24 '24 23:05 jlk9

@navidcy can we get a final approval/merge?

May 24 '24 23:05 wsmoses

@navidcy @glwagner bumping

May 29 '24 13:05 wsmoses

Let's do a simple test of CPU performance to make sure we won't have to revert this soon given the change to splatting (which was implemented to solve a 100x slow down a few months ago)

May 29 '24 15:05 glwagner

Billy mentioned it in some other comments but while we do splat args for some of the function calls, the function definitions use Varargs instead. This should avoid the catastrophic slowdown we saw with splatting earlier, but I agree that it should be tested. Do you have any good CPU performance tests set up @glwagner ?

Jun 05 '24 14:06 jlk9

Billy mentioned it in some other comments but while we do splat args for some of the function calls, the function definitions use Varargs instead. This should avoid the catastrophic slowdown we saw with splatting earlier, but I agree that it should be tested. Do you have any good CPU performance tests set up @glwagner ?

Any simple test will do, for example one of the examples

Jun 05 '24 15:06 glwagner