EAMxx: fix usage of layout extents in field methods

E3SM-Project · Oct 2, 2024 · 756ea4a · 756ea4a
1 parent a01889e
commit 756ea4a
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 7 deletions.
diff --git a/components/eamxx/src/share/field/field_impl.hpp b/components/eamxx/src/share/field/field_impl.hpp
@@ -306,10 +306,19 @@ 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();
 
-  using RangePolicy = typename KokkosTypes<DefaultDevice>::RangePolicy;
+  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,layout.size());
 
-  auto ext = layout.extents();
+  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:
       {
@@ -651,10 +660,18 @@ update_impl (const Field& x, const ST alpha, const ST beta, const ST fill_val)
       " - x layout: " + x_l.to_string() + "\n"
       " - y layout: " + y_l.to_string() + "\n");
 
-  using RangePolicy = typename KokkosTypes<DefaultDevice>::RangePolicy;
+  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());
 
-  auto ext = x_l.extents();
+  using extents_type = typename ekat::KokkosTypes<device_t>::template view_1d<int>;
+  extents_type ext;
+  if constexpr (HD==Device) {
+    ext = x_l.extents();
+  } else {
+    ext = x_l.extents_h();
+  }
 
   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;
 };