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

Use new ReactantState architecture in super_simple_simulation.jl #12

Closed
wants to merge 3 commits into from
Closed

Use new ReactantState architecture in super_simple_simulation.jl #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4639fa1
Use new `ReactantState` architecture in `super_simple_simulation.jl`
giordano Feb 8, 2025
8b2c516
Apply suggestions from code review
giordano Feb 8, 2025
613e956
Add clock to the Reactant model
giordano Feb 8, 2025
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 Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,5 +4,6 @@ CFTime = "179af706-886a-5703-950a-314cd64e0468"
ClimaOcean = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
OrthogonalSphericalShellGrids = "c2be9673-fb75-4747-82dc-aa2bb9f4aed0"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
41 changes: 29 additions & 12 deletions oceananigans-dynamical-core/super_simple_simulation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,53 +2,70 @@ using Oceananigans
using Reactant
using Libdl
using Reactant_jll
using Random

Reactant.Ops.DEBUG_MODE[] = true
ENV["JULIA_DEBUG"] = "Reactant_jll"
@show Reactant_jll.cuDriverGetVersion(dlopen("libcuda.so"))
# ENV["JULIA_DEBUG"] = "Reactant_jll"
# @show Reactant_jll.cuDriverGetVersion(dlopen("libcuda.so"))

arch = GPU() # CPU() to run on CPU
# Try to automatically guess whether we have a GPU available
arch = isempty(find_library(["libcuda.so.1", "libcuda.so"])) ? CPU() : GPU()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💪

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you also want to set the Reactant default to cpu vs gpu based on this check?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think Reactant does it automatically already.

r_arch = Oceananigans.ReactantState()
Nx, Ny, Nz = (360, 120, 100) # number of cells

grid = LatitudeLongitudeGrid(arch, size=(Nx, Ny, Nz), halo=(7, 7, 7),
longitude=(0, 360), latitude=(-60, 60), z=(-1000, 0))
r_grid = LatitudeLongitudeGrid(r_arch, size=(Nx, Ny, Nz), halo=(7, 7, 7),
longitude=(0, 360), latitude=(-60, 60), z=(-1000, 0))

FT = Float64
t = ConcreteRNumber(zero(FT))
iter = ConcreteRNumber(0)
stage = ConcreteRNumber(0)
last_Δt = ConcreteRNumber(zero(FT))
last_stage_Δt = ConcreteRNumber(zero(FT))
r_clock = Clock(; time=t, iteration=iter, stage, last_Δt, last_stage_Δt)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@glwagner to be clear, with CliMA/Oceananigans.jl#4096 how would I modify this? So that I can test it quickly locally

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

with that PR you won't need this anymore because it will happen by default. (But you can still override the default if you want)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CliMA/Oceananigans.jl#4096 basically just wraps this code into a default so that we retain existing behavior but also extend when using a ReactantState architecture


# One of the implest configurations we might consider:
model = HydrostaticFreeSurfaceModel(; grid, momentum_advection=WENO())
r_model = HydrostaticFreeSurfaceModel(; grid=r_grid, clock=r_clock, momentum_advection=WENO())

@assert model.free_surface isa SplitExplicitFreeSurface
@assert r_model.free_surface isa SplitExplicitFreeSurface

Random.seed!(123)
uᵢ(x, y, z) = randn()
set!(model, u=uᵢ, v=uᵢ)

# First form a Reactant model
r_model = Reactant.to_rarray(model)
Random.seed!(123)
set!(r_model, u=uᵢ, v=uᵢ)

# What we normally do:
@info "--> Running simulation with stock Oceananigans"
simulation = Simulation(model, Δt=60, stop_iteration=2)
run!(simulation)

#using CUDA
#CUDA.@device_code dir="cudajl" run!(simulation)

# What we want to do with Reactant:
@info "--> Running simulation with Reactant"
r_simulation = Simulation(r_model, Δt=60, stop_iteration=2)
pop!(r_simulation.callbacks, :nan_checker)
# @show @code_hlo optimize=:before_kernel run!(r_simulation)

r_run! = @compile sync = true run!(r_simulation)
r_run!(r_simulation)

@info "--> Re-running simulation with stock Oceananigans (with profiler)"
Reactant.with_profiler("./notrace4/") do
run!(simulation)
end

@info "--> Re-running simulation with Reactant (with profiler)"
Reactant.with_profiler("./retrace4/") do
r_run!(r_simulation)
end

using BenchmarkTools

@btime run!(simulation)
@btime r_run!(r_simulation)

@info "--> Benchmark stock Oceananigans"
@btime run!($simulation)
@info "--> Benchmark Reactant"
@btime r_run!($r_simulation)