Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Run Enzyme tests on the GPU on Sverdrup #3911

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Run Enzyme tests on the GPU on Sverdrup #3911

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a749d4b
remove test
simone-silvestri Nov 8, 2024
65b0b59
Get Enzyme tests running on GPU
glwagner Nov 8, 2024
fc25521
Merge branch 'main' into ss/remove-enzyme-GPU-tests
glwagner Nov 8, 2024
f7331ac
Update test_enzyme.jl
simone-silvestri Nov 9, 2024
e457022
Merge branch 'main' into ss/remove-enzyme-GPU-tests
simone-silvestri Dec 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .buildkite/pipeline.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -732,3 +732,4 @@ steps:
agents:
queue: Oceananigans
architecture: CPU

291 changes: 141 additions & 150 deletions test/test_enzyme.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,27 +67,6 @@ function stable_diffusion!(model, amplitude, diffusivity)
return sum_c²::Float64
end

@testset "Enzyme unit tests" begin
arch = CPU()
FT = Float64

N = 100
topo = (Periodic, Flat, Flat)
grid = RectilinearGrid(arch, FT, topology=topo, size=N, halo=2, x=(-1, 1), y=(-1, 1), z=(-1, 1))
fwd, rev = Enzyme.autodiff_thunk(ReverseSplitWithPrimal, Const{typeof(f)}, Duplicated, typeof(Const(grid)))
tape, primal, shadowp = fwd(Const(f), Const(grid))

@show tape primal shadowp

shadow = if shadowp isa Base.RefValue
shadowp[]
else
shadowp
end

@test size(primal) == size(shadow)
end

function set_initial_condition_via_launch!(model_tracer, amplitude)
# Set initial condition
amplitude = Ref(amplitude)
Expand All @@ -105,148 +84,160 @@ function set_initial_condition_via_launch!(model_tracer, amplitude)
return nothing
end

@testset "Enzyme + Oceananigans Initialization Broadcast Kernel" begin

Nx = Ny = 64
Nz = 8
@testset "Enzyme unit tests" begin
for arch in archs
FT = Float64

x = y = (-π, π)
z = (-0.5, 0.5)
topology = (Periodic, Periodic, Bounded)
N = 100
topo = (Periodic, Flat, Flat)
grid = RectilinearGrid(arch, FT, topology=topo, size=N, halo=2, x=(-1, 1), y=(-1, 1), z=(-1, 1))
fwd, rev = Enzyme.autodiff_thunk(ReverseSplitWithPrimal, Const{typeof(f)}, Duplicated, typeof(Const(grid)))
tape, primal, shadowp = fwd(Const(f), Const(grid))

grid = RectilinearGrid(size=(Nx, Ny, Nz); x, y, z, topology)
model = HydrostaticFreeSurfaceModel(; grid, tracers=:c)
model_tracer = model.tracers.c
@show tape primal shadowp

amplitude = 1.0
amplitude = Ref(amplitude)
cᵢ(x, y, z) = amplitude[]
temp = Base.broadcasted(Base.identity, FunctionField((Center, Center, Center), cᵢ, model_tracer.grid))
shadow = if shadowp isa Base.RefValue
shadowp[]
else
shadowp
end

temp = convert(Base.Broadcast.Broadcasted{Nothing}, temp)
grid = model_tracer.grid
arch = architecture(model_tracer)

if arch == CPU()
param = Oceananigans.Utils.KernelParameters(size(model_tracer),
map(Oceananigans.Fields.offset_index, model_tracer.indices))
dmodel_tracer = Enzyme.make_zero(model_tracer)

# Test the individual kernel launch
autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
Oceananigans.Utils.launch!,
Const(arch),
Const(grid),
Const(param),
Const(Oceananigans.Fields._broadcast_kernel!),
Duplicated(model_tracer, dmodel_tracer),
Const(temp))

# Test out differentiation of the broadcast infrastructure
autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
set_initial_condition_via_launch!,
Duplicated(model_tracer, dmodel_tracer),
Active(1.0))

# Test differentiation of the high-level set interface
dmodel = Enzyme.make_zero(model)
autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
set_initial_condition!,
Duplicated(model, dmodel),
Active(1.0))
@test size(primal) == size(shadow)
end
end

@testset "Enzyme + Oceananigans Initialization Broadcast Kernel" begin
for arch in archs
Nx = Ny = 4
Nz = 2

@testset "Enzyme for advection and diffusion with various boundary conditions" begin
Nx = Ny = 64
Nz = 8

Lx = Ly = L = 2π
Lz = 1

x = y = (-L/2, L/2)
z = (-Lz/2, Lz/2)
topology = (Periodic, Periodic, Bounded)

grid = RectilinearGrid(size=(Nx, Ny, Nz); x, y, z, topology)
diffusion = VerticalScalarDiffusivity(κ=0.1)

u = XFaceField(grid)
v = YFaceField(grid)

U = 1
u₀(x, y, z) = - U * cos(x + L/8) * sin(y) * (z + L/2)
v₀(x, y, z) = + U * sin(x + L/8) * cos(y) * (z + L/2)
x = y = (-π, π)
z = (-0.5, 0.5)
topology = (Periodic, Periodic, Bounded)

set!(u, u₀)
set!(v, v₀)
fill_halo_regions!(u)
fill_halo_regions!(v)
grid = RectilinearGrid(size=(Nx, Ny, Nz); x, y, z, topology)
model = HydrostaticFreeSurfaceModel(; grid, tracers=:c)
model_tracer = model.tracers.c

@inline function tracer_flux(i, j, grid, clock, model_fields, p)
c₀ = p.surface_tracer_concentration
u★ = p.piston_velocity
return - u★ * (c₀ - model_fields.c[i, j, p.level])
amplitude = 1.0
amplitude = Ref(amplitude)
cᵢ(x, y, z) = amplitude[]
temp = Base.broadcasted(Base.identity, FunctionField((Center, Center, Center), cᵢ, model_tracer.grid))

temp = convert(Base.Broadcast.Broadcasted{Nothing}, temp)
grid = model_tracer.grid
arch = architecture(model_tracer)

if arch == CPU()
param = Oceananigans.Utils.KernelParameters(size(model_tracer),
map(Oceananigans.Fields.offset_index, model_tracer.indices))
dmodel_tracer = Enzyme.make_zero(model_tracer)

# Test the individual kernel launch
autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
Oceananigans.Utils.launch!,
Const(arch),
Const(grid),
Const(param),
Const(Oceananigans.Fields._broadcast_kernel!),
Duplicated(model_tracer, dmodel_tracer),
Const(temp))

# Test out differentiation of the broadcast infrastructure
autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
set_initial_condition_via_launch!,
Duplicated(model_tracer, dmodel_tracer),
Active(1.0))

# Test differentiation of the high-level set interface
dmodel = Enzyme.make_zero(model)
autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
set_initial_condition!,
Duplicated(model, dmodel),
Active(1.0))
end
end
end

parameters = (surface_tracer_concentration = 1,
piston_velocity = 0.1,
level = Nz)

top_c_bc = FluxBoundaryCondition(tracer_flux; discrete_form=true, parameters)
c_bcs = FieldBoundaryConditions(top=top_c_bc)

# TODO:
# 1. Make the velocity fields evolve
# 2. Add surface fluxes
# 3. Do a problem where we invert for the tracer fluxes (maybe with CATKE)

model_no_bc = HydrostaticFreeSurfaceModel(; grid,
tracer_advection = WENO(),
tracers = :c,
velocities = PrescribedVelocityFields(; u, v),
closure = diffusion)

model_bc = HydrostaticFreeSurfaceModel(; grid,
tracer_advection = WENO(),
tracers = :c,
velocities = PrescribedVelocityFields(; u, v),
boundary_conditions = (; c=c_bcs),
closure = diffusion)

models = [model_no_bc, model_bc]

@show "Advection-diffusion results, first without then with flux BC"

for i in 1:2
# Compute derivative by hand
κ₁, κ₂ = 0.9, 1.1
c²₁ = stable_diffusion!(models[i], 1, κ₁)
c²₂ = stable_diffusion!(models[i], 1, κ₂)
dc²_dκ_fd = (c²₂ - c²₁) / (κ₂ - κ₁)
@inline function tracer_flux(i, j, grid, clock, model_fields, p)
c₀ = p.surface_tracer_concentration
u★ = p.piston_velocity
return - u★ * (c₀ - model_fields.c[i, j, p.level])
end

# Now for real
amplitude = 1.0
κ = 1.0
dmodel = Enzyme.make_zero(models[i])
set_diffusivity!(dmodel, 0)

dc²_dκ = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
stable_diffusion!,
Duplicated(models[i], dmodel),
Const(amplitude),
Active(κ))

@info """ \n
Advection-diffusion:
Enzyme computed $dc²_dκ
Finite differences computed $dc²_dκ_fd
"""

tol = 0.01
rel_error = abs(dc²_dκ[1][3] - dc²_dκ_fd) / abs(dc²_dκ_fd)
@test rel_error < tol
@testset "Enzyme for advection and diffusion with various boundary conditions" begin
for arch in archs
@info " Running Enzyme advection-diffusion tests on $arch..."
Nx = Ny = 4
Nz = 2
x = y = (-π, π)
z = (-1/2, 1/2)
topology = (Periodic, Periodic, Bounded)
grid = RectilinearGrid(arch, size=(Nx, Ny, Nz); x, y, z, topology)
diffusion = VerticalScalarDiffusivity(κ=0.1)

u = XFaceField(grid)
v = YFaceField(grid)

U = 1
L = grid.Lx
u₀(x, y, z) = - U * cos(x + L/8) * sin(y) * (z + L/2)
v₀(x, y, z) = + U * sin(x + L/8) * cos(y) * (z + L/2)

set!(u, u₀)
set!(v, v₀)
fill_halo_regions!(u)
fill_halo_regions!(v)
parameters = (surface_tracer_concentration=1, piston_velocity=0.1, level=Nz)
top_c_bc = FluxBoundaryCondition(tracer_flux; discrete_form=true, parameters)
c_bcs = FieldBoundaryConditions(top=top_c_bc)

model_no_bc = HydrostaticFreeSurfaceModel(; grid,
tracer_advection = WENO(),
tracers = :c,
velocities = PrescribedVelocityFields(; u, v),
closure = diffusion)

model_bc = HydrostaticFreeSurfaceModel(; grid,
tracer_advection = WENO(),
tracers = :c,
velocities = PrescribedVelocityFields(; u, v),
boundary_conditions = (; c=c_bcs),
closure = diffusion)

models = [model_no_bc, model_bc]

for model in models
bc = model.tracers.c.top.boundary_condition.condition
@info " Running advection-diffusion tests on $arch with $bc top tracer boundary condition..."
# Compute derivative by hand
κ₁, κ₂ = 0.9, 1.1
c²₁ = stable_diffusion!(model, 1, κ₁)
c²₂ = stable_diffusion!(model, 1, κ₂)
dc²_dκ_fd = (c²₂ - c²₁) / (κ₂ - κ₁)

# Now for real
amplitude = 1.0
κ = 1.0
dmodel = Enzyme.make_zero(model)
set_diffusivity!(dmodel, 0)

dc²_dκ = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
stable_diffusion!,
Duplicated(model, dmodel),
Const(amplitude),
Active(κ))

@info """ \n
Advection-diffusion:
Enzyme computed $dc²_dκ
Finite differences computed $dc²_dκ_fd
"""

tol = 0.01
rel_error = abs(dc²_dκ[1][3] - dc²_dκ_fd) / abs(dc²_dκ_fd)
@test rel_error < tol
end
end
end