From 15b3e553f8912371e75fa1dbca217534edd8b2a2 Mon Sep 17 00:00:00 2001 From: Luca Bertagna Date: Wed, 16 Jul 2025 16:22:19 -0600 Subject: [PATCH 1/3] EAMxx: skip fields without COL tag in horiz remappers operations --- .../share/grid/remap/coarsening_remapper.cpp | 30 ++++++++++++++----- .../grid/remap/horiz_interp_remapper_base.cpp | 23 ++++++++++++++ .../grid/remap/horiz_interp_remapper_base.hpp | 4 +++ .../grid/remap/refining_remapper_p2p.cpp | 29 ++++++++++-------- .../share/tests/coarsening_remapper_tests.cpp | 17 +++++++++-- .../tests/refining_remapper_p2p_tests.cpp | 27 +++++++++++++++++ 6 files changed, 109 insertions(+), 21 deletions(-) diff --git a/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp b/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp index 930d571c7323..73b8c80dce45 100644 --- a/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp +++ b/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp @@ -202,10 +202,15 @@ void CoarseningRemapper::remap_fwd_impl () return (ap.get_last_extent() % SCREAM_PACK_SIZE) == 0; }; - // Loop over each field + // First, perform the local mat-vec. Recall that in these y=Ax products, + // x is the src field, and y is the overlapped tgt field. for (int i=0; i field_col_size (m_num_fields); int sum_fields_col_sizes = 0; for (int i=0; iget_num_local_dofs(); const auto ov_gn = m_ov_coarse_grid->name(); @@ -114,6 +130,13 @@ void HorizInterpRemapperBase::create_ov_fields () for (int i=0; i is evil + std::vector m_needs_remap; + static std::map s_remapper_data; }; diff --git a/components/eamxx/src/share/grid/remap/refining_remapper_p2p.cpp b/components/eamxx/src/share/grid/remap/refining_remapper_p2p.cpp index 3791a56aa110..eebbc168d905 100644 --- a/components/eamxx/src/share/grid/remap/refining_remapper_p2p.cpp +++ b/components/eamxx/src/share/grid/remap/refining_remapper_p2p.cpp @@ -50,6 +50,12 @@ void RefiningRemapperP2P::remap_fwd_impl () // Loop over each field, perform mat-vec constexpr auto COL = ShortFieldTagsNames::COL; for (int i=0; i; - constexpr auto COL = ShortFieldTagsNames::COL; - auto export_pids = m_imp_exp->export_pids(); auto export_lids = m_imp_exp->export_lids(); auto ncols_send = m_imp_exp->num_exports_per_pid(); @@ -176,12 +183,12 @@ void RefiningRemapperP2P::pack_and_send () const int num_exports = export_pids.size(); const int total_col_size = m_fields_col_sizes_scan_sum.back(); for (int ifield=0; ifieldtgt later continue; - } + + const auto& f = m_src_fields[ifield]; + const auto& fl = f.get_header().get_identifier().get_layout(); const auto f_col_sizes_scan_sum = m_fields_col_sizes_scan_sum[ifield]; switch (fl.rank()) { case 1: @@ -316,8 +323,6 @@ void RefiningRemapperP2P::recv_and_unpack () using TeamMember = typename KT::MemberType; using ESU = ekat::ExeSpaceUtils; - constexpr auto COL = ShortFieldTagsNames::COL; - auto import_pids = m_imp_exp->import_pids(); auto import_lids = m_imp_exp->import_lids(); auto ncols_recv = m_imp_exp->num_imports_per_pid(); @@ -326,12 +331,12 @@ void RefiningRemapperP2P::recv_and_unpack () const int num_imports = import_pids.size(); const int total_col_size = m_fields_col_sizes_scan_sum.back(); for (int ifield=0; ifieldtgt later continue; - } + + auto& f = m_ov_fields[ifield]; + const auto& fl = f.get_header().get_identifier().get_layout(); const auto f_col_sizes_scan_sum = m_fields_col_sizes_scan_sum[ifield]; switch (fl.rank()) { case 1: diff --git a/components/eamxx/src/share/tests/coarsening_remapper_tests.cpp b/components/eamxx/src/share/tests/coarsening_remapper_tests.cpp index 799772f3d380..b204e16b89e3 100644 --- a/components/eamxx/src/share/tests/coarsening_remapper_tests.cpp +++ b/components/eamxx/src/share/tests/coarsening_remapper_tests.cpp @@ -113,6 +113,8 @@ Field create_field (const std::string& name, const LayoutType lt, const Abstract const auto& gn = grid.name(); Field f; switch (lt) { + case LayoutType::Scalar1D: + f = Field(FieldIdentifier(name,grid.get_vertical_layout(midpoints),u,gn)); break; case LayoutType::Scalar2D: f = Field(FieldIdentifier(name,grid.get_2d_scalar_layout(),u,gn)); break; case LayoutType::Vector2D: @@ -161,6 +163,10 @@ Field all_gather_field_impl (const Field& f, const ekat::Comm& comm) { constexpr auto COL = ShortFieldTagsNames::COL; const auto& fid = f.get_header().get_identifier(); const auto& fl = fid.get_layout(); + if (not fl.has_tag(COL)) { + // Not partitioned + return f; + } int col_size = fl.clone().strip_dim(COL).size(); auto tags = fl.tags(); auto dims = fl.dims(); @@ -310,6 +316,7 @@ TEST_CASE("coarsening_remap") // Here we will simplify and just remap a simple 2D horizontal field. auto tgt_grid = remap->get_coarse_grid(); + auto src_s1d = create_field("s1d", LayoutType::Scalar1D, *src_grid, true, engine); auto src_s2d = create_field("s2d", LayoutType::Scalar2D, *src_grid, false, engine); auto src_v2d = create_field("v2d", LayoutType::Vector2D, *src_grid, false, engine); auto src_t2d = create_field("t2d", LayoutType::Tensor2D, *src_grid, false, engine); @@ -320,6 +327,7 @@ TEST_CASE("coarsening_remap") auto src_t3d_m = create_field("t3d_m",LayoutType::Tensor3D, *src_grid, true, engine); auto src_t3d_i = create_field("t3d_i",LayoutType::Tensor3D, *src_grid, false, engine); + auto tgt_s1d = create_field("s1d", LayoutType::Scalar1D, *tgt_grid, true); auto tgt_s2d = create_field("s2d", LayoutType::Scalar2D, *tgt_grid, false); auto tgt_v2d = create_field("v2d", LayoutType::Vector2D, *tgt_grid, false); auto tgt_t2d = create_field("t2d", LayoutType::Tensor2D, *tgt_grid, false); @@ -330,8 +338,8 @@ TEST_CASE("coarsening_remap") auto tgt_t3d_m = create_field("t3d_m",LayoutType::Tensor3D, *tgt_grid, true ); auto tgt_t3d_i = create_field("t3d_i",LayoutType::Tensor3D, *tgt_grid, false); - std::vector src_f = {src_s2d,src_v2d,src_t2d,src_s3d_m,src_s3d_i,src_v3d_m,src_v3d_i,src_t3d_m,src_t3d_i}; - std::vector tgt_f = {tgt_s2d,tgt_v2d,tgt_t2d,tgt_s3d_m,tgt_s3d_i,tgt_v3d_m,tgt_v3d_i,tgt_t3d_m,tgt_t3d_i}; + std::vector src_f = {src_s1d,src_s2d,src_v2d,src_t2d,src_s3d_m,src_s3d_i,src_v3d_m,src_v3d_i,src_t3d_m,src_t3d_i}; + std::vector tgt_f = {tgt_s1d,tgt_s2d,tgt_v2d,tgt_t2d,tgt_s3d_m,tgt_s3d_i,tgt_v3d_m,tgt_v3d_i,tgt_t3d_m,tgt_t3d_i}; // -------------------------------------- // // Register fields in the remapper // @@ -373,6 +381,11 @@ TEST_CASE("coarsening_remap") root_print (msg + "\n",comm); bool ok = true; switch (l.type()) { + case LayoutType::Scalar1D: + { + CHECK ( views_are_equal(gsrc,gtgt) ); + ok &= catch_capture.lastAssertionPassed(); + } break; case LayoutType::Scalar2D: { const auto v_src = gsrc.get_view(); diff --git a/components/eamxx/src/share/tests/refining_remapper_p2p_tests.cpp b/components/eamxx/src/share/tests/refining_remapper_p2p_tests.cpp index 3454e7f60a8f..8022f69fd892 100644 --- a/components/eamxx/src/share/tests/refining_remapper_p2p_tests.cpp +++ b/components/eamxx/src/share/tests/refining_remapper_p2p_tests.cpp @@ -25,6 +25,8 @@ Field create_field (const std::string& name, const LayoutType lt, const Abstract const auto ndims = 2; Field f; switch (lt) { + case LayoutType::Scalar1D: + f = Field(FieldIdentifier(name,grid.get_vertical_layout(true),u,gn)); break; case LayoutType::Scalar2D: f = Field(FieldIdentifier(name,grid.get_2d_scalar_layout(),u,gn)); break; case LayoutType::Vector2D: @@ -64,6 +66,10 @@ Field all_gather_field (const Field& f, const ekat::Comm& comm) { constexpr auto COL = ShortFieldTagsNames::COL; const auto& fid = f.get_header().get_identifier(); const auto& fl = fid.get_layout(); + if (not fl.has_tag(COL)) { + // Not partitioned + return f; + } int col_size = fl.clone().strip_dim(COL).size(); auto tags = fl.tags(); auto dims = fl.dims(); @@ -205,17 +211,20 @@ TEST_CASE ("refining_remapper") { auto src_grid = r->get_src_grid(); auto bundle_src = create_field("bundle3d_src",LayoutType::Vector3D,*src_grid,engine); + auto s1d_src = create_field("s1d_src",LayoutType::Scalar1D,*src_grid,engine); auto s2d_src = create_field("s2d_src",LayoutType::Scalar2D,*src_grid,engine); auto v2d_src = create_field("v2d_src",LayoutType::Vector2D,*src_grid,engine); auto s3d_src = create_field("s3d_src",LayoutType::Scalar3D,*src_grid,engine); auto v3d_src = create_field("v3d_src",LayoutType::Vector3D,*src_grid,engine); auto bundle_tgt = create_field("bundle3d_tgt",LayoutType::Vector3D,*tgt_grid); + auto s1d_tgt = create_field("s1d_tgt",LayoutType::Scalar1D,*tgt_grid); auto s2d_tgt = create_field("s2d_tgt",LayoutType::Scalar2D,*tgt_grid); auto v2d_tgt = create_field("v2d_tgt",LayoutType::Vector2D,*tgt_grid); auto s3d_tgt = create_field("s3d_tgt",LayoutType::Scalar3D,*tgt_grid); auto v3d_tgt = create_field("v3d_tgt",LayoutType::Vector3D,*tgt_grid); + r->register_field(s1d_src,s1d_tgt); r->register_field(s2d_src,s2d_tgt); r->register_field(v2d_src,v2d_tgt); r->register_field(s3d_src,s3d_tgt); @@ -229,12 +238,14 @@ TEST_CASE ("refining_remapper") { r->remap_fwd(); // Gather global copies (to make checks easier) and check src/tgt fields + auto gs1d_src = all_gather_field(s1d_src,comm); auto gs2d_src = all_gather_field(s2d_src,comm); auto gv2d_src = all_gather_field(v2d_src,comm); auto gs3d_src = all_gather_field(s3d_src,comm); auto gv3d_src = all_gather_field(v3d_src,comm); auto gbundle_src = all_gather_field(bundle_src,comm); + auto gs1d_tgt = all_gather_field(s1d_tgt,comm); auto gs2d_tgt = all_gather_field(s2d_tgt,comm); auto gv2d_tgt = all_gather_field(v2d_tgt,comm); auto gs3d_tgt = all_gather_field(s3d_tgt,comm); @@ -242,6 +253,22 @@ TEST_CASE ("refining_remapper") { auto gbundle_tgt = all_gather_field(bundle_tgt,comm); Real avg; + // Scalar 1D + { + if (comm.am_i_root()) { + printf(" -> Checking 1d scalars .........\n"); + } + bool ok = true; + gs1d_src.sync_to_host(); + gs1d_tgt.sync_to_host(); + + CHECK (views_are_equal(gs1d_src,gs1d_tgt)); + ok &= catch_capture.lastAssertionPassed(); + if (comm.am_i_root()) { + printf(" -> Checking 1d scalars ......... %s\n",ok ? "PASS" : "FAIL"); + } + } + // Scalar 2D { if (comm.am_i_root()) { From 256b4782f6278f19ecf2510f48ccc046924ba46b Mon Sep 17 00:00:00 2001 From: Luca Bertagna Date: Wed, 16 Jul 2025 16:23:06 -0600 Subject: [PATCH 2/3] EAMxx: simplify handling of geo data in CoarserningRemapper --- .../share/grid/remap/coarsening_remapper.cpp | 21 +------------------ 1 file changed, 1 insertion(+), 20 deletions(-) diff --git a/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp b/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp index 73b8c80dce45..742e5a0b3b51 100644 --- a/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp +++ b/components/eamxx/src/share/grid/remap/coarsening_remapper.cpp @@ -34,26 +34,7 @@ CoarseningRemapper (const grid_ptr_type& src_grid, continue; } const auto& src_data = src_grid->get_geometry_data(name); - const auto& src_data_fid = src_data.get_header().get_identifier(); - const auto& layout = src_data_fid.get_layout(); - if (layout.tags().empty()) { - // This is a scalar field, so won't be coarsened. - // Simply copy it in the tgt grid, but we still need to assign the new grid name. - FieldIdentifier tgt_data_fid(src_data_fid.name(),src_data_fid.get_layout(),src_data_fid.get_units(),m_tgt_grid->name()); - auto tgt_data = m_coarse_grid->create_geometry_data(tgt_data_fid); - tgt_data.deep_copy(src_data); - } else if (layout.tags()[0]!=COL) { - // Not a field to be coarsened (perhaps a vertical coordinate field). - // Simply copy it in the tgt grid, but we still need to assign the new grid name. - FieldIdentifier tgt_data_fid(src_data_fid.name(),src_data_fid.get_layout(),src_data_fid.get_units(),m_tgt_grid->name()); - auto tgt_data = m_coarse_grid->create_geometry_data(tgt_data_fid); - tgt_data.deep_copy(src_data); - } else { - // This field needs to be remapped - auto tgt_data_fid = create_tgt_fid(src_data_fid); - auto tgt_data = m_coarse_grid->create_geometry_data(tgt_data_fid); - register_field(src_data,tgt_data); - } + register_field_from_src(src_data); } registration_ends(); if (get_num_fields()>0) { From 15ea0782830956fde427369cab668740c3c189ba Mon Sep 17 00:00:00 2001 From: Luca Bertagna Date: Wed, 16 Jul 2025 17:14:25 -0600 Subject: [PATCH 3/3] EAMxx: fix horiz remapper cleanup --- .../eamxx/src/share/grid/remap/horiz_interp_remapper_base.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/components/eamxx/src/share/grid/remap/horiz_interp_remapper_base.cpp b/components/eamxx/src/share/grid/remap/horiz_interp_remapper_base.cpp index 88c660a3b67f..29a51fa21368 100644 --- a/components/eamxx/src/share/grid/remap/horiz_interp_remapper_base.cpp +++ b/components/eamxx/src/share/grid/remap/horiz_interp_remapper_base.cpp @@ -278,6 +278,7 @@ void HorizInterpRemapperBase::clean_up () m_src_fields.clear(); m_tgt_fields.clear(); m_ov_fields.clear(); + m_needs_remap.clear(); // Reset the state of the base class m_state = RepoState::Clean;