multi-sliced subfield/subviews working

Author: mjschmidt271

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

@@ -128,8 +128,8 @@ Field Field::subfield(const std::string& sf_name,
   EKAT_REQUIRE_MSG(
       is_allocated(),
       "Error! Input field must be allocated in order to subview it.\n");
-  EKAT_REQUIRE_MSG(idim == 0 || idim == 1,
-                   "Error! Subview dimension index must be either 0 or 1.\n");
+  // EKAT_REQUIRE_MSG(idim == 0 || idim == 1,
+  //                  "Error! Subview dimension index must be either 0 or 1.\n");
 
   auto sf_layout =
       lt.clone_with_different_extent(idim, index_end - index_beg + 1);

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

@@ -136,9 +136,11 @@ class Field {
   get_strided_view_type<DT,HD>
   get_strided_view () const;
 
-  template<typename DT, HostOrDevice HD = Device>
-  Kokkos::View<Real****, Kokkos::LayoutStride>
-  get_strided_view (bool special) const;
+  // this is the same as above but for a multi-sliced view
+  // which is to say, also a strided view
+  template<typename DT, int N, HostOrDevice HD = Device>
+  auto get_multi_sliced_view () const
+    -> get_strided_view_type<data_nd_t<DT, N>, HD>;
 
   // These two getters are convenience function for commonly accessed metadata.
   // The same info can be extracted from the metadata stored in the FieldHeader
@@ -245,10 +247,9 @@ class Field {
   Field subfield (const std::string& sf_name, const int idim,
                        const int index, const bool dynamic = false) const;
   Field subfield (const int idim, const int k, const bool dynamic = false) const;
-  // get subfield to extract multiple slices in a continuous range of indices
+  // subfield fxn to extract multiple slices in a continuous range of indices
   // e.g., (in matlab syntax) subf = f.subfield(:, 1:3, :)
   // but NOT subf = f.subfield(:, [1, 3, 4], :)
-  // NOTE: use std::pair(index_beg, index_end) or a kokkos::pair()?
   Field subfield (const std::string& sf_name, const ekat::units::Units& sf_units,
                   const int idim, const int index_beg, const int index_end,
                   const bool dynamic = false) const;

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

@@ -77,22 +77,20 @@ subview (const int idim, const int k, const bool dynamic) const {
   return props;
 }
 
-FieldAllocProp FieldAllocProp::subview(const int idim, const int k_beg,
+FieldAllocProp FieldAllocProp::subview(const int idim,
+                                       const int k_beg,
                                        const int k_end,
                                        const bool dynamic) const {
   EKAT_REQUIRE_MSG(
       is_committed(),
       "Error! Subview requires alloc properties to be committed.\n");
-  EKAT_REQUIRE_MSG(
-      idim == 0 || idim == 1,
-      "Error! Subviewing is only allowed along first or second dimension.\n");
-  EKAT_REQUIRE_MSG(idim < m_layout.rank(),
+  EKAT_REQUIRE_MSG(idim <= m_layout.rank(),
                    "Error! Dimension index out of bounds.\n");
-  EKAT_REQUIRE_MSG(k_beg <= k_end,
+  EKAT_REQUIRE_MSG(k_beg < k_end,
                    "Error! Slice indices are invalid (non-increasing).\n");
   EKAT_REQUIRE_MSG(
       k_beg >= 0 && k_end < m_layout.dim(idim),
-      "Error! Index range along the dimension is out of bounds.\n");
+      "Error! Slice index range along the idim dimension is out of bounds.\n");
 
   // Set new layout basic stuff
   FieldAllocProp props(m_scalar_type_size);

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

@@ -106,7 +106,7 @@ class FieldAllocProp {
 
   FieldAllocProp& operator= (const FieldAllocProp&);
 
-  // Return allocation props of an unmanaged subivew of this field
+  // Return allocation props of an unmanaged subview of this field
   // at entry k along dimension idim. If dynamic = true, then it will be
   // possible to change the entry index (k) at runtime.
   FieldAllocProp subview (const int idim, const int k, const bool dynamic) const;

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

@@ -99,10 +99,6 @@ auto Field::get_view () const
   // Make sure input field is allocated
   EKAT_REQUIRE_MSG(is_allocated(),
       "Error! Cannot extract a field's view before allocation happens.\n");
-  // FIXME: add check
-  // EKAT_REQUIRE_MSG(true /*is_multiSlice_subview()*/,
-  //                  "Error! Multi-sliced subfield is incompatible--must employ "
-  //                  "get_strided_view().\n")
 
   EKAT_REQUIRE_MSG (not m_is_read_only || std::is_const<DstValueType>::value,
       "Error! Cannot get a view to non-const data if the field is read-only.\n");
@@ -207,26 +203,26 @@ auto Field::get_strided_view () const
   return DstView(get_ND_view<HD,DstValueType,1>());
 }
 
-template<typename DT, HostOrDevice HD>
-Kokkos::View<Real****, Kokkos::LayoutStride> Field::get_strided_view (bool special) const
+// NOTE: multi-slicing a view is only supported for strided view return type
+template<typename DT, int N, HostOrDevice HD>
+auto Field::get_multi_sliced_view () const
+ -> get_strided_view_type<data_nd_t<DT, N>, HD>
 {
   // The destination view type on correct mem space
-  using DstView = get_strided_view_type<DT,HD>;
+  using DstView = get_strided_view_type<data_nd_t<DT, N>, HD>;
   // The dst value types
   using DstValueType = typename DstView::traits::value_type;
-  // We only allow to reshape to a view of the correct rank
-  constexpr int DstRank = DstView::rank;
-  constexpr int DstRankDynamic = DstView::rank_dynamic;
 
   // Get src details
   const auto& alloc_prop = m_header->get_alloc_properties();
   const auto& fl = m_header->get_identifier().get_layout();
 
   // Checks
-  // EKAT_REQUIRE_MSG (DstRank==1 && fl.rank()==1,
-  //     "Error! Strided view only available for rank-1 fields.\n");
-  // EKAT_REQUIRE_MSG (DstRankDynamic==1,
+  // TODO: decide whether a dynamic-rank view is ok
+  // EKAT_REQUIRE_MSG (DstRankDynamic == 1,
   //     "Error! Strided view not allowed with compile-time dimensions.\n");
+  EKAT_REQUIRE_MSG (N == fl.rank(),
+      "Error! Input Rank must be equal to parent view's rank for multi-sliced subview.\n");
   EKAT_REQUIRE_MSG(is_allocated(),
       "Error! Cannot extract a field's view before allocation happens.\n");
   EKAT_REQUIRE_MSG (not m_is_read_only || std::is_const<DstValueType>::value,
@@ -236,51 +232,23 @@ Kokkos::View<Real****, Kokkos::LayoutStride> Field::get_strided_view (bool speci
 
   // Check if this field is a subview of another field
   const auto parent = m_header->get_parent().lock();
-  if (parent!=nullptr) {
-    // Parent field has correct layout to reinterpret the view into N+1-dim view
-    // So create the parent field on the fly, use it to get the N+1-dim view, then subview it.
-    // NOTE: we can set protected members, since f is the same type of this class.
+  if (parent != nullptr) {
     Field f;
     f.m_header = parent;
     f.m_data   = m_data;
 
-  //   // Take 2 dimensional view with normal LayoutRight
-  //   auto v_np1 = f.get_ND_view<HD,DstValueType,2>();
-      // const auto vfs_rank = f.get_header().get_identifier().get_layout().rank();
-      auto v_fullsize = f.get_ND_view<HD, DstValueType, 4>();
+    auto v_fullsize = f.get_ND_view<HD, DstValueType, N>();
 
     // Now we can subview v_np1 at the correct slice
     const auto& info = m_header->get_alloc_properties().get_subview_info();
-    const int idim = info.dim_idx;
-    const int k    = info.slice_idx;
-    const int k_end = info.slice_idx_end;
-
-  //   // So far we can only subview at first or second dimension.
-  //   EKAT_REQUIRE_MSG (idim==0 || idim==1,
-  //       "Error! Subview dimension index is out of bounds.\n");
-
-    // Use correct subview utility
-    if (idim==0) {
-      // FIXME: what's a better way to do this? can't use v_fullsize after the
-      // logic block b/c of scoping, and it's also not easy to know what type of
-      // exotic
-      // we know it's a subview, so if it has the same rank as its parent, then
-      // we know that it's a multi-slice subview
-      if (fl.rank() == f.get_header().get_identifier().get_layout().rank()) {
-        auto svs = ekat::subview(v_fullsize, Kokkos::make_pair<int, int>(k, k_end), idim);
-        // for (size_t i = 0; i < 4; i++)
-        // {
-        //   std::cout << "i = " << i << "\n";
-        //   std::cout << "svs.stride(i) = " << svs.stride(i) << "\n";
-        //   std::cout << "svs.extent(i) = " << svs.extent(i) << "\n";
-        // }
-        return svs;
-      } else {
-        // return DstView(ekat::subview(v_np1,k));
-      }
-    } else {
-      // return DstView(ekat::subview_1(v_np1,k));
-    }
+    const int idim   = info.dim_idx;
+    const int k      = info.slice_idx;
+    const int k_end  = info.slice_idx_end;
+
+    // this version of ekat::subview overloaded conveniently, so only the single
+    // version of the call is required here
+    return DstView(ekat::subview(v_fullsize,
+                                 Kokkos::make_pair<int, int>(k, k_end), idim));
   }
 }
 
@@ -838,7 +806,6 @@ auto Field::get_ND_view () const ->
   return ret_type (ptr,kl);
 }
 
-// TODO: will need to set this up for multi-sliced subfield
 template<HostOrDevice HD,typename T,int N>
 auto Field::get_ND_view () const ->
   if_t<N==MaxRank,get_view_type<data_nd_t<T,N>,HD>>
@@ -848,7 +815,6 @@ auto Field::get_ND_view () const ->
       "Error! Input Rank must either be 1 (flat array) or the actual field rank.\n");
 
   // Given that N==MaxRank, this field cannot be a subview of another field
-  // NOTE: this will not be true for multi-slice
   EKAT_REQUIRE_MSG (m_header->get_parent().expired(),
       "Error! A view of rank " + std::to_string(MaxRank) + " should not be the subview of another field.\n");
 

components/eamxx/src/share/tests/field_tests.cpp

@@ -312,38 +312,31 @@ TEST_CASE("field", "") {
     f1.allocate_view();
     randomize(f1,engine,pdf);
 
-    const int idim = 0;
-    const int ivar = 2;
-    const int sl_beg = 2;
-    const int sl_end = 4;
-
-    auto f3 = f1.subfield(idim, sl_beg, sl_end);
-
-    // const auto fl3 = f3.get_header().get_identifier().get_layout();
-    // const auto alprop3 = f3.get_header().get_alloc_properties();
-    // std::cout << "fl3.rank() = " << fl3.rank() << "\n";
-    // auto f3dim = fl3.dims();
-    // for (size_t i = 0; i < 4; i++)
-    // {
-    //   std::cout << "f3dim(i) = " << f3dim[i] << "\n";
-    // }
-    // auto f3ext = fl3.extents();
-    // for (size_t i = 0; i < 4; i++)
-    // {
-    //   std::cout << "f3ext(i) = " << f3ext(i) << "\n";
-    // }
-    // std::cout << "alprop3.is_subfield() = " << alprop3.is_subfield() << "\n";
-
-    auto v3_h = f3.get_strided_view<Real****, Host>(true);
+    const int idim[4] = {0, 1, 2, 3};
+    const int sl_beg[4] = {2, 3, 0, 9};
+    const int sl_end[4] = {4, 6, 1, 15};
+
     auto v4d_h = f1.get_view<Real****, Host>();
 
-    for (size_t i = sl_beg; i < sl_end; i++) {
-      for (size_t j = 0; j < d1[1]; j++) {
-        for (size_t k = 0; k < d1[2]; k++) {
-          for (size_t l = 0; l < d1[3]; l++) {
-            // std::cout << "v4d_h(i, j, k, l) = " << v4d_h(i, j, k, l) << "\n";
-            // std::cout << "v3_h(i - sl_beg, j, k, l) = " << v3_h(i - sl_beg, j, k, l) << "\n";
-            REQUIRE(v4d_h(i, j, k, l) == v3_h(i - sl_beg, j, k, l));
+    int i1, i2, j1, j2, k1, k2, l1, l2;
+
+    for (int ens = 0; ens < 4; ens++) {
+      auto sf = f1.subfield(idim[ens], sl_beg[ens], sl_end[ens]);
+      auto sv_h = sf.get_multi_sliced_view<Real, 4, Host>();
+      i1 = (ens == 0) ? sl_beg[0] : 0;
+      i2 = (ens == 0) ? sl_end[0] : d1[0];
+      j1 = (ens == 1) ? sl_beg[1] : 0;
+      j2 = (ens == 1) ? sl_end[1] : d1[1];
+      k1 = (ens == 2) ? sl_beg[2] : 0;
+      k2 = (ens == 2) ? sl_end[2] : d1[2];
+      l1 = (ens == 3) ? sl_beg[3] : 0;
+      l2 = (ens == 3) ? sl_end[3] : d1[3];
+      for (int i = i1; i < i2; i++) {
+        for (int j = j1; j < j2; j++) {
+          for (int k = k1; k < k2; k++) {
+            for (int l = l1; l < l2; l++) {
+              REQUIRE(v4d_h(i, j, k, l) == sv_h(i - i1, j - j1, k - k1, l - l1));
+            }
           }
         }
       }