Unify FM over all grids #6981

tcclevenger · 2025-02-07T01:17:22Z

Store a single FieldManager for all fields over all grids.

This solves the issue in #6789 where having a group as a subset of another, and allocated on different grids, was causing subview indices mismatch between grids.

Also

Remove Derivation type. Now groups contain the union of all fields registered over all grids. I.e., requesting "tracers" on GLL grid without registering any tracers on GLL will give all tracers registered on PG2 grid.
Remove "Preferred" bundling. You either require bundling or not.

Also - remove devivation type - simplify group logic to ensure subfields over grids have same index

tcclevenger

A few comments highlight the important pieces since this is such a large PR. So many of the changes come from every instance of the FM has to be used differently (passing grid names instead of only one grid being stored).

tcclevenger · 2025-02-07T01:21:33Z

components/eamxx/src/control/atmosphere_driver.cpp

-  //   we register a corresponding version of the field on the "target" FM.
-
-  // Helper lambda to reduce code duplication
-  auto process_imported_groups = [&](const std::set<GroupRequest>& group_requests) {


I moved this logic to FieldManager::pre_process_groups().

tcclevenger · 2025-02-07T01:27:12Z

components/eamxx/src/share/field/field.cpp

+Field::clone(const std::string& name) const {
+  return clone(name, get_header().get_identifier().get_grid_name());
 }

 Field
-Field::clone(const std::string& name) const {
+Field::clone(const std::string& name, const std::string& grid_name) const {
  // Create new field
  const auto& my_fid = get_header().get_identifier();
  FieldIdentifier fid(name,my_fid.get_layout(),my_fid.get_units(),
-                      my_fid.get_grid_name(),my_fid.data_type());
+                      grid_name,my_fid.data_type());


I needed to add a grid argument to Field::clone() because, in many cases, we are cloning a grid (giving it a new name), constructing a FM(new_grid), but using a field whose ID gives a grid name from the old grid, and the new FM was complaining that the grid was not in the internal grids manager.

tcclevenger · 2025-02-07T01:28:26Z

components/eamxx/src/share/field/field.cpp

@@ -185,7 +190,7 @@ Field Field::subfield(const std::string& sf_name,
      "Error! Input field must be allocated in order to subview it.\n");

  auto sf_layout = lt.clone();
-  sf_layout.reset_dim(idim, index_end - index_beg);
+  sf_layout.reset_dim(idim, index_end+1 - index_beg);


Kokkos::subview(..., beg, end, ...) has inclusive bounds, but we were treating the subview bound as [bed, end).

Wait, subview has inclusive bounds? That seems odd, given that the end keyword is virtually universally referring to "the first after the last". Also, I think the way we use this method by passing index_end="the first after the last". So when computing the layout dimension, we should not add 1...

Kokkos doc says exclusive which makes the most sense. Somewhere we are going wrong in EAMxx/EKAT, because doing exclusive bounds threw an error in ekat subview about dim being exceeded. I'll investigate.

I believe EKAT is falsely asserting that the end subview index is less than the dimension, when they can be equal (since end can be up to 1 past the last possible index) https://github.com/E3SM-Project/EKAT/blob/0eb00e5d017598aa29b27adc8fb13f51d15c4c1c/src/ekat/kokkos/ekat_subview_utils.hpp#L367C13-L367C44

tcclevenger · 2025-02-07T01:28:41Z

components/eamxx/src/share/field/field_alloc_prop.cpp

@@ -88,7 +88,7 @@ FieldAllocProp FieldAllocProp::subview(const int idim,
  EKAT_REQUIRE_MSG(index_beg < index_end,
                   "Error! Slice indices are invalid (non-increasing).\n");
  EKAT_REQUIRE_MSG(
-      index_beg >= 0 && index_end < m_layout.dim(idim),
+      index_beg >= 0 && index_end <= m_layout.dim(idim),


Same as comment above about Kokkos bounds

In EAMxx we use "std-like" convention of "end=past the last element".

I think this is correct since index_end can go up to one past the last possible index, which is equal to the dim.

components/eamxx/src/share/field/field_layout.hpp

tcclevenger · 2025-02-07T01:36:12Z

components/eamxx/src/share/field/field_manager.cpp

          info.m_subview_idx [*it] = std::distance(cluster_ordered_fields.begin(),it);
        }
        info.m_bundled = true;
-      }
-    }
-  }



Now indices info is independent of grid, guaranteeing each grid agrees.

tcclevenger · 2025-02-07T01:38:19Z

components/eamxx/src/share/field/field_manager.cpp

-  return it==m_fields.end() ? nullptr : it->second;
+FieldManager::get_field_ptr (const std::string& name, const std::string& grid_name) const {
+  auto it = m_fields.at(grid_name).find(name);
+  return it==m_fields.at(grid_name).end() ? nullptr : it->second;
 }

 void FieldManager::pre_process_group_requests () {


Important function for gathering which grids should register which group fields

tcclevenger · 2025-02-07T01:39:04Z

components/eamxx/src/share/field/field_manager.hpp

-  // the parent field. So at registration time, simply keep track of the subfields,
-  // and create them at registration_ends() time, after all other fields.
-  std::map<std::string,FieldRequest> m_subfield_requests;
-
  // The map group_name -> FieldGroupInfo
  group_info_map      m_field_groups;


Doesn't change, group info is the same for all grids.

tcclevenger · 2025-02-07T01:41:29Z

components/eamxx/src/share/grid/abstract_grid.cpp

@@ -115,6 +115,53 @@ AbstractGrid::get_3d_tensor_layout (const bool midpoints, const std::vector<int>
  return get_3d_tensor_layout(midpoints,cmp_dims,names);
 }

+FieldLayout
+AbstractGrid::equivalent_layout (const FieldLayout& template_layout) const


This is how we get the layout for tracer on a new grid.

tcclevenger · 2025-02-07T01:42:57Z

components/eamxx/src/share/grid/grids_manager.cpp

+std::set<std::string> GridsManager::
+get_grid_names () const {
+  std::set<std::string> names;
+  if (m_grids.size()==0) {
+    return names;
+  }
+  for (const auto& g : m_grids) {
+    names.emplace(g.second->name());
+  }
+  return names;
+}


I ended up using this function a lot, but I don't know that I love it. It doesn't take into account aliases. I could instead always loop over the grid manager's grid map and query the names from there.

The method is reasonable. Whether to include aliases depends on how you later intend to use this method. But I think the GM's get_grid method also scans if the input name is a valid grid alias, so maybe it would make sense to return aliases as well...

I use this as a way to look through the FM grids, so I think adding aliases might result in repeated work.

tcclevenger · 2025-02-07T01:51:37Z

Quite a few fails in the CI, odd some didn't trigger on my workstation (rrtmgp standalone for instance)... Leaving Draft for now.

bartgol

I have some preliminary thoughts.

components/eamxx/src/physics/mam/readfiles/tracer_reader_utils.hpp

bartgol · 2025-02-10T18:34:46Z

components/eamxx/src/share/field/field_manager.hpp

-  // Get information about the state of the repo
-  int size () const { return m_fields.size(); }
-  RepoState repository_state () const { return m_repo_state; }
+  // Get number of registered fields on particular grid


Do we even use this method? If not, we can prune it...

components/eamxx/src/share/field/field_manager.hpp

bartgol · 2025-02-10T19:00:11Z

components/eamxx/src/physics/nudging/tests/create_nudging_data.cpp

-  compute_fields(fm1,t0,comm,0);
-  compute_fields(fm2,t0,comm,nlevs_filled);
-  compute_fields(fm3,t0,comm,0);
+  compute_fields(fm1,gn,t0,comm,0);


See my comment on FM header: if we allow querying for a field without passing the grid name (if there's only 1 grid), we can revert these changes (here, as in many other unit tests)...

bartgol · 2025-02-10T19:53:07Z

components/eamxx/src/share/field/field.cpp

@@ -185,7 +190,7 @@ Field Field::subfield(const std::string& sf_name,
      "Error! Input field must be allocated in order to subview it.\n");

  auto sf_layout = lt.clone();
-  sf_layout.reset_dim(idim, index_end - index_beg);
+  sf_layout.reset_dim(idim, index_end+1 - index_beg);


Wait, subview has inclusive bounds? That seems odd, given that the end keyword is virtually universally referring to "the first after the last". Also, I think the way we use this method by passing index_end="the first after the last". So when computing the layout dimension, we should not add 1...

bartgol · 2025-02-10T22:25:51Z

components/eamxx/src/share/field/field_manager.cpp

      }
    }
  }
  ::scream::sort(groups_to_bundle);
-  ::scream::sort(copied_groups);


TODO:

The in-house scream::sort method was just to get around a gcc8+cuda10 issue (as explained in scream_utils.hpp). It is very well possible we no longer need this. We should investigate and possibly remove this fcn if we no longer have the issue (both gcc and cuda are much more recent than those versions).

bartgol · 2025-02-10T22:36:16Z

components/eamxx/src/share/field/field_manager.hpp


  // Constructor(s)
-  explicit FieldManager (const grid_ptr_type& grid);
+  explicit FieldManager (const std::shared_ptr<const AbstractGrid>& grid);
+  explicit FieldManager (const std::shared_ptr<const GridsManager>& grid);


Minor: either always use grids_mgr_type, or remove the typedef and just use the spelled out name (like you did for grid). I favor the latter, btw.

components/eamxx/src/share/field/field_request.hpp

bartgol · 2025-02-10T22:47:34Z

components/eamxx/src/share/field/field_manager.hpp

-  // When registering subfields, we might end up registering the subfield before
-  // the parent field. So at registration time, simply keep track of the subfields,
-  // and create them at registration_ends() time, after all other fields.
-  std::map<std::string,FieldRequest> m_subfield_requests;


Indeed, we were not correctly handling subfield requests. In particular, we were never setting up subfields at the end of registration_ends.

Since nobody has complained, I gather the incomplete "feature" was never used. Hence, I would also purge the reference to this possibility in FieldRequest (the subview-related members and the subview-related constructor).

bartgol · 2025-02-10T22:49:06Z

components/eamxx/src/share/grid/grids_manager.cpp

+std::set<std::string> GridsManager::
+get_grid_names () const {
+  std::set<std::string> names;
+  if (m_grids.size()==0) {
+    return names;
+  }
+  for (const auto& g : m_grids) {
+    names.emplace(g.second->name());
+  }
+  return names;
+}


The method is reasonable. Whether to include aliases depends on how you later intend to use this method. But I think the GM's get_grid method also scans if the input name is a valid grid alias, so maybe it would make sense to return aliases as well...

model_initial_np4

EKAT has an incorrect assert on the extents of subview slices, namely that the end index must be less than the dim extent, when it must be less than or equal to the last extent (since Kokkos uses exclusive end bounds). Will require an EKAT PR and update.

We need the changes to assert on subview bounds

… is size 1

Unify FM over all grids

769af94

Also - remove devivation type - simplify group logic to ensure subfields over grids have same index

tcclevenger marked this pull request as draft February 7, 2025 01:17

tcclevenger commented Feb 7, 2025

View reviewed changes

bartgol reviewed Feb 10, 2025

View reviewed changes

tcclevenger added 6 commits February 12, 2025 13:23

Fix rrtmgp standalone build

08cea74

Create shortcut functions in FM with a single grid

e660aac

model_initial_np4

No need to change name on grid clone

0c1bd13

Remove redundant field_layout::get_dim_names

dbc6088

Remove irrelevent change

4543958

Fix indexing for subfields

afed1e4

EKAT has an incorrect assert on the extents of subview slices, namely that the end index must be less than the dim extent, when it must be less than or equal to the last extent (since Kokkos uses exclusive end bounds). Will require an EKAT PR and update.

tcclevenger force-pushed the tcclevenger/eamxx/single_fm_multiple_grids branch from 952e6bf to cd557bc Compare February 13, 2025 15:52

tcclevenger mentioned this pull request Feb 13, 2025

Fix subview dim checks for multi slice E3SM-Project/EKAT#363

Merged

tcclevenger added 4 commits February 13, 2025 21:19

Update EKAT

8730c53

We need the changes to assert on subview bounds

remove unneeded "Preffered" bundling definition

af427b4

Fix unit test creation of FM, add check in AtmosphereInput that FM GM…

d717734

… is size 1

Allow to change grid name in field clone

1e69db5

tcclevenger force-pushed the tcclevenger/eamxx/single_fm_multiple_grids branch from c859596 to 1e69db5 Compare February 14, 2025 02:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Unify FM over all grids #6981

Unify FM over all grids #6981

tcclevenger commented Feb 7, 2025

tcclevenger left a comment

tcclevenger Feb 7, 2025

tcclevenger Feb 7, 2025

tcclevenger Feb 7, 2025

bartgol Feb 10, 2025

tcclevenger Feb 12, 2025

tcclevenger Feb 12, 2025

tcclevenger Feb 7, 2025

bartgol Feb 10, 2025

tcclevenger Feb 12, 2025

tcclevenger Feb 7, 2025

tcclevenger Feb 7, 2025

tcclevenger Feb 7, 2025

tcclevenger Feb 7, 2025

tcclevenger Feb 7, 2025

bartgol Feb 10, 2025

tcclevenger Feb 12, 2025

tcclevenger commented Feb 7, 2025

bartgol left a comment

bartgol Feb 10, 2025

bartgol Feb 10, 2025

bartgol Feb 10, 2025

bartgol Feb 10, 2025

bartgol Feb 10, 2025

bartgol Feb 10, 2025

bartgol Feb 10, 2025

Unify FM over all grids #6981

Are you sure you want to change the base?

Unify FM over all grids #6981

Conversation

tcclevenger commented Feb 7, 2025

tcclevenger left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tcclevenger commented Feb 7, 2025

bartgol left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment