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.

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

EAMxx: refactor field manipulation interfaces, and add new ones #6970

Open
wants to merge 12 commits into
base: master
Choose a base branch
from

Conversation

bartgol
Copy link
Contributor

@bartgol bartgol commented Feb 4, 2025

Simplify implementation of the scale/scale_inv method, and add new ones for max/min.

[BFB]

First, the current implementation had several methods implemented that could just be expressed as a single call to update, if only update accepted different combine mode values (like update_impl did). The solution was simple: template update on the combine mode, defaulting to the current behavior (i.e., ScaleUpdate). With this modification, scale, scale_impl and also deep_copy (from a Field rhs) could be rewritten as a 1-liner call to update.

Second, the PR adds two more values for CombineMode, namely Max and Min. The goal is to later use Field operations inside scorpio_output.cpp, to avoid manually handling the same kind of stuff in that class, simplifying its implementation.

Finally, allow do do stuff like y.scale(x) if y's data type is, say, double but x's data type is int (but not the other way around!).

@bartgol bartgol added EAMxx PRs focused on capabilities for EAMxx code usability code cleanup labels Feb 4, 2025
@bartgol bartgol requested a review from tcclevenger February 4, 2025 00:43
@bartgol bartgol self-assigned this Feb 4, 2025
Copy link
Contributor

@tcclevenger tcclevenger left a comment

Choose a reason for hiding this comment

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

int->double conversions are going to be nice!

@bartgol
Copy link
Contributor Author

bartgol commented Feb 4, 2025

@tcclevenger I added a few more improvements:

  • Use MDRange instead of Range+modularArithmetics: I chatted with Jim about this (he faced this choice in RRTMGP a lot) as well as kokkos folks. I think the idea is that (with some caveats) MD should be faster if the tiling is right, and especially for small kernels, where the cost of doing index arithmetic may be comparable (or larger) than the actual calculation. I hacked in some timing for a rank-3 field update (200 cols, 2 components, 128 levs) and indeed MD range was faster (almost 2x). It may be that the pb was too small, or other factors, but since it wasn't noticeably slower, I opted for clarity and less code. Maybe one day I'll do that in the field property checks too, and do away with the unflatten_idx utility, which is the only reason why FieldLayout stores a kokkos view for the extents.
  • Treat x and y separately when it comes to contiguity. This causes a duplication of template code instantiations, which is driving up compile time, but it makes sense to take advantage of contiguity if one of the two fields is contiguous, rather than use strided views for both...

The biggest rework compared to earlier is the CombineViewsHelper struct. I needed that on CUDA, since apparently you can't use an if constexpr inside a host-device lambda. So I had to create a functor.

@bartgol bartgol force-pushed the bartgol/eamxx/field-ops-upgrade branch from 09f6f77 to 094230b Compare February 5, 2025 00:03
@bartgol bartgol requested a review from tcclevenger February 5, 2025 20:11
}

template<int N,HostOrDevice HD>
using MDRange = Kokkos::MDRangePolicy<typename get_device<HD>::execution_space,
Copy link
Contributor

Choose a reason for hiding this comment

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

Seems like you are redefining MDRange in the scream::impl::CombineViewsHelper below instead of using this one?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
code cleanup code usability EAMxx PRs focused on capabilities for EAMxx
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants