From a19dbfb5f89fb6f6e25b06c040afdb73a276c2fe Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Tue, 30 Jun 2026 22:45:38 -0400 Subject: [PATCH 01/78] eval: observe-only predicted-footprint hook on CacheManager (trace level >= 2) Add a predict_hook to CacheManager, mirroring shaped_product_hook but observe-only: consulted at each binary Product node before materialization (self-gated on the eval trace level) to emit a "Predict" line naming the op and its predicted result footprint. This names an intermediate that exhausts memory even when it dies materializing (the post-hoc Eval line never prints for such an op). The hook is propagated into the batched scratch cache; batched products carry no shaped-product hook, so they are reported as unshaped(batched). The TA backend factory make_predict_hook() builds the hook from a TAEvalContext, reusing the shaped hook's result-outer-trange computation and a ToT average-inner-extent estimate to predict the footprint. --- .../eval/backends/tiledarray/eval_context.hpp | 118 ++++++++++++++++++ SeQuant/core/eval/cache_manager.hpp | 30 +++++ SeQuant/core/eval/eval.hpp | 22 +++- 3 files changed, 168 insertions(+), 2 deletions(-) diff --git a/SeQuant/core/eval/backends/tiledarray/eval_context.hpp b/SeQuant/core/eval/backends/tiledarray/eval_context.hpp index fd5310a978..975edb15b8 100644 --- a/SeQuant/core/eval/backends/tiledarray/eval_context.hpp +++ b/SeQuant/core/eval/backends/tiledarray/eval_context.hpp @@ -6,15 +6,57 @@ #include #include #include +#include #include +#include #include +#include #include #include namespace sequant { +namespace detail { + +/// Average nested-scalar count per outer element of a ToT operand result +/// (total scalar count / dense outer-element count); 1.0 for a flat or scalar +/// operand. Used only on the gated predicted-footprint trace path, so the +/// collective size reduction it performs (size_in_bytes sums over ranks) is +/// paid only when tracing. For a ToT * T contraction the surviving inner +/// (e.g. the PNO index) comes from the ToT operand, so this factor lifts the +/// result's outer-element prediction to a full nested-size prediction. +template +[[nodiscard]] double tot_avg_inner_factor(Result const& r) { + using ToTArray = TA::DistArray, PolicyT>; + using ToTResult = ResultTensorOfTensorTA; + if (!r.is()) return 1.0; + auto const& a = r.template get(); + std::size_t outer = 1; + for (std::size_t d = 0; d < a.trange().rank(); ++d) + outer *= a.trange().dim(d).extent(); + if (outer == 0) return 1.0; + double const total_scalars = static_cast(r.size_in_bytes()) / + static_cast(sizeof(NumericT)); + return total_scalars / static_cast(outer); +} + +/// Nonzero outer-element count over a result TiledRange, honoring an optional +/// SparseShape (sum of kept-tile element volumes); the full element volume when +/// the shape is absent (dense / unconstrained). O(total tiles). +[[nodiscard]] inline std::size_t pred_outer_elems( + TA::TiledRange const& tr, + std::optional> const& shape) { + if (!shape) return tr.elements_range().volume(); + std::size_t n = 0; + for (auto const& idx : tr.tiles_range()) + if (!shape->is_zero(idx)) n += tr.make_tile_range(idx).volume(); + return n; +} + +} // namespace detail + /// \brief Backend context for the TiledArray eval backend. /// /// Holds TiledArray-specific evaluation state that travels with the @@ -131,6 +173,82 @@ struct TAEvalContext { left, right, annot, *shape, de_nest); }; } + + /// Build an observe-only predicted-footprint hook for + /// CacheManager::set_predict_hook() from a TAEvalContext. + /// + /// The hook, consulted at a binary-Product node *before* materialization + /// (only when tracing), computes the result's outer TiledRange from the + /// operands, mirrors the shaping decision WITHOUT materializing (consulting + /// the result_shape_provider only when \p shapeable -- i.e. when the same + /// cache actually carries a shaped-product hook; the batched scratch does + /// not, so its products are predicted dense), estimates the result footprint + /// (outer-element count from the trange/shape times the ToT operand's average + /// inner extent), and emits a "Predict" trace line. It never alters the + /// result. Template parameters match make_hook(). + template > + static CacheManager>::predict_hook_type + make_predict_hook(TAEvalContext const& ctx) { + auto provider = ctx.result_shape_provider; + return [provider = std::move(provider)]( + std::any const& node_any, Result const& left, + Result const& right, std::array const& annot, + bool shapeable) -> void { + if (Logger::instance().eval.level < 2) return; + + auto const& node = + std::any_cast< + std::reference_wrapper const>>( + node_any) + .get(); + + using FlatArray = TA::DistArray, PolicyT>; + using ToTArray = TA::DistArray, PolicyT>; + using FlatResult = ResultTensorTA; + using ToTResult = ResultTensorOfTensorTA; + auto is_tensor_like = [](Result const& r) { + return r.is() || r.is(); + }; + if (!is_tensor_like(left) || !is_tensor_like(right)) return; + + auto const trange = + result_outer_trange_from_results( + left, right, annot); + + // Mirror the shaping decision without materializing: a shape is consulted + // (and reported) only when this cache will actually try to shape it. + std::optional> shape = + (shapeable && provider) ? provider(node, trange) : std::nullopt; + + double const inner = + node->tot() + ? std::max( + detail::tot_avg_inner_factor( + left), + detail::tot_avg_inner_factor( + right)) + : 1.0; + std::size_t const outer = detail::pred_outer_elems(trange, shape); + auto const& tiles = trange.tiles_range(); + std::size_t total_tiles = tiles.volume(), nz_tiles = total_tiles; + if (shape) { + nz_tiles = 0; + for (auto const& idx : tiles) + if (!shape->is_zero(idx)) ++nz_tiles; + } + std::size_t const pred_bytes = + static_cast(static_cast(outer) * inner * + static_cast(sizeof(NumericT))); + + char const* const shape_str = !shapeable ? "unshaped(batched)" + : shape ? "SHAPED" + : "plain"; + write_log(Logger::instance(), "Predict", " | ", node->label(), + " | shape=", shape_str, " | nnz_tiles=", nz_tiles, "/", + total_tiles, " | pred_result=", pred_bytes, "B", '\n'); + }; + } }; } // namespace sequant diff --git a/SeQuant/core/eval/cache_manager.hpp b/SeQuant/core/eval/cache_manager.hpp index 92fcdb8052..08e321488f 100644 --- a/SeQuant/core/eval/cache_manager.hpp +++ b/SeQuant/core/eval/cache_manager.hpp @@ -70,6 +70,23 @@ class CacheManager { std::any const& node, Result const& left, Result const& right, std::array const& annot)>; + /// An observe-only predicted-footprint hook. `evaluate()` consults it (if + /// set, and only when tracing) at each binary-Product node *before* the + /// product is computed. Unlike shaped_product_hook it NEVER replaces the + /// result: it emits a pre-materialization "Predict" trace line giving the + /// predicted result footprint (and, when \p shapeable, the result-shape + /// decision), so an op that exhausts memory is still named in the log -- the + /// post-hoc Eval line never prints for an OOMing op. \p shapeable tells the + /// hook whether this cache will actually try to shape the product (true iff a + /// shaped-product hook is set on the same cache): the batched scratch carries + /// the predict hook but no shaped hook, so its products materialize dense and + /// must be predicted as dense. All backend-specific types (TA tranges, + /// shapes) stay inside the hook's closure; the generic CacheManager and eval + /// see only Result/std::any. + using predict_hook_type = std::function const& annot, bool shapeable)>; + private: using hasher_type = TreeNodeHasher; using comparator_type = TreeNodeEqualityComparator; @@ -169,6 +186,8 @@ class CacheManager { shaped_product_hook_type shaped_product_hook_{}; + predict_hook_type predict_hook_{}; + public: /// Sets the custom evaluator (see custom_evaluator_type). Pass an empty /// std::function to clear it. @@ -193,6 +212,17 @@ class CacheManager { return shaped_product_hook_; } + /// Sets the predicted-footprint hook (see predict_hook_type). Pass an empty + /// std::function to clear it. + void set_predict_hook(predict_hook_type fn) noexcept { + predict_hook_ = std::move(fn); + } + + /// \return the predicted-footprint hook (empty if none is set). + [[nodiscard]] predict_hook_type const& predict_hook() const noexcept { + return predict_hook_; + } + /// Default persistence classifier: every entry is non-persistent (NP). struct all_non_persistent { bool operator()(key_type const&) const noexcept { return false; } diff --git a/SeQuant/core/eval/eval.hpp b/SeQuant/core/eval/eval.hpp index 6b3a21bf75..c34387fa7a 100644 --- a/SeQuant/core/eval/eval.hpp +++ b/SeQuant/core/eval/eval.hpp @@ -612,6 +612,19 @@ ResultPtr evaluate(Node const& node, // [&]() { result = left->sum(*right, ann); }); } else { SEQUANT_ASSERT(node->op_type() == EvalOp::Product); + // Observe-only predicted-footprint trace: consulted BEFORE materializing + // the product, so an op that exhausts memory is still named in the log + // (the post-hoc Eval line below never prints for an OOMing op). Gated on + // trace + level so it is inert by default; the batched scratch carries + // this hook too (see make_batched_scratch), so batched products -- which + // do NOT consult the shaped-product hook -- are still predicted (as + // dense, via the shapeable=false argument). + if constexpr (detail::trace(EvalTrace)) { + if (auto const& ph = cache.predict_hook(); + ph && log::printing() && Logger::instance().eval.level >= 2) + ph(std::any{std::cref(node)}, *left, *right, ann, + static_cast(cache.shaped_product_hook())); + } // Consult the shaped-product hook (if set) before evaluating the product. // The hook receives the node (wrapped in a std::any as a // std::reference_wrapper so the full IR node is inspectable) plus the @@ -1127,8 +1140,13 @@ template auto is_persistent = [seed_keys = std::move(seed_keys)](TreeNode const& n) { return seed_keys.contains(n); }; - return {CacheManager{std::move(reg), std::move(is_persistent)}, - std::move(seeds)}; + CacheManager scratch{std::move(reg), std::move(is_persistent)}; + // Carry the observe-only predicted-footprint hook into the scratch so batched + // products are predicted too; deliberately NOT the shaped-product hook -- + // batched products stay unshaped (existing behavior), which the predict hook + // reflects via shapeable=false (it has no shaped hook here). + scratch.set_predict_hook(real.predict_hook()); + return {std::move(scratch), std::move(seeds)}; } } // namespace detail From 691bb427d370b7b5711b945add08887305dbcff2 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Wed, 1 Jul 2026 19:52:58 -0400 Subject: [PATCH 02/78] eval: optional backend heap-stats suffix on the eval trace line Add Logger::eval.heap_stats (std::function); when set, its (already rank-reduced) return is appended to each Eval line after rss=, and omitted entirely when unset. Lets a backend report allocator-level memory (e.g. glibc all-arena in-use vs system) right at an RSS jump, to separate live heap from retained-free heap. --- SeQuant/core/eval/eval.hpp | 42 +++++++++++++++++++++++--------------- SeQuant/core/logger.hpp | 13 +++++++++++- 2 files changed, 38 insertions(+), 17 deletions(-) diff --git a/SeQuant/core/eval/eval.hpp b/SeQuant/core/eval/eval.hpp index c34387fa7a..ecbd4fca53 100644 --- a/SeQuant/core/eval/eval.hpp +++ b/SeQuant/core/eval/eval.hpp @@ -198,7 +198,7 @@ enum struct TermMode { Begin, End }; /// One log record per eval op. Line format: /// // clang-format off -/// Eval | |