Enzyme.jl
Enzyme.jl copied to clipboard
EnzymeAD
Method Error: no method matching MixedDuplicated for Oceananigans run with CATKE
Running this code:
https://github.com/DJ4Earth/Enzymanigans.jl/blob/jlk9/with-simulation/dynamical_core/autodiff_double_gyre.jl
Produces the following method error:
@wsmoses
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:
https://github.com/EnzymeAD/Enzyme.jl/pull/2297
https://github.com/DJ4Earth/Enzymanigans.jl/blob/jlk9/reduce-for-catke-0210/dynamical_core/autodiff_double_gyre.jl
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))
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))
Current MWE:
using Enzyme
using KernelAbstractions: @kernel, 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)
J[1, 1, 1] = J[1, 1, 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))
Yeah it makes me wonder if https://github.com/EnzymeAD/Enzyme.jl/blob/43654f904beeeed866305bd5b9e282818a098e50/src/rules/jitrules.jl#L187 should wrap the shadow in a Ref
but I don't know that code well enough.
No that would result in a correctness issue (as the rrror here is that someone provided a value that should have been a ref but was not a ref). The issue is figuring out who/where