Merge pull request #516 from ValeevGroup/evaleev/feature/size-of

evaleev · web-flow · commit 04b635a967f4 · 2025-03-27T21:15:17.000-04:00
`TA::size_of`
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -31,7 +31,7 @@ jobs:
       CCACHE_COMPRESS : true
       CCACHE_COMPRESSLEVEL : 6
       OMPI_MCA_btl_vader_single_copy_mechanism : none
-      PARSEC_MCA_runtime_bind_threads : 0
+      PARSEC_MCA_bind_threads : 0
       BUILD_CONFIG : >
         -DMADNESS_TASK_BACKEND=${{ matrix.task_backend }}
         -DCMAKE_BUILD_TYPE=${{ matrix.build_type }}
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -38,6 +38,7 @@ TiledArray/error.h
 TiledArray/initialize.h
 TiledArray/perm_index.h
 TiledArray/permutation.h
+TiledArray/platform.h
 TiledArray/proc_grid.h
 TiledArray/range.h
 TiledArray/range1.h
@@ -239,7 +240,6 @@ if(TILEDARRAY_HAS_HIP OR TILEDARRAY_HAS_CUDA)
           TiledArray/device/kernel/reduce_kernel.h
           TiledArray/device/kernel/thrust/mult_kernel.h
           TiledArray/device/kernel/thrust/reduce_kernel.h
-          TiledArray/device/platform.h
           TiledArray/device/thrust.h
           TiledArray/device/um_storage.h
   )
diff --git a/src/TiledArray/dense_shape.h b/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;
 }
diff --git a/src/TiledArray/device/btas.h b/src/TiledArray/device/btas.h
@@ -37,9 +37,9 @@
 
 #include <TiledArray/device/kernel/mult_kernel.h>
 #include <TiledArray/device/kernel/reduce_kernel.h>
-#include <TiledArray/device/platform.h>
 #include <TiledArray/device/um_storage.h>
 #include <TiledArray/math/gemm_helper.h>
+#include <TiledArray/platform.h>
 
 namespace TiledArray {
 
diff --git a/src/TiledArray/device/cpu_cuda_vector.h b/src/TiledArray/device/cpu_cuda_vector.h
@@ -4,8 +4,8 @@
 
 #include <btas/array_adaptor.h>
 
-#include <TiledArray/device/platform.h>
 #include <TiledArray/device/thrust.h>
+#include <TiledArray/platform.h>
 
 #include <madness/world/archive.h>
 
@@ -213,18 +213,18 @@ struct ArchiveLoadImpl<Archive, TiledArray::cpu_cuda_vector<T>> {
   static inline void load(const Archive& ar,
                           TiledArray::cpu_cuda_vector<T>& x) {
     typename TiledArray::cpu_cuda_vector<T>::size_type n(0);
-    ar& n;
+    ar & n;
     x.resize(n);
-    for (auto& xi : x) ar& xi;
+    for (auto& xi : x) ar & xi;
   }
 };
 
 template <class Archive, typename T>
 struct ArchiveStoreImpl<Archive, TiledArray::cpu_cuda_vector<T>> {
   static inline void store(const Archive& ar,
                            const TiledArray::cpu_cuda_vector<T>& x) {
-    ar& x.size();
-    for (const auto& xi : x) ar& xi;
+    ar & x.size();
+    for (const auto& xi : x) ar & xi;
   }
 };
 
diff --git a/src/TiledArray/device/um_storage.h b/src/TiledArray/device/um_storage.h
@@ -31,7 +31,7 @@
 #include <btas/array_adaptor.h>
 #include <btas/varray/varray.h>
 
-#include <TiledArray/device/platform.h>
+#include <TiledArray/platform.h>
 #include <TiledArray/utility.h>
 
 #include <madness/world/archive.h>
diff --git a/src/TiledArray/dist_array.h b/src/TiledArray/dist_array.h
@@ -1743,6 +1743,20 @@ class DistArray : public madness::archive::ParallelSerializableObject {
 
 };  // class DistArray
 
+/// \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;
+  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>;
diff --git a/src/TiledArray/platform.h b/src/TiledArray/platform.h
@@ -21,8 +21,12 @@
  *
  */
 
-#ifndef TILEDARRAY_DEVICE_PLATFORM_H__INCLUDED
-#define TILEDARRAY_DEVICE_PLATFORM_H__INCLUDED
+#ifndef TILEDARRAY_PLATFORM_H__INCLUDED
+#define TILEDARRAY_PLATFORM_H__INCLUDED
+
+#include <TiledArray/fwd.h>
+
+#include <TiledArray/type_traits.h>
 
 namespace TiledArray {
 
@@ -54,6 +58,40 @@ constexpr bool overlap(MemorySpace space1, MemorySpace space2) {
   return (space1 & space2) != MemorySpace::Null;
 }
 
+// customization point: is_constexpr_size_of_v<S,T> reports whether
+// size_of<S>(T) is the same for all T
+template <MemorySpace S, typename T>
+inline constexpr bool is_constexpr_size_of_v = detail::is_numeric_v<T>;
+
+// customization point: size_of<S>(O) -> std::size_t reports the number of
+// bytes occupied by O in S
+template <MemorySpace S, typename T,
+          typename = std::enable_if_t<is_constexpr_size_of_v<S, T>>>
+constexpr std::size_t size_of(const T& t) {
+  return sizeof(T);
+}
+
+// customization point: allocates_memory_space<S>(A) -> bool reports whether
+// allocator A allocates memory in space S
+template <MemorySpace S, typename T>
+constexpr bool allocates_memory_space(const std::allocator<T>& a) {
+  return S == MemorySpace::Host;
+}
+template <MemorySpace S, typename T>
+constexpr bool allocates_memory_space(const Eigen::aligned_allocator<T>& a) {
+  return S == MemorySpace::Host;
+}
+template <MemorySpace S, typename T>
+constexpr bool allocates_memory_space(const host_allocator<T>& a) {
+  return S == MemorySpace::Host;
+}
+#ifdef TILEDARRAY_HAS_DEVICE
+template <MemorySpace S, typename T>
+constexpr bool allocates_memory_space(const device_um_allocator<T>& a) {
+  return S == MemorySpace::Device_UM;
+}
+#endif
+
 /// enumerates the execution spaces
 enum class ExecutionSpace { Host, Device };
 
@@ -62,4 +100,4 @@ enum class ExecutionSpace { Host, Device };
 
 }  // namespace TiledArray
 
-#endif  // TILEDARRAY_DEVICE_PLATFORM_H__INCLUDED
+#endif  // TILEDARRAY_PLATFORM_H__INCLUDED
diff --git a/src/TiledArray/range.h b/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
diff --git a/src/TiledArray/range1.h b/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) {
diff --git a/src/TiledArray/sparse_shape.h b/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,33 @@ 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) {
+  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
 
 /// \tparam T the numeric type supporting the type of \c shape
diff --git a/src/TiledArray/tensor/tensor.h b/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,34 @@ class Tensor {
 
 };  // class Tensor
 
+/// \return the number of bytes used by \p t in memory space
+/// `S`
+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, T>) {
+        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_ =
diff --git a/src/TiledArray/tiled_range.h b/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.
diff --git a/src/TiledArray/tiled_range1.h b/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.
diff --git a/src/TiledArray/util/vector.h b/src/TiledArray/util/vector.h
diff --git a/tests/dist_array.cpp b/tests/dist_array.cpp
diff --git a/tests/tensor.cpp b/tests/tensor.cpp
