Merge pull request #2785 from E3SM-Project/mjs/eamxx/fix-#1312

Enable Subfield Slices Spanning a Range of indices (Issue #1312)

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

@@ -115,6 +115,48 @@ subfield (const int idim, const int index, const bool dynamic) const {
   return subfield(m_header->get_identifier().name(),idim,index,dynamic);
 }
 
+// slice at index idim, extracting the N = (index_end - index_beg) entries
+// written in math notation: [index_beg, index_end)
+// or equivalently, subF = F(index_beg, ... , index_beg + N)
+Field 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 {
+
+  const auto& id = m_header->get_identifier();
+  const auto& lt = id.get_layout();
+
+  // Sanity checks
+  EKAT_REQUIRE_MSG(
+      is_allocated(),
+      "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);
+  // Create identifier for subfield
+  FieldIdentifier sf_id(sf_name, sf_layout, sf_units, id.get_grid_name());
+
+  // Create empty subfield, then set header and views
+  // Note: we can access protected members, since it's the same type
+  Field sf;
+  sf.m_header = create_subfield_header(sf_id, m_header, idim, index_beg,
+                                       index_end);
+  sf.m_data = m_data;
+
+  return sf;
+}
+
+Field Field::subfield(const std::string& sf_name, const int idim,
+                      const int index_beg, const int index_end) const {
+  const auto& id = m_header->get_identifier();
+  return subfield(sf_name, id.get_units(), idim, index_beg, index_end);
+}
+
+Field Field::subfield(const int idim, const int index_beg,
+                      const int index_end) const {
+  return subfield(m_header->get_identifier().name(), idim, index_beg,
+                  index_end);
+}
+
 Field Field::
 get_component (const int i, const bool dynamic) {
   const auto& layout = get_header().get_identifier().get_layout();
@@ -127,11 +169,31 @@ get_component (const int i, const bool dynamic) {
   EKAT_REQUIRE_MSG (i>=0 && i<layout.dim(idim),
       "Error! Component index out of bounds [0," + std::to_string(layout.dim(idim)) + ").\n");
 
-  // Add _$idim to the field name, to avoid issues if the subfield is stored
+  // Add _$i to the field name, to avoid issues if the subfield is stored
   // in some structure that requires unique names (e.g., a remapper)
   return subfield (fname + "_" + std::to_string(i),idim,i,dynamic);
 }
 
+Field Field::get_components(const int beg, const int end) {
+  const auto& layout = get_header().get_identifier().get_layout();
+  const auto& fname = get_header().get_identifier().name();
+  EKAT_REQUIRE_MSG(layout.is_vector_layout(),
+                   "Error! 'get_component' available only for vector fields.\n"
+                   "       Layout of '" +
+                       fname + "': " + e2str(layout.type()) + "\n");
+
+  const int idim = layout.get_vector_component_idx();
+  EKAT_REQUIRE_MSG(beg >= 0 && end < layout.dim(idim),
+                   "Error! Component index range out of bounds [0," +
+                       std::to_string(layout.dim(idim)) + ").\n");
+  EKAT_REQUIRE_MSG(beg < end, "Error! Invalid component indices (beg >= end).\n");
+
+  // Add _$beg-$end to the field name, to avoid issues if the subfield is stored
+  // in some structure that requires unique names (e.g., a remapper)
+  return subfield(fname + "_" + std::to_string(beg) + "-" + std::to_string(end),
+                  idim, beg, end);
+}
+
 bool Field::equivalent(const Field& rhs) const
 {
   return (m_header==rhs.m_header &&

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

@@ -241,11 +241,20 @@ 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;
-
+  // extracts a subfield composed of 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], :)
+  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;
+  Field subfield (const std::string& sf_name, const int idim,
+                  const int index_beg, const int index_end) const;
+  Field subfield (const int idim, const int index_beg, const int index_end) const;
   // If this field is a vector field, get a subfield for the ith component.
   // If dynamic = true, it is possible to "reset" the component index at runtime.
   // Note: throws if this is not a vector field.
   Field get_component (const int i, const bool dynamic = false);
+  // version for slicing across multiple, contiguous indices
+  Field get_components (const int beg, const int end);
 
   // Checks whether the underlying view has been already allocated.
   bool is_allocated () const { return m_data.d_view.data()!=nullptr; }
@@ -303,25 +312,32 @@ class Field {
   if_t<(N<=2),
        get_view_type<data_nd_t<T,N-1>,HD>>
   get_subview_1 (const get_view_type<data_nd_t<T,N>,HD>&, const int) const {
-    EKAT_ERROR_MSG ("Error! Cannot subview a rank2 view along the second dimension without losing LayoutRight.\n");
+    EKAT_ERROR_MSG ("Error! Cannot subview a rank2 view along the second "
+                    "dimension without losing LayoutRight.\n");
   }
 
   template<HostOrDevice HD,typename T,int N>
   auto get_ND_view () const
-    -> if_t<N==MaxRank, get_view_type<data_nd_t<T,N>,HD>>;
+    -> if_t<(N < MaxRank), get_view_type<data_nd_t<T,N>,HD>>;
+
+  template<HostOrDevice HD,typename T,int N>
+  auto get_ND_view () const
+    -> if_t<N == MaxRank, get_view_type<data_nd_t<T,N>,HD>>;
 
+  // NOTE: DO NOT USE--it only returns an error and is here to protect
+  // against compiler errors related to sliced subviews in get_strided_view()
   template<HostOrDevice HD,typename T,int N>
   auto get_ND_view () const
-    -> if_t<(N<MaxRank), get_view_type<data_nd_t<T,N>,HD>>;
+    -> if_t<(N >= MaxRank + 1), get_view_type<data_nd_t<T,N>,HD>>;
 
-  // Metadata (name, rank, dims, customere/providers, time stamp, ...)
-  std::shared_ptr<header_type>            m_header;
+  // Metadata (name, rank, dims, customer/providers, time stamp, ...)
+  std::shared_ptr<header_type> m_header;
 
   // Actual data.
-  dual_view_t<char*>    m_data;
+  dual_view_t<char*> m_data;
 
   // Whether this field is read-only
-  bool                  m_is_read_only = false;
+  bool m_is_read_only = false;
 };
 
 // We use this to find a Field in a std container.

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

@@ -67,21 +67,63 @@ subview (const int idim, const int k, const bool dynamic) const {
     // and there is no packing
     props.m_last_extent = m_layout.dim(idim);
     props.m_pack_size_max = 1;
-    props.m_alloc_size = m_alloc_size / m_last_extent;
   } else {
     // We are keeping the last dim, so same last extent and max pack size
     props.m_last_extent = m_last_extent;
     props.m_pack_size_max = m_pack_size_max;
-    props.m_alloc_size = m_alloc_size / m_layout.dim(idim);
   }
+  // give it an invalid value because subviews don't get allocated
+  props.m_alloc_size = -1;
+  return props;
+}
+
+FieldAllocProp FieldAllocProp::subview(const int idim,
+                                       const int index_beg,
+                                       const int index_end) const {
+  EKAT_REQUIRE_MSG(
+      is_committed(),
+      "Error! Subview requires alloc properties to be committed.\n");
+  EKAT_REQUIRE_MSG(idim < m_layout.rank(),
+                   "Error! Dimension index out of bounds.\n");
+  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),
+      "Error! Slice index range along the idim dimension is out of bounds.\n");
+
+  // Set new layout basic stuff
+  FieldAllocProp props(m_scalar_type_size);
+  props.m_committed = true;
+  props.m_scalar_type_size = m_scalar_type_size;
+  props.m_layout = m_layout.clone();
+  props.m_layout.reset_dim(idim, index_end - index_beg);
+
+  // we do not currently have the capability for contiguous multi-slice views
+  props.m_contiguous = false;
+
+  props.m_subview_info =
+      SubviewInfo(idim, index_beg, index_end, m_layout.dim(idim));
+  // Figure out strides/packs
+  if (idim == (m_layout.rank() - 1)) {
+    // We're slicing the possibly padded dim, so everything else is as in the
+    // layout, and there is no packing
+    props.m_last_extent = index_end - index_beg;
+    props.m_pack_size_max = 1;
+  } else {
+    // We are keeping the last dim, so same last extent and max pack size
+    props.m_last_extent = m_last_extent;
+    props.m_pack_size_max = m_pack_size_max;
+  }
+  // give it an invalid value because subviews don't get allocated
+  props.m_alloc_size = -1;
   return props;
 }
 
 void FieldAllocProp::request_allocation (const int pack_size) {
   using ekat::ScalarTraits;
 
   EKAT_REQUIRE_MSG(!m_committed,
-      "Error! Cannot change allocation properties after they have been commited.\n");
+      "Error! Cannot change allocation properties after they have been committed.\n");
 
   const int vts = m_scalar_type_size*pack_size;
 

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

@@ -25,11 +25,11 @@ namespace scream
  *      For instance, say you declare a field with dimensions (3,10).
  *      In the field manager, this would be stored as a 1d field: View<Real*>.
  *      Let's say you have one class that uses the field as View<Real**>,
- *      and another class that used the field in a packed way, taht is,
- *      as View<Pack<Real,4>**>. The second view will have dimensions (3,4),
+ *      and another class that used the field in a packed way, that is,
+ *      as View<Pack<Real,4>**>. The second view will have dimensions (3,3),
         in terms of its value type (i.e., Pack<Real,4>).
  *      This means the field needs an allocation bigger than 30 Real's, namely
- *      3*4*4=36 Real's. This class does the book-keeping for the allocation size,
+ *      3*4*3=36 Real's. This class does the book-keeping for the allocation size,
  *      so that 1) the field can be allocated with enough memory to accommodate
  *      all requests, 2) customers of the field can check what the allocation
  *      is, so that they know whether there is padding in the field, and
@@ -61,21 +61,38 @@ namespace scream
 
 // Helper struct to store info related to the subviewing process
 struct SubviewInfo {
-  SubviewInfo () = default;
-  SubviewInfo (const int dim, const int slice, const int extent, const bool is_dynamic) :
-    dim_idx(dim), slice_idx(slice), dim_extent(extent), dynamic(is_dynamic) {}
-  SubviewInfo (const SubviewInfo&) = default;
-  SubviewInfo& operator= (const SubviewInfo&) = default;
-
-  int dim_idx    = -1;  // Dimension along which slicing happened
-  int slice_idx  = -1;  // Slice along dimension $dim_idx
-  int dim_extent = -1;  // Extent of dimension $dim_idx
-  bool dynamic   = false; // Whether this is a dynamic subview (slice_idx can change)
+  SubviewInfo() = default;
+  SubviewInfo(const int dim, const int slice, const int extent,
+              const bool is_dynamic)
+      : dim_idx(dim), slice_idx(slice), dim_extent(extent),
+        dynamic(is_dynamic) {}
+  // multi-slice subview across contiguous indices
+  SubviewInfo(const int dim, const int slice_beg, const int slice_end,
+              const int extent)
+      : dim_idx(dim), slice_idx(slice_beg), slice_idx_end(slice_end),
+        dim_extent(extent) {}
+  SubviewInfo(const SubviewInfo&) = default;
+  SubviewInfo& operator=(const SubviewInfo&) = default;
+
+  // Dimension along which slicing happened
+  int dim_idx = -1;
+  // in the case of single slice, this is the slice index
+  // in the case of multi-slice, this is the beginning index
+  int slice_idx = -1;
+  // ending slice index for multi-slice (remains == -1 if single-slice subview)
+  int slice_idx_end = -1;
+  // Extent of dimension $dim_idx
+  int dim_extent = -1;
+  // Whether this is a dynamic subview (slice_idx can change)
+  bool dynamic = false;
 };
 
 inline bool operator== (const SubviewInfo& lhs, const SubviewInfo& rhs) {
   return lhs.dim_idx==rhs.dim_idx &&
          lhs.slice_idx==rhs.slice_idx &&
+         //  slice_idx_end == -1 for single slice, and the ending index when
+         // it's a multi-slice
+         lhs.slice_idx_end==rhs.slice_idx_end &&
          lhs.dim_extent==rhs.dim_extent &&
          lhs.dynamic==rhs.dynamic;
 }
@@ -90,11 +107,14 @@ 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;
 
+  // multi-slice subview over contiguous indices
+  FieldAllocProp subview (const int idim, const int k_beg, const int k_end) const;
+
   // Request allocation able to accommodate a pack of ScalarType of the given pack size
   void request_allocation (const int pack_size = 1);
 

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

@@ -2,8 +2,7 @@
 
 #include "ekat/std_meta/ekat_std_utils.hpp"
 
-namespace scream
-{
+namespace scream {
 
 FieldHeader::FieldHeader (const identifier_type& id)
  : m_identifier (id)
@@ -33,10 +32,7 @@ set_extra_data (const std::string& key,
   }
 }
 
-std::shared_ptr<FieldHeader>
-FieldHeader::
-alias (const std::string& name) const
-{
+std::shared_ptr<FieldHeader> FieldHeader::alias(const std::string& name) const {
   auto fh = create_header(get_identifier().alias(name));
   fh->m_tracking = m_tracking;
   fh->m_alloc_prop = m_alloc_prop;
@@ -72,4 +68,33 @@ create_subfield_header (const FieldIdentifier& id,
   return fh;
 }
 
+// subfield with multiple, contiguous slices
+std::shared_ptr<FieldHeader>
+create_subfield_header(const FieldIdentifier& id,
+                       std::shared_ptr<FieldHeader> parent, const int idim,
+                       const int k_beg, const int k_end) {
+  // Sanity checks
+  EKAT_REQUIRE_MSG(parent != nullptr,
+                   "Error! Invalid pointer for parent header.\n");
+  EKAT_REQUIRE_MSG(k_end > k_beg,
+                   "Error! Slice indices are invalid (non-increasing).\n");
+
+  // Create header, and set up parent/child
+  auto fh = create_header(id);
+  fh->create_parent_child_link(parent);
+
+  // Create tracking, and set up parent/child
+  fh->m_tracking = create_tracking();
+  fh->m_tracking->create_parent_child_link(parent->get_tracking_ptr());
+  if (parent->get_tracking().get_time_stamp().is_valid()) {
+    fh->m_tracking->update_time_stamp(parent->get_tracking().get_time_stamp());
+  }
+
+  // Create alloc props
+  fh->m_alloc_prop = std::make_shared<FieldAllocProp>(
+      parent->get_alloc_properties().subview(idim, k_beg, k_end));
+
+  return fh;
+}
+
 } // namespace scream

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

@@ -91,6 +91,11 @@ class FieldHeader : public FamilyTracking<FieldHeader> {
   create_subfield_header (const FieldIdentifier&,
                           std::shared_ptr<FieldHeader>,
                           const int, const int, const bool);
+  // for creating multi-slice subfield (continuous indices)
+  friend std::shared_ptr<FieldHeader>
+  create_subfield_header (const FieldIdentifier&,
+                          std::shared_ptr<FieldHeader>,
+                          const int idim, const int k_beg, const int k_end);
 
   // NOTE: the identifier *cannot* be a shared_ptr, b/c we
   //       don't foresee sharing an identifier between two
@@ -169,6 +174,10 @@ std::shared_ptr<FieldHeader>
 create_subfield_header (const FieldIdentifier& id,
                         std::shared_ptr<FieldHeader> parent,
                         const int idim, const int k, const bool dynamic);
+std::shared_ptr<FieldHeader>
+create_subfield_header (const FieldIdentifier& id,
+                        std::shared_ptr<FieldHeader> parent,
+                        const int idim, const int k_beg, const int k_end);
 
 } // namespace scream