Make `swap_spaces` not allocating and more GPU compatible #249

valeriabarra · 2023-02-06T21:40:19Z

Purpose

The purspose of this PR is to improve swap_space so that it is not allocating and to make it GPU compatible.

Closes #248

To-do

~~Remove internal call to ClimaCore parent~~ This cannot be done internally in the function (see error: ERROR: LoadError: syntax: Global method definition around .... needs to be placed at the top level, or use "eval". because of the global method table). So it needs to be worked around wherever it's used.
Make it @inline?

Content

We remove the internal usage of ClimaCore parent by using ClimaCore's Fields.allow_mismatched_diagonalized_spaces() function as a temporary workaround, in the absence of a better ClimaCore's interface

Review checklist

I have:

followed the codebase contribution guide: https://clima.github.io/ClimateMachine.jl/latest/Contributing/
followed the style guide: https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/
followed the documentation policy: https://github.com/CliMA/policies/wiki/Documentation-Policy
checked that this PR does not duplicate an open PR.

In the Content, I have included

relevant unit tests, and integration tests,
appropriate docstrings on all functions, structs, and modules, and included relevant documentation.
I have read and checked the items on the review checklist.

charleskawczynski · 2023-02-16T16:14:08Z

Hi @valeriabarra,

To clarify, parent on ClimaCore Fields is not allocating, parent returns a reference to the underlying array. Ideally, we would still like to avoid using parent so that users do not write code that depends on the underlying data layout, because doing so may make such code incompatible if they were to change the underlying Space.

Calling allow_mismatched_diagonalized_spaces is really a kluge, and we shouldn't use it like this here (overloading the method table at runtime like this is generally not a good idea since it will trigger recompilation). If we can't simply remove the call to parent, then I suggest we revert the parent-related changes from the PR (the SDI says that parent related changes are potential future improvements anyway). We already have an issue in ClimaCore for this problem: CliMA/ClimaCore.jl#1120, so no need to open one. I assume that simply applying these changes didn't work, in which case, let's revert the parent-related changes and apply the allocation-related changes needed.

swap_space! is allocating because of this line:

ClimaCoupler.jl/src/Utilities.jl

Line 73 in b50046f

field_out = zeros(new_space)

, which allocates a field. It turns out that this is actually calling zeros defined on a ClimaCore space, so it may already be (or can be made on the ClimaCore-side) GPU-compatible-- but it's still allocating. I suggest you try to make swap_space! a 3-argument function, and pass field_out as an argument (and then make required changes to use this new 3-argument function). Doing so will allow us to push field_out outside of any loops and avoid repeated allocations.

valeriabarra · 2023-02-16T16:28:19Z

Thank you so much for the thorough suggestion. It was previously suggested to me to use this allow_mismatched_diagonalized_spaces work-around and eliminate the internal call to parent, but yes, I recognize that this was a clumsy work-around and we should proceed the way you suggested, which is much better. I will apply your suggestions as soon as possible.

valeriabarra · 2023-02-21T19:13:46Z

Hey @charleskawczynski I applied the suggestions. I made swap_space a 2-argument function (I really think 3 were not needed, as the space was passed just to create the zero field inside). This way, we just move the allocation from inside the swap_space to outside of it. There is no allocation gain, and only marginal flame graphs improvement. There are not explicit for loops where swap_space is called, so that the output field can be allocated only once out of the loop. But it is implicitly called by the coupling time stepping. Unless we want to pass the field_out all the way from the time-stepping routine, I am not sure how the proposed changes would improve things.

Can you please confirm that what I implemented aligns with what you had in mind? Or clarify otherwise? Thanks!

LenkaNovak

Thanks @valeriabarra . I left a comment on passing a boundary_field like we do with boundary_space right now. Hopefully that'll help with allocations. It would be good to demonstrate on a minimal working example (MWE) the approaches you tried to tackle this problem (as far as I know there are 1. the flag turning off the instance error, 2. creating a field dynamically as the swap_space is being called, and 3. creating a dummy field similar to the boundary_space as mentioned above). Since we are kind of blocked by this long-standing issue of mismatching ClimaCore space instances, we may need to defer this issue (hopefully it will be solved soon, now that it's tracked by CliMA/ClimaCore.jl#1120). That said, it would be good to justify this by summarising your findings using the MWE.

LenkaNovak · 2023-02-21T19:25:47Z

experiments/AMIP/moist_mpi_earth/coupler_utils/flux_calculator.jl

 the temperature of the atmosphere at the lowest level, and the heigh
 of the lowest level.
 """
 function set_ρ_sfc!(ρ_sfc, T_S, integrator)
    ts = integrator.p.ᶜts
    thermo_params = CAP.thermodynamics_params(integrator.p.params)
    ts_int = Spaces.level(ts, 1)
-    parent(ρ_sfc) .= parent(ρ_sfc_at_point.(thermo_params, ts_int, swap_space!(T_S, axes(ts_int))))
+    parent(ρ_sfc) .= parent(ρ_sfc_at_point.(thermo_params, ts_int, swap_space!(zeros(axes(ts_int)), T_S)))


Instead of creating a new field here, can we create it as part of the CouplerSimulation struct? Similar to how we treat boundary_space. @charleskawczynski is this what you meant?

Sure. Here is the MWE gist where I compare all the different approaches. Creating it as part of the CoupledSimulation struct is fine, but we need to be careful to create a temp copy, wherever we have repeated usages, such as in these conservation checks. I can proceed with the proposed changes (Case 4 in the MWE gist) if you agree that's the best option.

It actually seems like Case 2 in the gist is the best option for our purposes, so I will leave the current version of the code. I am going to squash/rebase and you're good to merge. Thanks for the review!

Output of the gist at first run:

[ Info: Case 1: alloc_u_ocn = 19493429 u_ocn = Float64-valued Field: [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0] alloc_u_ice = 0 u_ice = Float64-valued Field: [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0] [ Info: Case 2: alloc_u_ocn = 12800 u_ocn = Float64-valued Field: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0 … 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] alloc_u_ice = 12800 u_ice = Float64-valued Field: [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0] [ Info: Case 3: alloc_u_ocn = 33705 u_ocn = Float64-valued Field: [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0] alloc_u_ice = 1008 u_ice = Float64-valued Field: [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0] [ Info: Case 4: alloc_u_ocn = 465631 u_ocn = Float64-valued Field: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0 … 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] alloc_u_ice = 1008 u_ice = Float64-valued Field: [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0] Float64-valued Field: [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]

valeriabarra · 2023-02-23T00:46:38Z

bors try

LenkaNovak

Thanks, Valeria. After adjusting the alloc limit, feel free to merge.

valeriabarra · 2023-02-23T01:25:00Z

Thanks, Valeria. After adjusting the alloc limit, feel free to merge.

Thank you! Just to clarify, the alloc limit adjustment is not due to this PR but to upstream changes. Thanks, will do.

valeriabarra · 2023-02-23T01:27:54Z

bors r+

249: Make `swap_spaces` not allocating and more GPU compatible r=valeriabarra a=valeriabarra ## Purpose The purspose of this PR is to improve `swap_space` so that it is not allocating and to make it GPU compatible. Closes #248 ## To-do - [x] ~~Remove internal call to ClimaCore `parent`~~ This cannot be done internally in the function (see error: `ERROR: LoadError: syntax: Global method definition around .... needs to be placed at the top level, or use "eval".` because of the global method table). So it needs to be worked around wherever it's used. - [x] Make it ``@inline`?` ## Content We remove the internal usage of ClimaCore `parent` by using ClimaCore's `Fields.allow_mismatched_diagonalized_spaces()` function as a temporary workaround, in the absence of a better ClimaCore's interface Review checklist I have: - followed the codebase contribution guide: https://clima.github.io/ClimateMachine.jl/latest/Contributing/ - followed the style guide: https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/ - followed the documentation policy: https://github.com/CliMA/policies/wiki/Documentation-Policy - checked that this PR does not duplicate an open PR. In the Content, I have included - relevant unit tests, and integration tests, - appropriate docstrings on all functions, structs, and modules, and included relevant documentation. - [x] I have read and checked the items on the review checklist. Co-authored-by: Valeria Barra <[email protected]>

bors · 2023-02-23T01:38:33Z

try

Build failed:

buildkite/climacoupler-ci

bors · 2023-02-23T01:53:08Z

Build failed:

ci 1.8 - macOS-latest

valeriabarra · 2023-02-23T02:10:14Z

bors r+

249: Make `swap_spaces` not allocating and more GPU compatible r=valeriabarra a=valeriabarra ## Purpose The purspose of this PR is to improve `swap_space` so that it is not allocating and to make it GPU compatible. Closes #248 ## To-do - [x] ~~Remove internal call to ClimaCore `parent`~~ This cannot be done internally in the function (see error: `ERROR: LoadError: syntax: Global method definition around .... needs to be placed at the top level, or use "eval".` because of the global method table). So it needs to be worked around wherever it's used. - [x] Make it ``@inline`?` ## Content We remove the internal usage of ClimaCore `parent` by using ClimaCore's `Fields.allow_mismatched_diagonalized_spaces()` function as a temporary workaround, in the absence of a better ClimaCore's interface Review checklist I have: - followed the codebase contribution guide: https://clima.github.io/ClimateMachine.jl/latest/Contributing/ - followed the style guide: https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/ - followed the documentation policy: https://github.com/CliMA/policies/wiki/Documentation-Policy - checked that this PR does not duplicate an open PR. In the Content, I have included - relevant unit tests, and integration tests, - appropriate docstrings on all functions, structs, and modules, and included relevant documentation. - [x] I have read and checked the items on the review checklist. Co-authored-by: Valeria Barra <[email protected]>

bors · 2023-02-23T02:56:02Z

Build failed:

buildkite/climacoupler-ci

valeriabarra · 2023-02-23T03:48:15Z

bors r+

bors · 2023-02-23T05:42:15Z

Build succeeded:

valeriabarra requested a review from simonbyrne February 6, 2023 21:40

valeriabarra changed the title ~~Make swap_spaces not allocating~~ Use swap_spaces as not allocating Feb 6, 2023

valeriabarra self-assigned this Feb 6, 2023

valeriabarra marked this pull request as draft February 6, 2023 21:56

valeriabarra force-pushed the valeria/swap_spaces branch from 8cc5767 to 26aa4a6 Compare February 6, 2023 22:02

valeriabarra mentioned this pull request Feb 6, 2023

Some improvements to interfaces and performance #205

Closed

14 tasks

valeriabarra force-pushed the valeria/swap_spaces branch 5 times, most recently from ab0213d to 528ec61 Compare February 8, 2023 00:39

valeriabarra changed the title ~~Use swap_spaces as not allocating~~ Make swap_spaces not allocating and GPU compatible Feb 9, 2023

valeriabarra force-pushed the valeria/swap_spaces branch 3 times, most recently from 8e488c8 to 213bd2a Compare February 13, 2023 21:50

valeriabarra marked this pull request as ready for review February 16, 2023 01:52

valeriabarra requested review from LenkaNovak and charleskawczynski February 16, 2023 02:38

valeriabarra changed the title ~~Make swap_spaces not allocating and GPU compatible~~ Make swap_spaces not allocating and more GPU compatible Feb 16, 2023

valeriabarra removed the request for review from simonbyrne February 16, 2023 02:38

valeriabarra force-pushed the valeria/swap_spaces branch 2 times, most recently from 02acd0d to b052662 Compare February 21, 2023 00:38

LenkaNovak requested changes Feb 21, 2023

View reviewed changes

valeriabarra force-pushed the valeria/swap_spaces branch from b052662 to 7afd5b0 Compare February 22, 2023 21:48

bors bot added a commit that referenced this pull request Feb 23, 2023

Try #249:

cf7a5aa

LenkaNovak self-requested a review February 23, 2023 01:11

LenkaNovak approved these changes Feb 23, 2023

View reviewed changes

valeriabarra force-pushed the valeria/swap_spaces branch from 7afd5b0 to 41de7df Compare February 23, 2023 01:27

Make swap_spaces not allocating and more GPU friendly

8ac167d

valeriabarra force-pushed the valeria/swap_spaces branch from 41de7df to 8ac167d Compare February 23, 2023 03:48

bors bot merged commit 5b7dfcf into main Feb 23, 2023

bors bot deleted the valeria/swap_spaces branch February 23, 2023 05:42

valeriabarra mentioned this pull request Jun 15, 2023

Performance tips #323

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make `swap_spaces` not allocating and more GPU compatible #249

Make `swap_spaces` not allocating and more GPU compatible #249

valeriabarra commented Feb 6, 2023 •

edited

Loading

charleskawczynski commented Feb 16, 2023

valeriabarra commented Feb 16, 2023

valeriabarra commented Feb 21, 2023

LenkaNovak left a comment •

edited

Loading

LenkaNovak Feb 21, 2023

valeriabarra Feb 22, 2023

valeriabarra Feb 22, 2023

valeriabarra Feb 22, 2023

valeriabarra commented Feb 23, 2023

LenkaNovak left a comment

valeriabarra commented Feb 23, 2023

valeriabarra commented Feb 23, 2023

bors bot commented Feb 23, 2023

bors bot commented Feb 23, 2023

valeriabarra commented Feb 23, 2023

bors bot commented Feb 23, 2023

valeriabarra commented Feb 23, 2023

bors bot commented Feb 23, 2023

Make swap_spaces not allocating and more GPU compatible #249

Make swap_spaces not allocating and more GPU compatible #249

Conversation

valeriabarra commented Feb 6, 2023 • edited Loading

Purpose

To-do

Content

charleskawczynski commented Feb 16, 2023

valeriabarra commented Feb 16, 2023

valeriabarra commented Feb 21, 2023

LenkaNovak left a comment • edited Loading

Choose a reason for hiding this comment

LenkaNovak Feb 21, 2023

Choose a reason for hiding this comment

valeriabarra Feb 22, 2023

Choose a reason for hiding this comment

valeriabarra Feb 22, 2023

Choose a reason for hiding this comment

valeriabarra Feb 22, 2023

Choose a reason for hiding this comment

valeriabarra commented Feb 23, 2023

LenkaNovak left a comment

Choose a reason for hiding this comment

valeriabarra commented Feb 23, 2023

valeriabarra commented Feb 23, 2023

bors bot commented Feb 23, 2023

try

bors bot commented Feb 23, 2023

valeriabarra commented Feb 23, 2023

bors bot commented Feb 23, 2023

valeriabarra commented Feb 23, 2023

bors bot commented Feb 23, 2023

Make `swap_spaces` not allocating and more GPU compatible #249

Make `swap_spaces` not allocating and more GPU compatible #249

valeriabarra commented Feb 6, 2023 •

edited

Loading

LenkaNovak left a comment •

edited

Loading