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

Jump to bottom

Add test of tracer advection-diffusion for Enzyme #3480

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

Add test of tracer advection-diffusion for Enzyme #3480

wants to merge 41 commits into from

Conversation

jlk9
Copy link
Collaborator

@jlk9 jlk9 commented Feb 14, 2024

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

@wsmoses wsmoses mentioned this pull request Feb 14, 2024
Merged
@navidcy navidcy added testing 🧪 Tests get priority in case of emergency evacuation extensions 🧬 labels Feb 14, 2024
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
navidcy
navidcy reviewed Feb 14, 2024
View reviewed changes
test/test_enzyme.jl Outdated Show resolved Hide resolved
@navidcy
Copy link
Collaborator

navidcy commented Feb 15, 2024

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

@navidcy navidcy self-requested a review February 15, 2024 12:50
navidcy
navidcy requested changes Feb 15, 2024
View reviewed changes
Copy link
Collaborator

@navidcy navidcy left a comment

Choose a reason for hiding this comment

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

added a few suggestions, most of them to improve clarity

but most importantly the tests seem to fail. I don't know how to interpret the Enzyme error... :(

@glwagner glwagner changed the title Added test of tracer advection-diffusion for Enzyme Add test of tracer advection-diffusion for Enzyme Feb 16, 2024
jlk9 and others added 2 commits February 16, 2024 10:36
@jlk9 @navidcy
Update test/test_enzyme.jl
1ce8978 
Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
@jlk9 @navidcy
Update test/test_enzyme.jl
bc988a0 
Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
@wsmoses
Copy link
Collaborator

wsmoses commented Feb 16, 2024

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

@glwagner
Copy link
Member

Just to summarize where we are here:

  1. The new test fails.
  2. We also expect that changes on (0.90.7) Remove argument splatting in hydrostatic free surface tendency kernel entry functions #3477 will cause the test to fail. A possible solution to the changes that (0.90.7) Remove argument splatting in hydrostatic free surface tendency kernel entry functions #3477 will incur is to replace the slurped locations with args::Varargs{N, T} @wsmoses
  3. We need to merge (0.90.7) Remove argument splatting in hydrostatic free surface tendency kernel entry functions #3477 to restore normal CPU performance. Otherwise, Oceananigans can't be used even merely to demonstrate capabilities on the CPU.

@glwagner
Copy link
Member

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.

jlk9 and others added 5 commits February 16, 2024 10:41
@jlk9 @navidcy
Update test/test_enzyme.jl
2e061d8 
Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
@jlk9 @navidcy
Update test/test_enzyme.jl
31cc3c2 
Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
@jlk9
Added elaboration to comment on sum
aea14b8
@jlk9 @navidcy
Update test/test_enzyme.jl
4331906 
Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
@jlk9
Added include statement and changed @test assertion with typo
e308701
@glwagner
Copy link
Member

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

glwagner
glwagner reviewed Feb 27, 2024
View reviewed changes
src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl Outdated
@@ -283,7 +283,7 @@ end
##### Tendency calculators for an explicit free surface
#####

""" Calculate the right-hand-side of the free surface displacement (``η``) equation. """
""" Calculate the right-hand-side of the free surface displacement equation. (``η``) """
Copy link
Member

@glwagner glwagner Feb 27, 2024 •
edited

Choose a reason for hiding this comment

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

The first one is right (η is the symbol we use for the free surface displacement, and is sometimes referred to as the "η-equation").

@wsmoses wsmoses mentioned this pull request May 11, 2024
Open
jlk9 and others added 4 commits May 20, 2024 15:04
@jlk9
Merge branch 'main' into jlk9/enzyme-test
2661e49
@jlk9 @wsmoses @glwagner
Jlk9/enzyme constructor any test (#3603)
53686c2 
* Adding test to check for Oceananigans' broadcasted arrays (via KA) breaking with AD

* Update test_enzyme.jl

* Update test_enzyme.jl

* Update test_enzyme.jl

* Update test_enzyme.jl

* Update Project.toml

* Up KernelAbstractions

* Update test_enzyme.jl

* Made a few problematic objects inactive in OceananigansEnzymeExt

* Accidentally pushed some bug reductions, undoing those

---------

Co-authored-by: William Moses <gh@wsmoses.com>
Co-authored-by: Gregory Wagner <wagner.greg@gmail.com>
@jlk9
fixed test setup error
663622f
@jlk9
Disabled test that is currently bugged
23037d7
@jlk9
Copy link
Collaborator Author

jlk9 commented May 23, 2024

@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.

@jlk9
Copy link
Collaborator Author

jlk9 commented May 24, 2024

@wsmoses @glwagner It passed all tests!

@wsmoses
Copy link
Collaborator

wsmoses commented May 24, 2024

@navidcy can we get a final approval/merge?

@wsmoses
Copy link
Collaborator

wsmoses commented May 29, 2024

@navidcy @glwagner bumping

glwagner
glwagner reviewed May 29, 2024
View reviewed changes
src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
@@ -154,7 +154,7 @@ function compute_free_surface_tendency!(grid, model, kernel_parameters)

launch!(arch, grid, kernel_parameters,
compute_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η,
grid, args)
grid, args...)
Copy link
Member

@glwagner glwagner May 29, 2024 •
edited

Choose a reason for hiding this comment

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

Does this produce a catastrophic slow down on CPU or no?

glwagner
glwagner reviewed May 29, 2024
View reviewed changes
src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
@@ -122,7 +122,7 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
c_tendency,
grid,
active_cells_map,
args;
args...;
Copy link
Member

Choose a reason for hiding this comment

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

This went in previously to fix an O(100)x slow down on CPU, so we should test whether the slow down is present in this PR

@glwagner
Copy link
Member

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)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@wsmoses wsmoses wsmoses left review comments

@glwagner glwagner glwagner left review comments

@navidcy navidcy navidcy requested changes

Requested changes must be addressed to merge this pull request.

Assignees
No one assigned
Labels
extensions 🧬 testing 🧪 Tests get priority in case of emergency evacuation
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
4 participants
@jlk9 @navidcy @wsmoses @glwagner