many size_of improvements

- introduced size_of for std::vector, boost::container::small_vector, Range1, Range, TiledRange1, TiledRange
- improved/fixed size_of(Tensor) and size_of(SparseShape)
- stronger unit test for size_of(DistArray)

Author: evaleev

src/TiledArray/dist_array.h

@@ -1743,8 +1743,9 @@ class DistArray : public madness::archive::ParallelSerializableObject {
 
 };  // class DistArray
 
-/// \return the number of bytes used by \p t in this rank's memory space
-/// `S`
+/// \return the approximate number of bytes used by \p t in this rank's
+/// memory space `S`
+/// \note this does not account for the TiledRange and some other metadata
 template <MemorySpace S, typename Tile, typename Policy>
 std::size_t size_of(const DistArray<Tile, Policy>& da) {
   std::size_t result = 0;

src/TiledArray/range.h

@@ -21,6 +21,7 @@
 #define TILEDARRAY_RANGE_H__INCLUDED
 
 #include <TiledArray/permutation.h>
+#include <TiledArray/platform.h>
 #include <TiledArray/range1.h>
 #include <TiledArray/range_iterator.h>
 #include <TiledArray/size_array.h>
@@ -1247,6 +1248,20 @@ class Range {
     return ordinal(last) - ordinal(first);
   }
 
+  template <MemorySpace S>
+  friend constexpr std::size_t size_of(const Range& r) {
+    std::size_t sz = 0;
+    if constexpr (S == MemorySpace::Host) {
+      sz += sizeof(r);
+    }
+    // correct for optional dynamic allocation of datavec_
+    if constexpr (S == MemorySpace::Host) {
+      sz -= sizeof(r.datavec_);
+    }
+    sz += size_of<S>(r.datavec_);
+    return sz;
+  }
+
 };  // class Range
 
 // lift Range::index_type and Range::index_view_type into user-land

src/TiledArray/range1.h

@@ -26,6 +26,7 @@
 #include <boost/iterator/iterator_facade.hpp>
 
 #include <TiledArray/error.h>
+#include <TiledArray/platform.h>
 
 namespace TiledArray {
 
@@ -138,7 +139,7 @@ struct Range1 {
     /// \brief dereferences this iterator
     /// \return const reference to the current index
     const auto& dereference() const { return v; }
-  };
+  };  // class Iterator
   friend class Iterator;
 
   typedef Iterator const_iterator;  ///< Iterator type
@@ -201,6 +202,15 @@ struct Range1 {
   void serialize(Archive& ar) const {
     ar & first & second;
   }
+
+  template <MemorySpace S>
+  friend constexpr std::size_t size_of(const Range1& r) {
+    std::size_t sz = 0;
+    if constexpr (S == MemorySpace::Host) {
+      sz += sizeof(r);
+    }
+    return sz;
+  }
 };
 
 inline bool operator==(const Range1& x, const Range1& y) {

src/TiledArray/sparse_shape.h

@@ -1734,7 +1734,19 @@ typename SparseShape<T>::value_type SparseShape<T>::threshold_ =
 
 template <MemorySpace S, typename T>
 std::size_t size_of(const SparseShape<T>& shape) {
-  return size_of<S>(shape.tile_norms_);
+  std::size_t sz = 0;
+  if constexpr (S == MemorySpace::Host) {
+    sz += sizeof(shape);
+  }
+  // account for dynamically-allocated content
+  if constexpr (S == MemorySpace::Host) {
+    sz -= sizeof(shape.tile_norms_);
+  }
+  sz += size_of<S>(shape.tile_norms_);
+  if (shape.tile_norms_unscaled_) {
+    sz += size_of<S>(*(shape.tile_norms_unscaled_));
+  }
+  return sz;
 }
 
 /// Add the shape to an output stream

src/TiledArray/tensor/tensor.h

@@ -2779,17 +2779,21 @@ class Tensor {
 
 /// \return the number of bytes used by \p t in memory space
 /// `S`
-/// \warning footprint of Range is approximated, will not be exact for
-/// tensor orders highers than `TA_MAX_SOO_RANK_METADATA`
 template <MemorySpace S, typename T, typename A>
 std::size_t size_of(const Tensor<T, A>& t) {
   std::size_t result = 0;
   if constexpr (S == MemorySpace::Host) {
     result += sizeof(t);
   }
+  // correct for optional dynamic allocation of Range
+  if constexpr (S == MemorySpace::Host) {
+    result -= sizeof(Range);
+  }
+  result += size_of<S>(t.range());
+
   if (allocates_memory_space<S>(A{})) {
     if (!t.empty()) {
-      if constexpr (is_constexpr_size_of_v<S, Tensor<T, A>>) {
+      if constexpr (is_constexpr_size_of_v<S, T>) {
         result += t.size() * sizeof(T);
       } else {
         result += std::accumulate(

src/TiledArray/tiled_range.h

@@ -398,6 +398,24 @@ class TiledRange {
   range_type elements_range_;  ///< Range of element indices
   Ranges ranges_;  ///< tiled (1d) range, aka TiledRange1, for each mode
                    ///< `*this` is a direct product of these tilings
+
+  template <MemorySpace S>
+  friend constexpr std::size_t size_of(const TiledRange& r) {
+    std::size_t sz = 0;
+    if constexpr (S == MemorySpace::Host) {
+      sz += sizeof(r);
+    }
+    // correct for optional dynamic allocation of range_ and elements_range_
+    if constexpr (S == MemorySpace::Host) {
+      sz -= sizeof(r.range_);
+      sz -= sizeof(r.elements_range_);
+      sz -= sizeof(r.ranges_);
+    }
+    sz += size_of<S>(r.range_);
+    sz += size_of<S>(r.elements_range_);
+    sz += size_of<S>(r.ranges_);
+    return sz;
+  }
 };
 
 /// TiledRange permutation operator.

src/TiledArray/tiled_range1.h

@@ -23,6 +23,7 @@
 #include <TiledArray/config.h>
 
 #include <TiledArray/error.h>
+#include <TiledArray/platform.h>
 #include <TiledArray/range1.h>
 #include <TiledArray/type_traits.h>
 #include <TiledArray/utility.h>
@@ -493,6 +494,27 @@ class TiledRange1 {
   mutable std::shared_ptr<const index1_type[]>
       elem2tile_;  ///< maps element index to tile index (memoized data).
 
+  template <MemorySpace S>
+  friend constexpr std::size_t size_of(const TiledRange1& r) {
+    std::size_t sz = 0;
+    if constexpr (S == MemorySpace::Host) {
+      sz += sizeof(r);
+    }
+    // correct for optional dynamic allocation of range_ and elements_range_
+    if constexpr (S == MemorySpace::Host) {
+      sz -= sizeof(r.range_);
+      sz -= sizeof(r.elements_range_);
+      sz -= sizeof(r.tiles_ranges_);
+    }
+    sz += size_of<S>(r.range_);
+    sz += size_of<S>(r.elements_range_);
+    sz += size_of<S>(r.tiles_ranges_);
+    if (r.elem2tile_) {
+      sz += r.elements_range_.extent() * sizeof(index1_type);
+    }
+    return sz;
+  }
+
 };  // class TiledRange1
 
 /// Exchange the data of the two given ranges.

src/TiledArray/util/vector.h

@@ -45,6 +45,7 @@
 #include <TiledArray/utility.h>
 #include <madness/world/archive.h>
 #include "TiledArray/error.h"
+#include "TiledArray/platform.h"
 
 namespace TiledArray {
 
@@ -55,6 +56,66 @@ using vector = std::vector<T>;
 template <typename T, std::size_t N = TA_MAX_SOO_RANK_METADATA>
 using svector = boost::container::small_vector<T, N>;
 
+}  // namespace container
+
+// size_of etc.
+
+template <MemorySpace S, typename T, typename A>
+constexpr std::size_t size_of(const std::vector<T, A>& t) {
+  std::size_t sz = 0;
+  if constexpr (S == MemorySpace::Host) {
+    sz += sizeof(t);
+  }
+  if constexpr (allocates_memory_space<S>(t.get_allocator())) {
+    if constexpr (is_constexpr_size_of_v<S, T>) {
+      sz += sizeof(T) * t.capacity();
+    } else {
+      sz += std::accumulate(
+          t.begin(), t.end(), std::size_t{0},
+          [](const std::size_t s, const T& t) { return s + size_of<S>(t); });
+      sz += (t.capacity() - t.size()) * sizeof(T);
+    }
+  }
+  return sz;
+}
+
+template <MemorySpace S, typename T, typename VoidAlloc, typename Options>
+constexpr bool allocates_memory_space(
+    const boost::container::small_vector_allocator<T, VoidAlloc, Options>& a) {
+  return S == MemorySpace::Host;
+}
+
+template <MemorySpace S, typename T, std::size_t N, typename A>
+constexpr std::size_t size_of(
+    const boost::container::small_vector<T, N, A>& t) {
+  std::size_t sz = 0;
+  if constexpr (S == MemorySpace::Host) {
+    sz += sizeof(t);
+  }
+  if (allocates_memory_space<S>(t.get_allocator())) {
+    std::size_t data_sz = 0;
+    if constexpr (is_constexpr_size_of_v<S, T>) {
+      data_sz += sizeof(T) * t.capacity();
+    } else {
+      data_sz += std::accumulate(
+          t.begin(), t.end(), std::size_t{0},
+          [](const std::size_t s, const T& t) { return s + size_of<S>(t); });
+      data_sz += (t.capacity() - t.size()) * sizeof(T);
+    }
+    if (t.capacity() > N ||
+        S != MemorySpace::Host) {  // dynamically allocated buffer => account
+                                   // for tne entirety of data_sz
+      sz += data_sz;
+    } else {  // data in internal buffer => account for the
+              // dynamically-allocated part of data_sz
+      sz += (data_sz - t.capacity() * sizeof(T));
+    }
+  }
+  return sz;
+}
+
+namespace container {
+
 template <typename Range>
 std::enable_if_t<detail::is_integral_range_v<Range> &&
                      detail::is_sized_range_v<Range>,

tests/dist_array.cpp

@@ -1000,10 +1000,11 @@ BOOST_AUTO_TEST_CASE(size_of) {
 
   auto sz0 = TiledArray::size_of<MemorySpace::Host>(array);
   world.gop.sum(sz0);
-  if (world.size() == 1)
-    BOOST_REQUIRE(sz0 == 56688);
-  else
-    BOOST_REQUIRE(sz0 > 56688);
+  const auto sz0_expected =
+      /* size on 1 rank */ 56728 +
+      /* size of shape on ranks 1 ... N-1 */ (world.size() - 1) *
+          TiledArray::size_of<MemorySpace::Host>(array.shape());
+  BOOST_REQUIRE(sz0 == sz0_expected);
 }
 
 BOOST_AUTO_TEST_SUITE_END()