Merge pull request #3021 from E3SM-Project/bartgol/eamxx/remove-runti…

…me-allocs Avoid allocating small views at runtime inside Field methods
E3SM-Project · Oct 4, 2024 · 2763997 · 2763997
2 parents a269ef9 + 756ea4a
commit 2763997
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 18 deletions.
diff --git a/components/eamxx/src/share/field/field_impl.hpp b/components/eamxx/src/share/field/field_impl.hpp
@@ -306,32 +306,34 @@ deep_copy_impl (const Field& src) {
   auto src_alloc_props = src.get_header().get_alloc_properties();
   auto tgt_alloc_props =     get_header().get_alloc_properties();
 
-  // If a manual parallel_for is required (b/c of alloc sizes difference),
-  // we need to create extents (rather than just using the one in layout),
-  // since we don't know if we're running on host or device
   using device_t = typename Field::get_device<HD>;
   using exec_space = typename device_t::execution_space;
   using RangePolicy = Kokkos::RangePolicy<exec_space>;
-  using extents_type = typename ekat::KokkosTypes<device_t>::template view_1d<int>;
-  extents_type ext ("",rank);
-  Kokkos::deep_copy(ext,layout.extents());
+
   auto policy = RangePolicy(0,layout.size());
 
+  using extents_type = typename ekat::KokkosTypes<device_t>::template view_1d<int>;
+  extents_type ext;
+  if constexpr (HD==Device) {
+    ext = layout.extents();
+  } else {
+    ext = layout.extents_h();
+  }
   switch (rank) {
     case 1:
       {
         if (src_alloc_props.contiguous() and tgt_alloc_props.contiguous()) {
           auto v     =     get_view<      ST*,HD>();
           auto v_src = src.get_view<const ST*,HD>();
           Kokkos::parallel_for(policy,KOKKOS_LAMBDA(const int idx) {
-              v(idx) = v_src(idx);
-            });
+            v(idx) = v_src(idx);
+          });
         } else {
           auto v     =     get_strided_view<      ST*,HD>();
           auto v_src = src.get_strided_view<const ST*,HD>();
           Kokkos::parallel_for(policy,KOKKOS_LAMBDA(const int idx) {
-              v(idx) = v_src(idx);
-            });
+            v(idx) = v_src(idx);
+          });
         }
       }
       break;
@@ -661,14 +663,16 @@ update_impl (const Field& x, const ST alpha, const ST beta, const ST fill_val)
   using device_t = typename Field::get_device<HD>;
   using exec_space = typename device_t::execution_space;
   using RangePolicy = Kokkos::RangePolicy<exec_space>;
+  auto policy = RangePolicy(0,x_l.size());
 
-  // Need to create extents (rather than just using the one in x_l),
-  // since we don't know if we're running on host or device
   using extents_type = typename ekat::KokkosTypes<device_t>::template view_1d<int>;
-  extents_type ext ("",x_l.rank());
-  Kokkos::deep_copy(ext,x_l.extents());
+  extents_type ext;
+  if constexpr (HD==Device) {
+    ext = x_l.extents();
+  } else {
+    ext = x_l.extents_h();
+  }
 
-  auto policy = RangePolicy(0,x_l.size());
   switch (x_l.rank()) {
     case 0:
       {

diff --git a/components/eamxx/src/share/field/field_layout.cpp b/components/eamxx/src/share/field/field_layout.cpp
@@ -250,9 +250,9 @@ FieldLayout& FieldLayout::reset_dim (const FieldTag t, const int extent, const b
 
 void FieldLayout::set_extents () {
   m_extents = decltype(m_extents)("",m_rank);
-  auto extents_h = Kokkos::create_mirror_view(m_extents);
-  std::copy_n(m_dims.begin(),m_rank,extents_h.data());
-  Kokkos::deep_copy(m_extents,extents_h);
+  m_extents_h = Kokkos::create_mirror_view(m_extents);
+  std::copy_n(m_dims.begin(),m_rank,m_extents_h.data());
+  Kokkos::deep_copy(m_extents,m_extents_h);
 }
 
 void FieldLayout::compute_type () {

diff --git a/components/eamxx/src/share/field/field_layout.hpp b/components/eamxx/src/share/field/field_layout.hpp
@@ -97,6 +97,7 @@ class FieldLayout {
   int dim (const int idim) const;
   const std::vector<int>& dims () const { return m_dims; }
   const extents_type& extents () const { return m_extents; }
+  const extents_type::HostMirror& extents_h () const { return m_extents_h; }
 
   long long  size () const;
 
@@ -156,6 +157,7 @@ class FieldLayout {
   std::vector<std::string>  m_names;
   std::vector<int>          m_dims;
   extents_type              m_extents;
+  extents_type::HostMirror  m_extents_h;
 
   LayoutType                m_type;
 };