introduced size_of

Author: evaleev

src/TiledArray/dense_shape.h

@@ -27,7 +27,10 @@
 #define TILEDARRAY_DENSE_SHAPE_H__INCLUDED
 
 #include <TiledArray/config.h>
+
+#include <TiledArray/platform.h>
 #include <TiledArray/type_traits.h>
+
 #include <cstdint>
 
 namespace madness {
@@ -391,6 +394,11 @@ class DenseShape {
       std::numeric_limits<value_type>::epsilon();
 };  // class DenseShape
 
+template <MemorySpace S>
+std::size_t size_of(const DenseShape& shape) {
+  return sizeof(shape);
+}
+
 constexpr inline bool operator==(const DenseShape& a, const DenseShape& b) {
   return true;
 }

src/TiledArray/dist_array.h

@@ -1743,6 +1743,19 @@ class DistArray : public madness::archive::ParallelSerializableObject {
 
 };  // class DistArray
 
+/// \return the number of bytes used by \p t in this rank's memory space
+/// `S`
+template <MemorySpace S, typename Tile, typename Policy>
+std::size_t size_of(const DistArray<Tile, Policy>& da) {
+  std::size_t result = 0;
+  result += size_of<S>(da.shape());
+  // add up local tile's contributions
+  for (const auto& tile_ref : da) {
+    result += size_of<S>(tile_ref.get());
+  }
+  return result;
+}
+
 #ifndef TILEDARRAY_HEADER_ONLY
 
 extern template class DistArray<Tensor<double>, DensePolicy>;

src/TiledArray/sparse_shape.h

@@ -28,6 +28,7 @@
 
 #include <TiledArray/fwd.h>
 
+#include <TiledArray/platform.h>
 #include <TiledArray/tensor.h>
 #include <TiledArray/tensor/shift_wrapper.h>
 #include <TiledArray/tensor/tensor_interface.h>
@@ -1721,13 +1722,21 @@ class SparseShape {
     return cast_abs_factor;
   }
 
+  template <MemorySpace S, typename T_>
+  friend std::size_t size_of(const SparseShape<T_>& shape);
+
 };  // class SparseShape
 
 // Static member initialization
 template <typename T>
 typename SparseShape<T>::value_type SparseShape<T>::threshold_ =
     std::numeric_limits<T>::epsilon();
 
+template <MemorySpace S, typename T>
+std::size_t size_of(const SparseShape<T>& shape) {
+  return size_of<S>(shape.tile_norms_);
+}
+
 /// Add the shape to an output stream
 
 /// \tparam T the numeric type supporting the type of \c shape

src/TiledArray/tensor/tensor.h

@@ -24,6 +24,7 @@
 
 #include "TiledArray/external/umpire.h"
 #include "TiledArray/host/env.h"
+#include "TiledArray/platform.h"
 
 #include "TiledArray/math/blas.h"
 #include "TiledArray/math/gemm_helper.h"
@@ -2776,6 +2777,30 @@ class Tensor {
 
 };  // 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);
+  }
+  if (allocates_memory_space<S>(A{})) {
+    if (!t.empty()) {
+      if constexpr (is_constexpr_size_of_v<S, Tensor<T, A>>) {
+        result += t.size() * sizeof(T);
+      } else {
+        result += std::accumulate(
+            t.begin(), t.end(), std::size_t{0},
+            [](const std::size_t s, const T& t) { return s + size_of<S>(t); });
+      }
+    }
+  }
+  return result;
+}
+
 #ifdef TA_TENSOR_MEM_TRACE
 template <typename T, typename A>
 std::size_t Tensor<T, A>::trace_if_larger_than_ =

tests/dist_array.cpp

@@ -944,4 +944,63 @@ BOOST_AUTO_TEST_CASE(reduction) {
   BOOST_REQUIRE(array_norm = std::sqrt(TA::dot(array, array)));
 }
 
+BOOST_AUTO_TEST_CASE(size_of) {
+  using Numeric = double;
+  using T = Tensor<Numeric>;
+  using ToT = Tensor<T>;
+  using Policy = SparsePolicy;
+  using ArrayToT = DistArray<ToT, Policy>;
+
+  auto unit_T = [](Range const& rng) { return T(rng, Numeric{1}); };
+
+  auto unit_ToT = [unit_T](Range const& rngo, Range const& rngi) {
+    return ToT(rngo, unit_T(rngi));
+  };
+
+  size_t constexpr nrows = 3;
+  size_t constexpr ncols = 4;
+  TiledRange const trange({{0, 2, 5, 7}, {0, 5, 7, 10, 12}});
+  TA_ASSERT(trange.tiles_range().extent().at(0) == nrows &&
+                trange.tiles_range().extent().at(1) == ncols,
+            "Following code depends on this condition.");
+
+  // this Range is used to construct all inner tensors of the tile with
+  // tile index @c tix.
+  auto inner_dims = [nrows, ncols](Range::index_type const& tix) -> Range {
+    static std::array<size_t, nrows> const rows{7, 8, 9};
+    static std::array<size_t, ncols> const cols{7, 8, 9, 10};
+
+    TA_ASSERT(tix.size() == 2, "Only rank-2 tensor expected.");
+    return Range({rows[tix.at(0) % nrows], cols[tix.at(1) % ncols]});
+  };
+
+  // let's make all 'diagonal' tiles zero
+  auto zero_tile = [](Range::index_type const& tix) -> bool {
+    return tix.at(0) == tix.at(1);
+  };
+
+  auto make_tile = [inner_dims,  //
+                    zero_tile,   //
+                    &trange,     //
+                    unit_ToT](auto& tile, auto const& rng) {
+    auto&& tix = trange.element_to_tile(rng.lobound());
+    if (zero_tile(tix))
+      return 0.;
+    else {
+      tile = unit_ToT(rng, inner_dims(tix));
+      return tile.norm();
+    }
+  };
+
+  auto& world = get_default_world();
+
+  // all non-zero inner tensors of this ToT array are unit (ie all
+  // inner tensors' elements are 1.)
+  auto array = make_array<ArrayToT>(world, trange, make_tile);
+
+  auto sz0 = TiledArray::size_of<MemorySpace::Host>(array);
+  world.gop.sum(sz0);
+  BOOST_REQUIRE(sz0 == 56688);
+}
+
 BOOST_AUTO_TEST_SUITE_END()

tests/tensor.cpp

@@ -783,4 +783,32 @@ BOOST_AUTO_TEST_CASE(print) {
   // std::cout << tb;
 }
 
+BOOST_AUTO_TEST_CASE(size_of) {
+  auto sz0h = TiledArray::size_of<TiledArray::MemorySpace::Host>(TensorN{});
+  BOOST_REQUIRE(sz0h == sizeof(TensorN));
+
+  auto sz0d = TiledArray::size_of<TiledArray::MemorySpace::Device>(TensorN{});
+  BOOST_REQUIRE(sz0d == 0);
+
+  auto sz0um =
+      TiledArray::size_of<TiledArray::MemorySpace::Device_UM>(TensorN{});
+  BOOST_REQUIRE(sz0um == 0);
+
+  auto sz1 = TiledArray::size_of<TiledArray::MemorySpace::Host>(
+      TensorZ(Range(2, 3, 4)));
+  BOOST_REQUIRE(sz1 ==
+                sizeof(TensorZ) + 2 * 3 * 4 * sizeof(TensorZ::value_type));
+
+  using TTD = Tensor<Tensor<double>>;
+  auto sz2 =
+      TiledArray::size_of<TiledArray::MemorySpace::Host>(TTD(Range(2, 3, 4)));
+  BOOST_REQUIRE(sz2 == sizeof(TTD) + 2 * 3 * 4 * sizeof(TTD::value_type));
+
+  TTD ttd(Range(2, 3, 4));
+  ttd(0, 0, 0) = TensorD(Range(5, 6));
+  auto sz3 = TiledArray::size_of<TiledArray::MemorySpace::Host>(ttd);
+  BOOST_REQUIRE(sz3 == sizeof(TTD) + 2 * 3 * 4 * sizeof(TTD::value_type) +
+                           5 * 6 * sizeof(TTD::value_type::value_type));
+}
+
 BOOST_AUTO_TEST_SUITE_END()