Method Error: no method matching MixedDuplicated for Oceananigans run with CATKE

[jlk9]

Running this code:

https://github.com/DJ4Earth/Enzymanigans.jl/blob/jlk9/with-simulation/dynamical_core/autodiff_double_gyre.jl

Produces the following method error:

error_catke_0203.txt

@wsmoses

[SouthEndMusic]

I can't open the link so I cannot test whether it works now, is Enzymanigans.jl a private repository? You could try yourself on the branch of my PR:

#2297

[jlk9]

https://github.com/DJ4Earth/Enzymanigans.jl/blob/jlk9/reduce-for-catke-0210/dynamical_core/autodiff_double_gyre.jl

[wsmoses]

using Enzyme

using KernelAbstractions: @kernel, @index, CPU as KCPU

Enzyme.API.looseTypeAnalysis!(true)
Enzyme.API.strictAliasing!(true)

function compute_diffusivities!(J, tracers)

    top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers))

    worksize = (16, 16)
    loop!    = compute_average_surface_buoyancy_flux!(KCPU(), worksize, worksize)

    loop!(J, top_tracer_bcs)

    return nothing
end

@kernel function compute_average_surface_buoyancy_flux!(J, top_tracer_bcs)
    i, j = @index(Global, NTuple)
    J[i, j, 1] = J[i, j, 1]
end

struct CPU     end
struct Bounded end

struct GridHomemade{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY, Arch}
    architecture :: Arch
    Nx :: Int
    Ny :: Int
    Nz :: Int
    Hx :: Int
    Hy :: Int
    Hz :: Int
    Lx :: FT
    Ly :: FT
    Lz :: FT
    # All directions can be either regular (FX, FY, FZ) <: Number
    # or stretched (FX, FY, FZ) <: AbstractVector
    Δλᶠᵃᵃ :: FX
    Δλᶜᵃᵃ :: FX
    λᶠᵃᵃ  :: VX
    λᶜᵃᵃ  :: VX
    Δφᵃᶠᵃ :: FY
    Δφᵃᶜᵃ :: FY
    φᵃᶠᵃ  :: VY
    φᵃᶜᵃ  :: VY
    z     :: CZ
    # Precomputed metrics M <: Nothing means metrics will be computed on the fly
    Δxᶠᶜᵃ :: M
    Δxᶜᶠᵃ :: M
    Δxᶠᶠᵃ :: M
    Δxᶜᶜᵃ :: M
    Δyᶠᶜᵃ :: MY
    Δyᶜᶠᵃ :: MY
    Azᶠᶜᵃ :: M
    Azᶜᶠᵃ :: M
    Azᶠᶠᵃ :: M
    Azᶜᶜᵃ :: M
    # Spherical radius
    radius :: FT

    GridHomemade{TX, TY, TZ}(architecture::Arch,
                                      Nλ, Nφ, Nz,
                                      Hλ, Hφ, Hz,
                                      Lλ :: FT, Lφ :: FT, Lz :: FT,
                                      Δλᶠᵃᵃ :: FX, Δλᶜᵃᵃ :: FX,
                                       λᶠᵃᵃ :: VX,  λᶜᵃᵃ :: VX,
                                      Δφᵃᶠᵃ :: FY, Δφᵃᶜᵃ :: FY,
                                       φᵃᶠᵃ :: VY,  φᵃᶜᵃ :: VY,
                                       z :: CZ,
                                      Δxᶠᶜᵃ :: M,  Δxᶜᶠᵃ :: M,
                                      Δxᶠᶠᵃ :: M,  Δxᶜᶜᵃ :: M,
                                      Δyᶠᶜᵃ :: MY, Δyᶜᶠᵃ :: MY,
                                      Azᶠᶜᵃ :: M,  Azᶜᶠᵃ :: M, 
                                      Azᶠᶠᵃ :: M,  Azᶜᶜᵃ :: M,
                                     radius :: FT) where {Arch, FT, TX, TY, TZ,
                                                           FX, FY, CZ, VX, VY,
                                                           M, MY} =
    new{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY, Arch}(architecture,
                                                         Nλ, Nφ, Nz,
                                                         Hλ, Hφ, Hz,
                                                         Lλ, Lφ, Lz,
                                                         Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ,
                                                         Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ,
                                                         z,
                                                         Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ,
                                                         Δyᶠᶜᵃ, Δyᶜᶠᵃ,
                                                         Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ, radius)
end

# Use autodiff to compute a gradient
grid_homemade = GridHomemade{Bounded, Bounded, Bounded}(CPU(),
                                      16, 16, 2,
                                      2, 2, 2,
                                      60.0, 60.0, 1800.0,
                                      3.75, 3.75,
                                      [0, 1.0],  [0, 1.0],
                                      3.75, 3.75,
                                      [0, 1.0],  [0, 1.0],
                                      ([0,1.0], [0,1.0], [0,1.0]),
                                      [0, 1.0],  [0, 1.0],
                                      [0, 1.0],  [0, 1.0],
                                      1.0, 1.0,
                                      [0, 1.0],  [0, 1.0], 
                                      [0, 1.0],  [0, 1.0],
                                      6.371e6)


struct HomeField{D, B}
    data :: D
    boundary_conditions :: B
end

struct HomeGridField{G}
    grid :: G
end

struct Flux end

mutable struct HomeFieldBoundaryConditions{B, T}
    bottom :: B
    top :: T
end

struct HomeBoundaryCondition{C, T}
    classification :: C
    condition :: T
end

struct HomeDiscreteBoundaryFunction{P, F}
    func :: F
    parameters :: P
end

struct HomeTKETopBoundaryConditionParameters{C, U}
    top_tracer_boundary_conditions :: C
    top_velocity_boundary_conditions :: U
end

@inline top_tke_flux(i, j, grid, clock, fields, parameters, closure, buoyancy) = zero(grid)

# Use autodiff to compute a gradient
zero_flux = HomeBoundaryCondition(Flux(), nothing)
u_top_bc  = HomeBoundaryCondition(Flux(), HomeGridField(grid_homemade))

etop_condition_parameters = HomeTKETopBoundaryConditionParameters((b=zero_flux, e=zero_flux), (u=u_top_bc, v=zero_flux))
etop_condition            = HomeDiscreteBoundaryFunction(top_tke_flux, etop_condition_parameters)

hero_flux = HomeBoundaryCondition(Flux(), nothing)
etop_flux = HomeBoundaryCondition(Flux(), etop_condition)

b_bc = HomeFieldBoundaryConditions(zero_flux, zero_flux)
e_bc = HomeFieldBoundaryConditions(zero_flux, etop_flux)

b = HomeField(0.0, b_bc)
e = HomeField(0.0, e_bc)

tracers = (; b, e)

dtracers = Enzyme.make_zero(tracers)

J  = zeros(Float64, 400)
dJ = zeros(Float64, 400)

dedν = autodiff(set_runtime_activity(Enzyme.Reverse),
                compute_diffusivities!,
                Duplicated(J, dJ),
                Duplicated(tracers, dtracers))

[wsmoses]

using Enzyme

using KernelAbstractions: @kernel, @index, CPU as KCPU

Enzyme.API.printall!(true)

# Enzyme.API.looseTypeAnalysis!(true)
# Enzyme.API.strictAliasing!(true)

@kernel function compute_average_surface_buoyancy_flux!(top_tracer_bcs)
end

function compute_diffusivities!(tracers)

    top_tracer_bcs = NamedTuple(c => tracers[c] for c in propertynames(tracers))

    worksize = (16, 16)
    loop!    = compute_average_surface_buoyancy_flux!(KCPU(), worksize, worksize)

    loop!(top_tracer_bcs)

    return nothing
end

struct GridHomemade
    Lx :: Float64
    v::Vector{Float64}
end

# Use autodiff to compute a gradient
grid_homemade = GridHomemade(0.0, [0.0])

struct HomeGridField{G}
    grid :: G
end

struct Flux end

struct HomeBoundaryCondition{C, T}
    classification :: C
    condition :: T
end

struct HomeDiscreteBoundaryFunction{P, F}
    func :: F
    parameters :: P
end

struct HomeTKETopBoundaryConditionParameters{C, U}
    top_tracer_boundary_conditions :: C
    top_velocity_boundary_conditions :: U
end

@inline top_tke_flux(i, j, grid, clock, fields, parameters, closure, buoyancy) = zero(grid)

# Use autodiff to compute a gradient
zero_flux = HomeBoundaryCondition(Flux(), nothing)
u_top_bc  = HomeBoundaryCondition(Flux(), HomeGridField(grid_homemade))

etop_condition_parameters = HomeTKETopBoundaryConditionParameters((b=zero_flux, e=zero_flux), (u=u_top_bc, v=zero_flux))
etop_condition            = HomeDiscreteBoundaryFunction(top_tke_flux, etop_condition_parameters)

etop_flux = HomeBoundaryCondition(Flux(), etop_condition)

b = zero_flux
e = etop_flux

tracers = (; b, e)

dtracers = Enzyme.make_zero(tracers)

dedν = autodiff(set_runtime_activity(Enzyme.Reverse),
                compute_diffusivities!,
                Duplicated(tracers, dtracers))