feat: size_of over all live DistArrays of a type in a World

Adds a ground-truth memory-accounting facility that finds every live
DistArray of a given type by walking a World's WorldObject registry,
rather than summing size_of over a set of handles. Because each array's
tile storage is a single DistributedStorage WorldObject, an array
referenced by N shallow-copy handles is counted exactly once -- handle
summation double-counts shared storage, which makes it unsuitable as a
cross-check.

API (in TiledArray namespace, dist_array.h):
  - size_of<DistArrayT, S>(World&)
      Walks world.get_object_ids(), recovers each registered pointer as
      the common polymorphic base madness::WorldObjectBase, and
      dynamic_casts to the DistributedStorage matching DistArrayT's tile
      type; non-matching objects are skipped. Sums
      DistributedStorage::local_size_in_bytes<S>().
  - size_of_live_distarrays<S, DistArrayTs...>(worlds)
      Returns a [world][type] matrix of the above.

DistributedStorage gains local_size_in_bytes<S>(): sums
TiledArray::size_of<S>(tile) over locally-owned, set futures -- the same
tile set size_of(DistArray) iterates.

Type-safety rests on WorldObjectBase sitting at offset 0 of every
registered WorldObject; verified that across MADNESS, TiledArray, and
MPQC no WorldObject-derived class has WorldObject as a non-primary base
(the single-inheritance "class X : public WorldObject<X>" idiom). The
recovered base is dynamic_cast, so a wrong type yields nullptr, not UB.

Counts only locally-owned set tiles; excludes TiledRange/shape metadata
and remote-tile caches. Call at a quiescent point (after a fence).

Test (array_suite/size_of_live_in_world): builds two distinct arrays
plus a shallow copy of one, checks the World walk equals the two-array
(deduplicated) tile total -- not the three-handle sum -- that a ToT-typed
walk does not pick up regular-tensor arrays, and the variadic matrix
form. Passes at np=1; np=2 left to CI (the local macOS ASan+mpirun
environment crashes all tests, touched or not).

Author: evaleev

src/TiledArray/dist_array.h

@@ -2015,6 +2015,64 @@ std::size_t size_of(const DistArray<Tile, Policy>& da) {
   return result;
 }
 
+/// \return the per-rank bytes (in memory space `S`) held by *all* live
+/// `DistArray<Tile,Policy>` of the requested type currently registered in
+/// \p world, discovered by walking the World's `WorldObject` registry.
+///
+/// Each array's tile storage is a single `detail::DistributedStorage`
+/// `WorldObject`, so an array referenced by N shallow-copy handles is counted
+/// exactly once — unlike summing `size_of` over a set of handles, which
+/// double-counts shared storage. This makes the result suitable as ground
+/// truth for validating handle-based accounting.
+///
+/// Discovery is type-safe: each registered pointer is recovered as the common
+/// polymorphic base `madness::WorldObjectBase` and `dynamic_cast` to the
+/// `DistributedStorage` matching `DistArrayT`'s tile type; non-matching
+/// objects (other tile types, MADNESS containers) are skipped. Assumes the
+/// registered `WorldObject`s place `WorldObjectBase` at offset 0 (true for
+/// the single-inheritance `class X : public WorldObject<X>` idiom TA uses).
+///
+/// \note Counts only locally-owned tiles whose futures are set — the same set
+/// `size_of(DistArray)` iterates. Excludes TiledRange/shape metadata and
+/// remote-tile caches. Call at a quiescent point (after a fence).
+/// \tparam DistArrayT the `DistArray` specialization to look for
+/// \tparam S the memory space to report (default `Host`)
+template <typename DistArrayT, MemorySpace S = MemorySpace::Host>
+std::size_t size_of(World& world) {
+  using tile_type = typename DistArrayT::value_type;
+  using storage_type = detail::DistributedStorage<tile_type>;
+  std::size_t result = 0;
+  for (const auto& id : world.get_object_ids()) {
+    auto base_opt = world.template ptr_from_id<madness::WorldObjectBase>(id);
+    if (!base_opt || !*base_opt) continue;
+    if (auto* storage = dynamic_cast<storage_type*>(*base_opt)) {
+      result += storage->template local_size_in_bytes<S>();
+    }
+  }
+  return result;
+}
+
+/// \return a matrix of per-rank live-`DistArray` byte totals indexed
+/// `[world_index][type_index]`: for each `World` in \p worlds (rows) and each
+/// `DistArray` type in the pack `DistArrayTs` (columns), the value of
+/// `size_of<DistArrayT, S>(world)`. Lets a caller inventory which array types
+/// hold how much in which world at a checkpoint, deduplicated across
+/// shallow-copy handles.
+///
+/// \note `S` is the leading template argument (it has a default but precedes
+/// the type pack), so callers must spell it out:
+/// `size_of_live_distarrays<MemorySpace::Host, ArrayA, ArrayB>(worlds)`.
+template <MemorySpace S = MemorySpace::Host, typename... DistArrayTs>
+std::vector<std::array<std::size_t, sizeof...(DistArrayTs)>>
+size_of_live_distarrays(const std::vector<World*>& worlds) {
+  std::vector<std::array<std::size_t, sizeof...(DistArrayTs)>> result;
+  result.reserve(worlds.size());
+  for (World* w : worlds) {
+    result.push_back({size_of<DistArrayTs, S>(*w)...});
+  }
+  return result;
+}
+
 #ifndef TILEDARRAY_HEADER_ONLY
 
 extern template class DistArray<Tensor<double>, DensePolicy>;

src/TiledArray/distributed_storage.h

@@ -20,6 +20,7 @@
 #ifndef TILEDARRAY_DISTRIBUTED_STORAGE_H__INCLUDED
 #define TILEDARRAY_DISTRIBUTED_STORAGE_H__INCLUDED
 
+#include <TiledArray/platform.h>
 #include <TiledArray/pmap/pmap.h>
 
 namespace TiledArray {
@@ -360,6 +361,25 @@ class DistributedStorage : public madness::WorldObject<DistributedStorage<T>> {
   /// \throw nothing
   size_type size() const { return data_.size(); }
 
+  /// Number of bytes occupied by the locally-owned tiles in memory space \p S.
+
+  /// Sums `TiledArray::size_of<S>(tile)` over every locally-stored element
+  /// whose future is already set; pending (unset) and remote-cached elements
+  /// are skipped. No communication; intended to be called at a quiescent
+  /// point (e.g. after a fence). This is the per-rank local footprint, the
+  /// same set of tiles that `size_of(DistArray)` iterates.
+  /// \tparam S the memory space to report
+  /// \return local tile bytes in space \p S
+  template <MemorySpace S>
+  std::size_t local_size_in_bytes() const {
+    std::size_t result = 0;
+    for (auto it = data_.begin(); it != data_.end(); ++it) {
+      const future& f = it->second;
+      if (f.probe()) result += TiledArray::size_of<S>(f.get());
+    }
+    return result;
+  }
+
   /// Max size accessor
 
   /// The maximum size is the total number of elements that can be held by

tests/dist_array.cpp

@@ -1020,6 +1020,65 @@ BOOST_AUTO_TEST_CASE(size_of) {
   BOOST_REQUIRE(sz0 == sz0_expected);
 }
 
+BOOST_AUTO_TEST_CASE(size_of_live_in_world) {
+  using T = Tensor<double>;
+  using ToT = Tensor<T>;
+  using Policy = SparsePolicy;
+  using ArrayT = DistArray<T, Policy>;
+  using ArrayToT = DistArray<ToT, Policy>;
+
+  auto& world = get_default_world();
+  world.gop.fence();
+
+  // arrays from earlier test cases may still be registered (destruction is
+  // deferred to the next fence), so measure a baseline and compare deltas
+  auto const base_T = TiledArray::size_of<ArrayT>(world);
+  auto const base_ToT = TiledArray::size_of<ArrayToT>(world);
+
+  TiledRange const trange({{0, 2, 5, 7}, {0, 5, 7, 10, 12}});
+
+  // two distinct regular arrays
+  auto a1 = make_array<ArrayT>(world, trange, [](T& tile, Range const& rng) {
+    tile = T(rng, 1.0);
+    return tile.norm();
+  });
+  auto a2 = make_array<ArrayT>(world, trange, [](T& tile, Range const& rng) {
+    tile = T(rng, 2.0);
+    return tile.norm();
+  });
+  // shallow copy: shares a1's storage WorldObject, must NOT be double-counted
+  ArrayT a1_copy = a1;
+  BOOST_REQUIRE(a1_copy.trange() == a1.trange());  // keep a1_copy alive & used
+
+  world.gop.fence();
+
+  // per-array local-tile bytes = size_of(array) - size_of(shape)
+  auto tiles_only = [](ArrayT const& a) {
+    return TiledArray::size_of<MemorySpace::Host>(a) -
+           TiledArray::size_of<MemorySpace::Host>(a.shape());
+  };
+
+  // the World walk counts each distinct DistributedStorage once: a1 + a2,
+  // NOT a1 + a2 + a1_copy
+  auto const expected_T = tiles_only(a1) + tiles_only(a2);
+  auto const got_T = TiledArray::size_of<ArrayT>(world) - base_T;
+  BOOST_CHECK_EQUAL(got_T, expected_T);
+
+  // the ToT-typed walk must not pick up the regular (T) arrays
+  auto const got_ToT_delta = TiledArray::size_of<ArrayToT>(world) - base_ToT;
+  BOOST_CHECK_EQUAL(got_ToT_delta, 0u);
+
+  // variadic matrix: one world (one row), two types (two columns)
+  auto const mat =
+      TiledArray::size_of_live_distarrays<MemorySpace::Host, ArrayT, ArrayToT>(
+          std::vector<World*>{&world});
+  BOOST_REQUIRE_EQUAL(mat.size(), 1u);
+  BOOST_CHECK_EQUAL(mat[0][0], TiledArray::size_of<ArrayT>(world));
+  BOOST_CHECK_EQUAL(mat[0][1], TiledArray::size_of<ArrayToT>(world));
+
+  world.gop.fence();
+}
+
 BOOST_FIXTURE_TEST_CASE(fill_zero_sparse, ArrayFixture) {
   // construct a sparse array with some non-zero tiles and fill it
   SpArrayN as(world, tr, TiledArray::SparseShape<float>(shape_tensor, tr));