From 9f938f157dfbd26ec77518a56af80906d9f22573 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Wed, 29 Jul 2026 12:55:26 -0400 Subject: [PATCH 001/213] eval: dry-run cost-profile prediction backend CostProfile (peak/flops/exec) over the factorized IR via a zero-data dry-run evaluation, driving the batched cost model's predictions. --- .../backends/dryrun/cost_model_object.hpp | 142 +++++ .../eval/backends/dryrun/cost_profile.hpp | 327 ++++++++++ .../core/eval/backends/dryrun/eval_expr.hpp | 87 +++ SeQuant/core/eval/backends/dryrun/result.hpp | 586 ++++++++++++++++++ .../core/eval/backends/dryrun/size_regime.hpp | 79 +++ 5 files changed, 1221 insertions(+) create mode 100644 SeQuant/core/eval/backends/dryrun/cost_model_object.hpp create mode 100644 SeQuant/core/eval/backends/dryrun/cost_profile.hpp create mode 100644 SeQuant/core/eval/backends/dryrun/eval_expr.hpp create mode 100644 SeQuant/core/eval/backends/dryrun/result.hpp create mode 100644 SeQuant/core/eval/backends/dryrun/size_regime.hpp diff --git a/SeQuant/core/eval/backends/dryrun/cost_model_object.hpp b/SeQuant/core/eval/backends/dryrun/cost_model_object.hpp new file mode 100644 index 0000000000..48cc230d50 --- /dev/null +++ b/SeQuant/core/eval/backends/dryrun/cost_model_object.hpp @@ -0,0 +1,142 @@ +#ifndef SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_COST_MODEL_OBJECT_HPP +#define SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_COST_MODEL_OBJECT_HPP + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace sequant::eval::dryrun { + +/// Per-index extent OVERRIDE table: narrows specific indices (by identity, so +/// it survives reshaping across prod/sum/permute -- the same shared/ +/// contracted Index object may occupy different tensor modes at different +/// nodes) to a runtime-realized element count. Populated by +/// Result::slice_mode()/mode_batches() call sites (see result.hpp); empty => +/// no override, the regime's nominal extent applies. This table -- not a +/// second cost model -- is what lets a zero-data DryRun Result report the +/// REALIZED (possibly runtime-sliced) size rather than always the full +/// regime extent, which is exactly the signal Task 6's replay witnesses. +using ExtentOverrides = container::map; + +/// +/// \brief Bundles the optimizer's own cost closures (memsize/flops/roofline) +/// behind one value type so dry-run Results report MODEL size (not an +/// allocated size), and the harness can additionally read FLOPs and +/// projected execution cost per operation. +/// +/// This is a thin wrapper: all arithmetic is delegated verbatim to +/// \c sequant::opt::detail::memsize_counter / \c flops_counter / \c +/// roofline_op_cost (see \c core/optimize/single_term_detail.hpp and \c +/// core/optimize/cost_model.hpp) -- no parallel cost model is implemented +/// here. The only thing this class adds is the ExtentOverrides indirection: +/// each query builds a fresh (cheap; no heap allocation beyond the closure +/// itself) index-to-extent callable that consults \p overrides before +/// falling back to the SizeRegime's nominal extent, then hands that callable +/// to the counter. +/// +class CostModel { + public: + explicit CostModel(SizeRegime regime, RooflineParams roofline = {}) + : regime_{std::move(regime)}, roofline_{roofline} {} + + /// + /// \brief Bytes for a tensor with these (literal, canon-order) indices, + /// honoring any per-index extent override (a runtime slice_mode()/ + /// mode_batches() narrowing). + /// + /// Delegates the extent-product / composite-moment math to \c + /// memsize_counter, invoked with \p idxset as the sole (`lhs`) operand and + /// empty `rhs`/`result` -- an empty operand's tot_indices() split + /// accumulates the starting product of 1.0, which memsize_counter itself + /// special-cases to contribute zero bytes, so this reproduces exactly the + /// single-operand byte count \c memsize_counter is designed to report per + /// operand. + /// + [[nodiscard]] std::size_t memsize( + container::svector const& idxset, + ExtentOverrides const& overrides = {}) const { + auto const ext = make_extent_fn(overrides); + auto const mc = + sequant::opt::detail::memsize_counter(ext, regime_.inner_pow_fn()); + double const elems = + mc(idxset, container::svector{}, container::svector{}); + return static_cast(elems * numeric_size_); + } + + /// + /// \brief Multiply-add count for a contraction whose free (result) indices + /// are \p out and whose contracted (summed-over) indices are + /// \p contracted. + /// + /// Delegates to \c flops_counter, which prices the union of its (lhs, rhs, + /// result) arguments; passing (\p out, \p contracted, {}) makes that union + /// exactly `out U contracted` -- the full index set touched by the + /// contraction, since by construction `contracted` holds precisely the + /// indices present in both operands but absent from the result. + /// + [[nodiscard]] double flops(container::svector const& out, + container::svector const& contracted, + ExtentOverrides const& overrides = {}) const { + auto const ext = make_extent_fn(overrides); + auto const fc = + sequant::opt::detail::flops_counter(ext, regime_.inner_pow_fn()); + return fc(out, contracted, container::svector{}); + } + + /// + /// \brief Roofline-projected execution cost of one contraction (see + /// \c sequant::opt::detail::roofline_op_cost). + /// + /// \p left_bytes / \p right_bytes are operand footprints in BYTES (as + /// reported by \c Result::size_in_bytes()); converted to elements (the + /// counter's native unit) via \c numeric_size before delegating. + /// + [[nodiscard]] double exec_cost(double flops_count, std::size_t left_bytes, + std::size_t right_bytes) const { + double const traffic_elems = + static_cast(left_bytes + right_bytes) / numeric_size_; + return sequant::opt::detail::roofline_op_cost( + flops_count, traffic_elems, roofline_.machine_balance, + roofline_.fast_mem_elems, roofline_.block_tiles, + roofline_.block_prefactor); + } + + [[nodiscard]] SizeRegime const& regime() const noexcept { return regime_; } + + private: + // Index-to-extent callable consulting `overrides` first, else the + // regime's nominal extent. The returned std::function captures `overrides` + // (and `this`) BY REFERENCE and is only ever used -- never stored -- + // within the (memsize/flops) call that constructs it, so the reference + // stays valid for its entire lifetime. Explicit (non-deduced) return type + // so this can be called from memsize()/flops(), which appear earlier in + // the class body (a deduced `auto` return type would require the + // definition to precede every use, even within the same class). + [[nodiscard]] std::function make_extent_fn( + ExtentOverrides const& overrides) const { + return [this, &overrides](Index const& ix) -> std::size_t { + if (auto it = overrides.find(ix); it != overrides.end()) + return it->second; + return regime_.extent(ix); + }; + } + + SizeRegime regime_; + RooflineParams roofline_; + // sizeof(double); see doc/dev/plans/2026-07-04-dryrun-eval-backend.md Task 2 + // note on OptimizeOptions::numeric_size (hardcoded here, matching the C60 + // trace's real-only CSV-CCk path; complex CSV-CCk is out of scope, see the + // plan's carried-minor N4). + double numeric_size_ = 8.0; +}; + +} // namespace sequant::eval::dryrun + +#endif // SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_COST_MODEL_OBJECT_HPP diff --git a/SeQuant/core/eval/backends/dryrun/cost_profile.hpp b/SeQuant/core/eval/backends/dryrun/cost_profile.hpp new file mode 100644 index 0000000000..45584753c9 --- /dev/null +++ b/SeQuant/core/eval/backends/dryrun/cost_profile.hpp @@ -0,0 +1,327 @@ +#ifndef SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_COST_PROFILE_HPP +#define SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_COST_PROFILE_HPP + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace sequant::eval::dryrun { + +/// Configuration for a faithful (gated) dry-run cache: the same footprint gate +/// and free-batchable-axis veto the real batched eval loop applies, so a +/// free-batchable giant (a `mu~`/`K`-carrying DF intermediate) is NOT cached +/// whole but recomputed sliced under each consumer's batch trigger. +/// +/// The element types mirror the gated \c sequant::cache_manager overload +/// (\c cache_manager.hpp): \c is_volatile is invoked on every \c TreeNode +/// (deduced as \c EvalNodeDryRun for the dry-run backend), \c +/// is_batchable_index on every result \c Index. +/// +/// This struct lives here (rather than only in the test) because Task 4's +/// \c cost_profile() entry point consumes it. +struct CacheConfig { + /// Footprint gate (bytes): a node whose result footprint exceeds this is not + /// cached. 0 (default) disables the gate. + double max_footprint = 0.; + /// Minimum non-persistent repeats to cache an internal node (CSE rule). + std::size_t min_repeats = 2; + /// `bool(EvalNodeDryRun const&)`: true if the node is intrinsically volatile + /// (typically the amplitude leaves). Empty => nothing is volatile. + std::function is_volatile; + /// `bool(Index const&)`: true for an index the runtime batched evaluator + /// slices over (e.g. DF aux `K` / PAO `mu~`). A node whose result carries + /// such a free index is vetoed from caching. Empty => nothing is batchable. + /// + /// ADVISORY when passed to \c cost_profile(): that entry point OVERWRITES + /// this field with \c policy.is_batchable_contracted_index (the + /// CONTRACTED-role building block) before building the cache. The cache veto + /// is contracted-stamp-only, so it is decoupled from the replay evaluator's + /// accept (the derived role union). Only \c build_dryrun_cache() called + /// directly honors this field as-is. + std::function is_batchable_index; +}; + +/// Builds a gated dry-run cache from an eval-node range, a \p cfg, and a +/// \p regime that supplies the moment-aware node-size model used for the +/// footprint gate. +/// +/// The footprint functor sizes a node's result (its \c canon_indices()) with +/// the SAME moment-aware counter the DryRun \c Result uses +/// (\c memsize_counter over \c regime.idx_to_extent()/inner_pow_fn()), scaled +/// to bytes, so the gate compares like-for-like against \c cfg.max_footprint. +/// +/// Unlike the SIMPLE \c cache_manager(nodes) factory the ad-hoc dry-run test +/// sites use, this routes through the GATED overload so free-batchable-axis +/// giants are vetoed (matching the real run). Call \c CacheManager::reset() on +/// the returned cache between summands to drop per-term non-persistent scratch +/// while keeping persistent (cross-term) entries. +/// +/// \param nodes the evaluation forest (a range of \c EvalNodeDryRun). +/// \param cfg footprint/repeat/volatility/batchability configuration. +/// \param regime the size regime supplying extents and CSV moment tables. +/// \return a \c CacheManager over \c EvalNodeDryRun. +template +auto build_dryrun_cache(NodeRange const& nodes, CacheConfig const& cfg, + SizeRegime const& regime) { + auto memsize = sequant::opt::detail::memsize_counter(regime.idx_to_extent(), + regime.inner_pow_fn()); + + // Footprint (bytes) of a node's RESULT: canon_indices() fed to the + // moment-aware counter (as the counter's `result` slot; the empty lhs/rhs + // contribute nothing) times 8 bytes/element. Same arithmetic as the DryRun + // Result::size_in_bytes(), so the gate is faithful. + auto footprint_of = + [memsize = std::move(memsize)](EvalNodeDryRun const& n) -> double { + std::vector const result(n->canon_indices().begin(), + n->canon_indices().end()); + return memsize(std::vector{}, std::vector{}, result) * 8.0; + }; + + // Default the predicates so the gated factory never invokes an empty + // std::function (nothing volatile / nothing batchable leaves those gates + // inert, matching the factory's own defaults). + std::function is_volatile = + cfg.is_volatile ? cfg.is_volatile + : std::function( + [](EvalNodeDryRun const&) { return false; }); + std::function is_batchable_index = + cfg.is_batchable_index ? cfg.is_batchable_index + : std::function( + [](Index const&) { return false; }); + + // Note: the gated sequant::cache_manager() overload below stamps the + // cross-occurrence lifetime mask on `nodes` itself before its DAG walk / + // veto (cache_manager.hpp), so this call site does not need to do so. + return sequant::cache_manager(nodes, std::move(is_volatile), cfg.min_repeats, + std::move(footprint_of), cfg.max_footprint, + std::move(is_batchable_index)); +} + +/// Summary of the modeled cost of a factorized dry-run eval forest, as produced +/// by \c cost_profile(). All quantities are summed/maxed over every summand +/// tree in the forest. +struct CostProfile { + /// Predicted peak working-set (bytes): the max over summands of the + /// batched-scratch high-watermark folded by the Task-3 \c PeakSink and the + /// outer gated cache's \c working_set_hwmark(). + /// + /// This is computed as \c max(batched-inner scratch high-watermark, + /// outer cross-term cached residency), NOT their sum. When a persistent + /// cross-term cache entry co-resides in memory with a batched-inner + /// transient at the same instant, the two are actually additive, so this + /// value is a LOWER BOUND on the true peak in that case. It is exact + /// whenever one of the two terms dominates the other (e.g. the C60 4-PNO + /// \c W case, where the batched-inner scratch dwarfs any persistent + /// cross-term residency). + double peak_bytes = 0; + /// Summed unweighted static contraction FLOPs over all internal nodes. + /// NOT CSE-aware across summands: a cross-term shared intermediate is + /// walked (and its FLOPs counted) once per occurrence, not once overall. + double flops = 0; + /// Summed roofline-projected execution cost over all internal nodes. Same + /// per-occurrence (not CSE-deduplicated) accounting caveat as \c flops. + double exec_cost = 0; + /// Number of internal (contraction) nodes across the forest. + std::size_t n_ops = 0; +}; + +/// Replays a factorized eval forest zero-data through the real eval loop -- +/// with a gated cache built from \p cfg (Task 2) and a \c PeakSink threaded +/// through the batched evaluator (Task 3) -- and, alongside, does a static walk +/// of the forest to accumulate FLOPs / roofline exec cost / op count. This is +/// the single reusable entry point both SeQuant tests and MPQC call. +/// +/// \par The printing gate +/// \c CacheManager::working_set_hwmark() only accumulates while +/// \c sequant::eval::log::printing() is true (the hwmark update sits on the +/// trace-printing path). This routine therefore FORCES the eval logger's level +/// > 0 around the replay -- discarding the narrow trace to a null sink when no +/// \p trace is requested -- and restores the previous logger state afterward, +/// so \c peak_bytes is non-zero even with no trace stream. +/// +/// \par Global state / threading +/// This routine mutates the process-global \c Logger::instance().eval state +/// (\c level and \c stream) for the duration of the replay (restored on every +/// exit path, including exceptions). Because that state is a singleton shared +/// by the whole process, \c cost_profile() MUST be called single-threaded -- +/// e.g. as a pre-flight step before, or a post-hoc step after, the real +/// multi-threaded eval -- never concurrently with other code that reads or +/// writes \c Logger::instance().eval (including another concurrent +/// \c cost_profile() call). +/// +/// \par FLOPs / exec_cost accounting +/// \c CostProfile::flops and \c CostProfile::exec_cost are accumulated by a +/// static walk that sums a contribution per BINARIZED internal node of the +/// forest; they are NOT CSE-aware across summands. A shared intermediate +/// that recurs across multiple summand trees (or multiple times within one) +/// is counted once per occurrence, not once overall -- unlike \c peak_bytes, +/// which is driven by the gated cache and so does reflect cross-term reuse. +/// +/// \param forest per-summand optimized+binarized eval forest (the real IR). +/// \param policy the batch policy driving the replay evaluator; its +/// \c is_batchable_contracted_index is COPIED over +/// \c cfg.is_batchable_index internally (the cache veto is +/// contracted-stamp-only), while the evaluator accept is the derived +/// role union \c policy.is_batchable_index(). The \c cfg field is +/// advisory here. +/// \param cfg gated-cache config (footprint gate, axis veto, volatile, +/// repeats). +/// \param regime the size regime supplying extents and CSV moment tables; +/// the internal \c CostModel and \c DryRunLeafEvaluator are built from +/// it. +/// \param trace optional per-op trace sink (nullptr = no trace). When +/// non-null, the eval loop's narrow trace is transcoded (UTF-8) into it. +/// \return the accumulated \c CostProfile. +inline CostProfile cost_profile(std::vector const& forest, + BatchPolicy const& policy, + CacheConfig const& cfg, + SizeRegime const& regime, + std::wostream* trace = nullptr) { + CostProfile profile; + + auto cm = std::make_shared(regime); + DryRunLeafEvaluator const leaf{cm}; + + // ---- static cost walk (independent of the replay) -------------------- + // For every internal node: flops = flops_counter(left, right, result); the + // roofline exec cost uses the left operand's footprint as the transferred + // bytes and the arena convention (4096) the [dryrun-costmodel] test fixes. + auto const flops_of = sequant::opt::detail::flops_counter( + regime.idx_to_extent(), regime.inner_pow_fn()); + std::function walk = + [&](EvalNodeDryRun const& n) { + if (n.leaf()) return; + profile.n_ops += 1; + double const node_flops = + flops_of(n.left()->canon_indices(), n.right()->canon_indices(), + n->canon_indices()); + profile.flops += node_flops; + container::svector const left(n.left()->canon_indices().begin(), + n.left()->canon_indices().end()); + profile.exec_cost += cm->exec_cost(node_flops, cm->memsize(left), 4096); + walk(n.left()); + walk(n.right()); + }; + for (auto const& root : forest) walk(root); + + // ---- peak replay through the real eval loop -------------------------- + // Route the two consumers by intent (they are decoupled): + // - the cache VETO takes the CONTRACTED-role building block: only a mode + // actually sliced AT a node (BatchModeType::Contracted) vetoes caching, + // so the veto must never see the union (an external-only mode is + // batch-invariant and stays cacheable). + // - the replay EVALUATOR's accept is the derived role union, applied inside + // make_evaluator(policy) via policy.is_batchable_index(). + CacheConfig local_cfg = cfg; + local_cfg.is_batchable_index = policy.is_batchable_contracted_index; + auto cache = build_dryrun_cache(forest, local_cfg, regime); + + auto& logger = Logger::instance(); + // RAII guard restoring the process-global Logger::eval state on EVERY exit + // path from this point on -- normal return, early return, or an exception + // unwinding out of the replay loop below -- not just the two trailing + // assignments a plain save/restore would rely on. Without this, a throw + // from anything in the loop OTHER than evaluate (e.g. + // std::bad_alloc from make_evaluator/set_custom_evaluator/ + // working_set_hwmark/cache.reset()) would unwind past the local + // `trace_capture` destructor while `logger.eval.stream` still points at it, + // leaving a dangling pointer in the process-global singleton with + // level == 2 still set. + struct LoggerEvalGuard { + decltype(logger.eval)& eval; + std::size_t const prev_level; + std::ostream* const prev_stream; + ~LoggerEvalGuard() { + eval.level = prev_level; + eval.stream = prev_stream; + } + } logger_eval_guard{logger.eval, logger.eval.level, logger.eval.stream}; + + // Force printing() on so working_set_hwmark() accumulates. The eval logger + // stream is narrow; capture into a narrow buffer only when a (wide) trace + // sink was requested, else discard to a null stream. + std::ostringstream trace_capture; + logger.eval.level = 2; + logger.eval.stream = trace ? &trace_capture : nullptr; + + std::atomic peak{0.0}; + for (auto const& root : forest) { + cache.set_custom_evaluator(sequant::make_evaluator( + policy, leaf, sequant::make_no_scope_guard{}, &peak)); + try { + (void)sequant::evaluate(root, leaf, cache); + } catch (std::exception const&) { + // A zero-data DryRun sizing throw must not mask the peak read. + } + // Fold the outer cached residency BEFORE reset() (which zeroes the + // hwmark). `peak` folds every batched scratch high-watermark across all + // summands via std::max, so its running load() is the global scratch peak. + profile.peak_bytes = std::max( + {profile.peak_bytes, peak.load(), double(cache.working_set_hwmark())}); + cache.reset(); // drop per-term non-persistent scratch; keep persistent + } + + // logger_eval_guard's destructor restores logger.eval.{level,stream} at + // function exit (see above); no manual restore needed here. + + // If a wide trace sink was requested, transcode the captured narrow (UTF-8) + // eval trace into it (the eval loop writes only to the narrow logger stream; + // index labels such as mu~/K are multi-byte, so a plain widen would corrupt + // them -- decode UTF-8 to code points instead). + if (trace) { + std::string const s = trace_capture.str(); + std::wstring w; + w.reserve(s.size()); + for (std::size_t i = 0; i < s.size();) { + unsigned char const c = static_cast(s[i]); + char32_t cp; + std::size_t len; + if (c < 0x80) { + cp = c; + len = 1; + } else if ((c >> 5) == 0x6) { + cp = c & 0x1Fu; + len = 2; + } else if ((c >> 4) == 0xE) { + cp = c & 0x0Fu; + len = 3; + } else if ((c >> 3) == 0x1E) { + cp = c & 0x07u; + len = 4; + } else { + cp = c; // invalid lead byte: pass through + len = 1; + } + for (std::size_t k = 1; k < len && i + k < s.size(); ++k) + cp = (cp << 6) | (static_cast(s[i + k]) & 0x3Fu); + w.push_back(static_cast(cp)); + i += len; + } + *trace << w; + } + + return profile; +} + +} // namespace sequant::eval::dryrun + +#endif // SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_COST_PROFILE_HPP diff --git a/SeQuant/core/eval/backends/dryrun/eval_expr.hpp b/SeQuant/core/eval/backends/dryrun/eval_expr.hpp new file mode 100644 index 0000000000..a7bcad51f9 --- /dev/null +++ b/SeQuant/core/eval/backends/dryrun/eval_expr.hpp @@ -0,0 +1,87 @@ +#ifndef SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_EVAL_EXPR_HPP +#define SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_EVAL_EXPR_HPP + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace sequant::eval::dryrun { + +/// +/// \brief Extends EvalExpr with an annot() method so DryRun eval nodes can be +/// evaluated. +/// +/// Unlike \c EvalExprTAPP (opaque \c int64_t hashes of index labels -- see +/// \c backends/tapp/eval_expr.hpp), DryRun's annotation IS the plain literal +/// (canon-order) index list itself: \c Result::prod/sum/permute need each +/// index's actual space/extent (via \c CostModel), not just its identity, to +/// compute a modeled size. +/// +class EvalExprDryRun final : public EvalExpr { + public: + using annot_t = dryrun::annot_t; // container::svector + + template >> + explicit EvalExprDryRun(Args&&... args) + : EvalExpr{std::forward(args)...} { + annot_ = canon_indices() | ranges::to; + } + + /// + /// \return Annotation (container::svector) for DryRun tensors. + /// + [[nodiscard]] annot_t const& annot() const noexcept { return annot_; } + + private: + annot_t annot_; +}; + +/// Type alias for DryRun evaluation nodes +using EvalNodeDryRun = EvalNode; + +static_assert(meta::eval_node); +static_assert(meta::can_evaluate); + +/// +/// \brief Leaf yielder: turns each IR leaf (a tensor/constant/variable node) +/// into a zero-data DryRun Result. This is the `F` in +/// \c evaluate(node, layout, F, cache). +/// +/// A tensor leaf's literal (canon-order) index list decides flat vs nested: +/// \c make_dryrun_result builds a flat \c ResultDryRun if none of the leaf's +/// indices are proto-indexed, or a nested \c ResultDryRunNested (a CSV/PNO +/// amplitude or coefficient) if any are -- and threads that SAME literal list +/// through as the nested result's canon-order position map, so a later +/// \c slice_mode()/\c mode_batches() call (which the batched runtime only +/// ever issues against a LEAF's result) resolves its positional `mode` +/// argument correctly regardless of the leaf's flat/nested-ness. +/// +struct DryRunLeafEvaluator { + std::shared_ptr cm; + + [[nodiscard]] ResultPtr operator()(EvalNodeDryRun const& leaf) const { + SEQUANT_ASSERT(leaf.leaf()); + if (!leaf->is_tensor()) { + // Constant / Variable leaf: a bare scalar. No real numeric value is + // ever tracked by this zero-data backend (only sizes/costs), so 1.0 is + // a placeholder never meant to be read as a physical result. + return eval_result>(1.0); + } + container::svector idx = leaf->canon_indices() | ranges::to; + return make_dryrun_result(std::move(idx), cm); + } +}; + +} // namespace sequant::eval::dryrun + +#endif // SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_EVAL_EXPR_HPP diff --git a/SeQuant/core/eval/backends/dryrun/result.hpp b/SeQuant/core/eval/backends/dryrun/result.hpp new file mode 100644 index 0000000000..03b3c8631a --- /dev/null +++ b/SeQuant/core/eval/backends/dryrun/result.hpp @@ -0,0 +1,586 @@ +#ifndef SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_RESULT_HPP +#define SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_RESULT_HPP + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace sequant::eval::dryrun { + +/// +/// \brief Annotation type DryRun's Result ops decode from the eval engine's +/// std::any [l,r,res] / [pre,post] triples/pairs. +/// +/// Unlike \c EvalExprTAPP (opaque \c int64_t index-label hashes -- see +/// \c backends/tapp/eval_expr.hpp), DryRun's annotation IS the plain literal +/// (canon-order) index list itself: \c Result::prod/sum/permute need each +/// index's actual space/extent (via \c CostModel), not just an opaque +/// identity, to compute a modeled size. +/// +using annot_t = container::svector; + +/// Per-mode assembled element coverage recorded by write_into_slice(): maps an +/// outer mode position to the contiguous `[lo, hi)` element range filled so far +/// by scattered blocks. Lets a zero-data DryRun destination report the REALIZED +/// (assembled) size along a partitioned mode and detect gaps/overlaps between +/// blocks -- the assemble-side analogue of ExtentOverrides for slice_mode(). +using AssembledCoverage = + container::map>; + +class ResultDryRun; +class ResultDryRunNested; + +/// +/// \brief Builds whichever concrete DryRun Result type matches \p idx's +/// content: a nested \c ResultDryRunNested if any index in \p idx is +/// proto-indexed (a CSV/PNO composite leg, e.g. a CSV amplitude's PNO +/// domain leg `a_1`), otherwise a flat \c ResultDryRun. +/// +/// Dispatch is by CONTENT of the decoded result annotation, not by either +/// operand's concrete type -- exactly mirroring how the real eval engine +/// itself decides tensor-of-tensor-ness (\c EvalExpr::tot(), from the same +/// proto-indexed-leg criterion). This is what lets \c prod()/sum() freely +/// combine a flat operand (e.g. a bare 3-center DF integral) with a nested +/// one (e.g. a CSV/PNO coefficient), exactly as real CSV-CCSD terms do, +/// without either side needing to know the other's concrete type. +/// +[[nodiscard]] inline ResultPtr make_dryrun_result( + container::svector idx, std::shared_ptr cm, + ExtentOverrides overrides = {}); + +namespace detail { + +[[nodiscard]] inline bool has_proto(container::svector const& idx) { + return std::any_of(idx.begin(), idx.end(), + [](Index const& ix) { return ix.has_proto_indices(); }); +} + +[[nodiscard]] inline ExtentOverrides merge_overrides(ExtentOverrides const& a, + ExtentOverrides const& b) { + ExtentOverrides out = a; + for (auto const& [ix, n] : b) out[ix] = n; + return out; +} + +// Uniform read access to a DryRun Result's (index list, overrides, cost +// model) regardless of which concrete DryRun type `r` is. `is()`/`as()` +// are public Result methods, so no friendship is needed; declared here (and +// defined below, after both concrete classes) purely because their bodies +// need the concrete classes' definitions. +[[nodiscard]] container::svector indices_of(Result const& r); +[[nodiscard]] ExtentOverrides overrides_of(Result const& r); + +/// +/// \brief Shared op bodies for the two DryRun Result concrete types. +/// +/// Both \c ResultDryRun and \c ResultDryRunNested carry exactly an (index +/// list, ExtentOverrides, CostModel) triple and differ only in what they +/// additionally expose (\c ResultDryRunNested splits its index list into +/// outer()/inner() views for CSV-composite-aware inspection/testing). +/// Implemented once here so the two classes' prod/sum/permute/slice_mode/ +/// mode_batches bodies are one-line forwards, not near-duplicated logic. +/// +struct DryRunOps { + [[nodiscard]] static ResultPtr sum(container::svector const& idx, + ExtentOverrides const& ov, + std::shared_ptr const& cm, + Result const& other, + std::array const& annot) { + auto const a = Annot{annot}; + auto merged = merge_overrides(ov, overrides_of(other)); + return make_dryrun_result( + container::svector(a.this_annot.begin(), a.this_annot.end()), cm, + std::move(merged)); + } + + [[nodiscard]] static ResultPtr prod( + container::svector const& idx, ExtentOverrides const& ov, + std::shared_ptr const& cm, Result const& other, + std::array const& annot) { + if (other.is>()) { + // Scalar * tensor: shape (and any accumulated slicing) unchanged. + return make_dryrun_result(idx, cm, ov); + } + auto const a = Annot{annot}; + auto merged = merge_overrides(ov, overrides_of(other)); + + // Emit the cost model's OWN flops / roofline exec_cost for THIS op into the + // eval trace (gated on the eval log level), interleaved right before the + // generic engine's `Eval | Product` line for the same op. This lets trace + // post-processing weight avoidable recomputation by MODELLED TIME without + // re-deriving the cost downstream (which would silently drift from the + // model). The (out, contracted) index sets and the realized (sliced) extent + // overrides feed the same CostModel closures the static cost_profile() walk + // uses, so per-op costs are consistent with the whole-forest totals. + if (Logger::instance().eval.level > 0) { + container::svector const rhs = indices_of(other); + container::svector out(a.this_annot.begin(), a.this_annot.end()); + container::svector contracted; + for (auto const& ix : idx) + if (std::find(rhs.begin(), rhs.end(), ix) != rhs.end()) + contracted.push_back(ix); + double const flops = cm->flops(out, contracted, merged); + double const exec = cm->exec_cost(flops, cm->memsize(idx, ov), 4096); + write_log(Logger::instance(), "OpCost", std::format(" | {}", flops), + std::format(" | {}", exec), '\n'); + } + + if (a.this_annot.empty()) { + // Full contraction -> scalar. No real numeric value is ever tracked by + // this zero-data backend (only sizes/costs), so the placeholder 0.0 + // is never meant to be read as a physical result. + return eval_result>(0.0); + } + return make_dryrun_result( + container::svector(a.this_annot.begin(), a.this_annot.end()), cm, + std::move(merged)); + } + + [[nodiscard]] static ResultPtr permute( + container::svector const& /*idx*/, ExtentOverrides const& ov, + std::shared_ptr const& cm, + std::array const& ann) { + auto const post = std::any_cast(ann[1]); + return make_dryrun_result( + container::svector(post.begin(), post.end()), cm, ov); + } + + [[nodiscard]] static ResultPtr slice_mode( + container::svector const& idx, ExtentOverrides const& ov, + std::shared_ptr const& cm, std::size_t mode, + std::size_t elem_lo, std::size_t elem_hi) { + SEQUANT_ASSERT(mode < idx.size()); + auto merged = ov; + merged[idx[mode]] = elem_hi - elem_lo; + return make_dryrun_result(idx, cm, std::move(merged)); + } + + /// Scatter \p block into the `[block_lo, block_hi)` element slice of the + /// destination's mode \p mode -- the inverse of slice_mode(). Zero-data: + /// updates only the destination's modelled size and assembled-coverage + /// bookkeeping. \p ov and \p cov are the destination's (mutated in place). + static void write_into_slice(container::svector const& idx, + ExtentOverrides& ov, AssembledCoverage& cov, + std::shared_ptr const& cm, + Result const& block, std::size_t mode, + std::size_t block_lo, std::size_t block_hi) { + SEQUANT_ASSERT(mode < idx.size()); + SEQUANT_ASSERT(block_lo < block_hi); + Index const& mix = idx[mode]; + // Tile/width consistency: the block's own modelled extent on the shared + // mode index must equal the slice width it is being written into. + auto const bov = overrides_of(block); + std::size_t const block_extent = [&] { + if (auto it = bov.find(mix); it != bov.end()) return it->second; + return cm->regime().extent(mix); + }(); + SEQUANT_ASSERT(block_extent == block_hi - block_lo); + // Merge the block's range into the assembled coverage, requiring + // contiguity: a block that neither appends after nor prepends before the + // filled range would leave a gap or overlap another block (a + // double-count). This is what makes disjoint gap-free tiling the only + // accepted assembly. + if (auto it = cov.find(mode); it == cov.end()) { + cov.emplace(mode, + std::pair{block_lo, block_hi}); + } else { + auto& lohi = it->second; + bool const append = block_lo == lohi.second; + bool const prepend = block_hi == lohi.first; + SEQUANT_ASSERT(append || prepend); + if (append) + lohi.second = block_hi; + else + lohi.first = block_lo; + } + // Reflect the assembled element width (hi - lo, lobound preserved) as the + // realized extent of the batch mode so size_in_bytes() tracks the + // reconstructed footprint. + auto const& lohi = cov.at(mode); + ov[mix] = lohi.second - lohi.first; + } + + /// Build a zero-data destination shaped like \p idx but with mode \p mode's + /// index widened to \p axis_src's FULL (unsliced) extent at \p + /// axis_src_mode -- the dry-run analogue of \c ResultTensorTA:: and \c + /// ResultTensorOfTensorTA::pre_sized_zeros_over_mode's outer-\c + /// TiledRange1 swap. \p axis_src is the unsliced carrier leaf's token (its + /// OWN recorded override for the axis index, if any, else the CostModel + /// regime's natural extent, is the "full" extent -- mirroring how the TA + /// backend reads the swapped-in dimension straight off \p axis_src rather + /// than assuming a fixed default). The width recorded here is a structural + /// (shape) fact, queryable via \c size_in_bytes()/overrides() immediately + /// -- exactly as the TA test \c batched_scratch_tot_presize_scatter checks + /// \c trange().dim(mode) right after presizing, before any block is + /// written. Once the scatter loop's \c write_into_slice() calls begin, the + /// AssembledCoverage bookkeeping there takes over the mode's reported + /// extent (the REALIZED/assembled width so far); by the time every + /// disjoint block has been written the assembled width converges back to + /// this same full extent. + [[nodiscard]] static ResultPtr pre_sized_zeros_over_mode( + container::svector const& idx, ExtentOverrides const& ov, + std::shared_ptr const& cm, std::size_t mode, + Result const& axis_src, std::size_t axis_src_mode) { + SEQUANT_ASSERT(mode < idx.size()); + auto const src_idx = indices_of(axis_src); + SEQUANT_ASSERT(axis_src_mode < src_idx.size()); + Index const& axis_ix = src_idx[axis_src_mode]; + auto const src_ov = overrides_of(axis_src); + std::size_t const full_extent = [&] { + if (auto it = src_ov.find(axis_ix); it != src_ov.end()) return it->second; + return cm->regime().extent(axis_ix); + }(); + auto merged = ov; + merged[idx[mode]] = full_extent; + return make_dryrun_result(idx, cm, std::move(merged)); + } + + [[nodiscard]] static container::svector> + mode_batches(container::svector const& idx, ExtentOverrides const& ov, + std::shared_ptr const& cm, std::size_t mode, + std::size_t target_batch_size) { + SEQUANT_ASSERT(mode < idx.size()); + Index const& ix = idx[mode]; + std::size_t extent; + if (auto it = ov.find(ix); it != ov.end()) + extent = it->second; + else + extent = cm->regime().extent(ix); + + container::svector> out; + if (target_batch_size == 0 || extent == 0) { + out.push_back({0, extent}); + return out; + } + for (std::size_t lo = 0; lo < extent; lo += target_batch_size) + out.push_back({lo, std::min(extent, lo + target_batch_size)}); + return out; + } +}; + +} // namespace detail + +/// +/// \brief Flat (non-CSV) zero-data tensor token. +/// +/// Carries only its own literal outer index list (canon order -- the same +/// order \c EvalExpr::canon_indices()/annot() use, so \c slice_mode()/ +/// \c mode_batches()'s positional `mode` argument indexes it correctly), an +/// \c ExtentOverrides table recording any runtime \c slice_mode()/ +/// \c mode_batches() narrowing (keyed by Index so it survives reshaping +/// across prod/sum/permute), and a shared \c CostModel. No tensor data is +/// ever allocated or copied; every op is index-set bookkeeping plus a +/// CostModel query. Mirrors \c ResultTensorTAPP's structure +/// (backends/tapp/result.hpp) with every real-tensor line replaced by that +/// bookkeeping. +/// +class ResultDryRun final : public Result { + public: + using Result::id_t; + + ResultDryRun(container::svector idxset, + std::shared_ptr cm, + ExtentOverrides overrides = {}) + : Result{Payload{}}, + indices_{std::move(idxset)}, + cm_{std::move(cm)}, + overrides_{std::move(overrides)} {} + + [[nodiscard]] container::svector const& indices() const noexcept { + return indices_; + } + [[nodiscard]] ExtentOverrides const& overrides() const noexcept { + return overrides_; + } + + /// The contiguous `[lo, hi)` element range of outer mode \p mode assembled so + /// far by write_into_slice() (empty `{0, 0}` if nothing written). + [[nodiscard]] std::pair assembled_range( + std::size_t mode) const { + if (auto it = assembled_.find(mode); it != assembled_.end()) + return it->second; + return {0, 0}; + } + + private: + struct Payload {}; + + [[nodiscard]] id_t type_id() const noexcept override { + return id_for_type(); + } + + [[nodiscard]] ResultPtr sum( + Result const& other, + std::array const& annot) const override { + return detail::DryRunOps::sum(indices_, overrides_, cm_, other, annot); + } + + [[nodiscard]] ResultPtr prod(Result const& other, + std::array const& annot, + DeNest /*DeNestFlag*/) const override { + return detail::DryRunOps::prod(indices_, overrides_, cm_, other, annot); + } + + [[nodiscard]] ResultPtr permute( + std::array const& ann) const override { + return detail::DryRunOps::permute(indices_, overrides_, cm_, ann); + } + + [[nodiscard]] ResultPtr adjoint( + std::array const& ann) const override { + return detail::DryRunOps::permute(indices_, overrides_, cm_, ann); + } + + [[nodiscard]] ResultPtr slice_mode(std::size_t mode, std::size_t elem_lo, + std::size_t elem_hi) const override { + return detail::DryRunOps::slice_mode(indices_, overrides_, cm_, mode, + elem_lo, elem_hi); + } + + [[nodiscard]] container::svector> + mode_batches(std::size_t mode, std::size_t target_batch_size) const override { + return detail::DryRunOps::mode_batches(indices_, overrides_, cm_, mode, + target_batch_size); + } + + void write_into_slice(Result const& block, std::size_t mode, + std::size_t block_lo, std::size_t block_hi) override { + detail::DryRunOps::write_into_slice(indices_, overrides_, assembled_, cm_, + block, mode, block_lo, block_hi); + } + + [[nodiscard]] ResultPtr pre_sized_zeros_over_mode( + std::size_t mode, Result const& axis_src, + std::size_t axis_src_mode) const override { + return detail::DryRunOps::pre_sized_zeros_over_mode( + indices_, overrides_, cm_, mode, axis_src, axis_src_mode); + } + + void add_inplace(Result const& other) override { + SEQUANT_ASSERT(other.is() || other.is()); + overrides_ = + detail::merge_overrides(overrides_, detail::overrides_of(other)); + } + + [[nodiscard]] ResultPtr symmetrize() const override { + return eval_result(indices_, cm_, overrides_); + } + + [[nodiscard]] ResultPtr antisymmetrize(size_t /*bra_rank*/) const override { + return eval_result(indices_, cm_, overrides_); + } + + [[nodiscard]] ResultPtr mult_by_phase(std::int8_t /*factor*/) const override { + return eval_result(indices_, cm_, overrides_); + } + + [[nodiscard]] std::size_t size_in_bytes() const final { + return cm_->memsize(indices_, overrides_); + } + + container::svector indices_; + std::shared_ptr cm_; + ExtentOverrides overrides_; + AssembledCoverage assembled_; +}; + +/// +/// \brief CSV/PNO tensor-of-tensor zero-data token. +/// +/// Like \c ResultDryRun, but additionally exposes an outer()/inner() split of +/// its (canon-order) index list -- inner = the proto-indexed (composite) +/// legs, e.g. a CSV amplitude's PNO domain leg `a_1`; outer = every +/// other (plain) leg, e.g. the PAO index `mu~_1`. The split is purely an +/// observability/testing convenience: \c size_in_bytes()'s arithmetic is +/// IDENTICAL to \c ResultDryRun's (\c CostModel::memsize already routes any +/// index list containing a proto-indexed entry through the moment-aware +/// `inner_pow` path internally, via \c tot_indices/inner_aware_volume -- +/// content-driven, not type-driven), so tests that want to confirm "this used +/// the k-th moment, not extent^k" can inspect inner() directly. +/// +/// Position semantics for \c slice_mode()/\c mode_batches(): the `mode` +/// argument the runtime passes is always resolved against the FULL +/// canon-order list (an optional trailing constructor argument, defaulting to +/// `outer ++ inner` when the caller does not need position accuracy, e.g. a +/// hand-built test instance); the \c DryRunLeafEvaluator (eval_expr.hpp) +/// always supplies the leaf's true \c canon_indices() order there, since only +/// LEAF-constructed instances are ever sliced by the runtime (\c slice_mode() +/// is invoked only inside the batched evaluator's leaf-wrapping closure, never +/// on a prod()/sum()-produced intermediate). +/// +class ResultDryRunNested final : public Result { + public: + using Result::id_t; + + ResultDryRunNested(container::svector outer, + container::svector inner, + std::shared_ptr cm, + ExtentOverrides overrides = {}, + container::svector canon_order = {}) + : Result{Payload{}}, + outer_{std::move(outer)}, + inner_{std::move(inner)}, + indices_{canon_order.empty() + ? [this] { + container::svector c = outer_; + c.insert(c.end(), inner_.begin(), inner_.end()); + return c; + }() + : std::move(canon_order)}, + cm_{std::move(cm)}, + overrides_{std::move(overrides)} {} + + [[nodiscard]] container::svector const& outer() const noexcept { + return outer_; + } + [[nodiscard]] container::svector const& inner() const noexcept { + return inner_; + } + [[nodiscard]] container::svector const& indices() const noexcept { + return indices_; + } + [[nodiscard]] ExtentOverrides const& overrides() const noexcept { + return overrides_; + } + + /// The contiguous `[lo, hi)` element range of outer mode \p mode assembled so + /// far by write_into_slice() (empty `{0, 0}` if nothing written). + [[nodiscard]] std::pair assembled_range( + std::size_t mode) const { + if (auto it = assembled_.find(mode); it != assembled_.end()) + return it->second; + return {0, 0}; + } + + private: + struct Payload {}; + + [[nodiscard]] id_t type_id() const noexcept override { + return id_for_type(); + } + + [[nodiscard]] ResultPtr sum( + Result const& other, + std::array const& annot) const override { + return detail::DryRunOps::sum(indices_, overrides_, cm_, other, annot); + } + + [[nodiscard]] ResultPtr prod(Result const& other, + std::array const& annot, + DeNest /*DeNestFlag*/) const override { + return detail::DryRunOps::prod(indices_, overrides_, cm_, other, annot); + } + + [[nodiscard]] ResultPtr permute( + std::array const& ann) const override { + return detail::DryRunOps::permute(indices_, overrides_, cm_, ann); + } + + [[nodiscard]] ResultPtr adjoint( + std::array const& ann) const override { + return detail::DryRunOps::permute(indices_, overrides_, cm_, ann); + } + + [[nodiscard]] ResultPtr slice_mode(std::size_t mode, std::size_t elem_lo, + std::size_t elem_hi) const override { + return detail::DryRunOps::slice_mode(indices_, overrides_, cm_, mode, + elem_lo, elem_hi); + } + + [[nodiscard]] container::svector> + mode_batches(std::size_t mode, std::size_t target_batch_size) const override { + return detail::DryRunOps::mode_batches(indices_, overrides_, cm_, mode, + target_batch_size); + } + + void write_into_slice(Result const& block, std::size_t mode, + std::size_t block_lo, std::size_t block_hi) override { + detail::DryRunOps::write_into_slice(indices_, overrides_, assembled_, cm_, + block, mode, block_lo, block_hi); + } + + [[nodiscard]] ResultPtr pre_sized_zeros_over_mode( + std::size_t mode, Result const& axis_src, + std::size_t axis_src_mode) const override { + return detail::DryRunOps::pre_sized_zeros_over_mode( + indices_, overrides_, cm_, mode, axis_src, axis_src_mode); + } + + void add_inplace(Result const& other) override { + SEQUANT_ASSERT(other.is() || other.is()); + overrides_ = + detail::merge_overrides(overrides_, detail::overrides_of(other)); + } + + [[nodiscard]] ResultPtr symmetrize() const override { + return eval_result(outer_, inner_, cm_, overrides_, + indices_); + } + + [[nodiscard]] ResultPtr antisymmetrize(size_t /*bra_rank*/) const override { + return eval_result(outer_, inner_, cm_, overrides_, + indices_); + } + + [[nodiscard]] ResultPtr mult_by_phase(std::int8_t /*factor*/) const override { + return eval_result(outer_, inner_, cm_, overrides_, + indices_); + } + + [[nodiscard]] std::size_t size_in_bytes() const final { + return cm_->memsize(indices_, overrides_); + } + + container::svector outer_; + container::svector inner_; + container::svector indices_; // canon order; outer_++inner_ content + std::shared_ptr cm_; + ExtentOverrides overrides_; + AssembledCoverage assembled_; +}; + +[[nodiscard]] inline ResultPtr make_dryrun_result( + container::svector idx, std::shared_ptr cm, + ExtentOverrides overrides) { + if (!detail::has_proto(idx)) + return eval_result(std::move(idx), std::move(cm), + std::move(overrides)); + container::svector outer, inner; + for (auto const& ix : idx) + (ix.has_proto_indices() ? inner : outer).push_back(ix); + return eval_result(std::move(outer), std::move(inner), + std::move(cm), std::move(overrides), + std::move(idx)); +} + +namespace detail { + +[[nodiscard]] inline container::svector indices_of(Result const& r) { + if (r.is()) return r.as().indices(); + SEQUANT_ASSERT(r.is()); + return r.as().indices(); +} + +[[nodiscard]] inline ExtentOverrides overrides_of(Result const& r) { + if (r.is()) return r.as().overrides(); + SEQUANT_ASSERT(r.is()); + return r.as().overrides(); +} + +} // namespace detail + +} // namespace sequant::eval::dryrun + +#endif // SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_RESULT_HPP diff --git a/SeQuant/core/eval/backends/dryrun/size_regime.hpp b/SeQuant/core/eval/backends/dryrun/size_regime.hpp new file mode 100644 index 0000000000..9faf1aa834 --- /dev/null +++ b/SeQuant/core/eval/backends/dryrun/size_regime.hpp @@ -0,0 +1,79 @@ +#ifndef SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_SIZE_REGIME_HPP +#define SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_SIZE_REGIME_HPP + +#include + +#include +#include +#include +#include +#include +#include + +namespace sequant::eval::dryrun { + +/// Per-space extents and per-rank CSV moment tables that define one size +/// regime for a dry-run replay. Extents are element counts; CSV moments are +/// power means over occupied pairs (PNO) or singles (OSV). +struct SizeRegime { + std::map space_extent; + + // csv_pno_moment[k] / csv_osv_moment[k] hold the k-th POWER MEAN + // M_k = (mean_over_pairs d^k)^(1/k) of the per-pair PNO / per-orbital OSV + // domain size d, for k in [1,4] (index 0 is unused, set to 1). inner_pow() + // returns M_k so that inner_aware_volume's per-member product over a + // k-composite group is M_k^k = mean(d^k), and outer_nocc^N * M_k^k equals + // the true block-sparse volume Sum_pairs d^k. Do NOT store raw moments + // mean(d^k) here: that would over-count k-composite groups by a further + // power of k. For a constant domain d, M_k = d for all k. + std::array csv_pno_moment{1.0, 1.0, 1.0, 1.0, 1.0}; + std::array csv_osv_moment{1.0, 1.0, 1.0, 1.0, 1.0}; + + // Moment tables for CSV cluster ranks >= 3 (CSV-CCSDT triples and beyond), + // keyed by cluster rank (= number of proto indices). csv_moment_by_rank[r][k] + // is the k-th power mean of the rank-r cluster domain. A rank not present + // falls back to csv_pno_moment (the rank-2 table) in inner_pow(), preserving + // the pre-rank-general behavior where every proto-rank >= 2 used the PNO + // table. Ranks 1 and 2 are held by csv_osv_moment / csv_pno_moment above and + // are NOT expected here (an entry for 1 or 2 is ignored by inner_pow()). + std::map> csv_moment_by_rank; + + /// \return the flat extent of \p ix's space; throws \c std::out_of_range + /// if the space is not present in \c space_extent (fail loud rather + /// than silently defaulting to 1). + [[nodiscard]] std::size_t extent(Index const& ix) const { + return space_extent.at(std::wstring{ix.space().base_key()}); + } + + /// \return the k-th power-mean moment for a proto-indexed CSV/PNO composite + /// index (\p k clamped to 0..4), or \c pow(extent, k) for a plain + /// (non-composite) index. Rank is determined by the number of proto + /// indices: 1 => OSV (occupied single), 2 => PNO (occupied pair), + /// >= 3 => the rank-specific csv_moment_by_rank table if present, + /// else the PNO (rank-2) table. + [[nodiscard]] double inner_pow(Index const& composite, std::size_t k) const { + if (k > 4) k = 4; + auto const& protos = composite.proto_indices(); + if (protos.empty()) + return std::pow(static_cast(extent(composite)), + static_cast(k)); + auto const rank = protos.size(); + if (rank <= 1) return csv_osv_moment[k]; + if (rank == 2) return csv_pno_moment[k]; + auto const it = csv_moment_by_rank.find(rank); + return (it != csv_moment_by_rank.end()) ? it->second[k] : csv_pno_moment[k]; + } + + [[nodiscard]] std::function idx_to_extent() const { + return [this](Index const& ix) { return extent(ix); }; + } + + [[nodiscard]] std::function inner_pow_fn() + const { + return [this](Index const& ix, std::size_t k) { return inner_pow(ix, k); }; + } +}; + +} // namespace sequant::eval::dryrun + +#endif // SEQUANT_CORE_EVAL_BACKENDS_DRYRUN_SIZE_REGIME_HPP From 5bcc97f6dcdf04d731bb0b7b9559d6e133d97b99 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Wed, 29 Jul 2026 12:55:28 -0400 Subject: [PATCH 002/213] optimize: batched cost model for multimode batching Perf-first (DenseTimeSpace) objective with peak_threshold as a ceiling; role-split (contracted/external) batchability; order-aware placement over the combined nest; per-node batch annotations consumed by the evaluator. --- SeQuant/core/batch_policy.hpp | 69 +- SeQuant/core/optimize/cost_model.hpp | 1387 ++++++++++++++++-- SeQuant/core/optimize/optimize.cpp | 78 +- SeQuant/core/optimize/optimize.hpp | 14 + SeQuant/core/optimize/options.hpp | 172 ++- SeQuant/core/optimize/single_term.hpp | 121 +- SeQuant/core/optimize/single_term_detail.hpp | 129 +- 7 files changed, 1728 insertions(+), 242 deletions(-) diff --git a/SeQuant/core/batch_policy.hpp b/SeQuant/core/batch_policy.hpp index 48c3c5cd45..d42168e9c1 100644 --- a/SeQuant/core/batch_policy.hpp +++ b/SeQuant/core/batch_policy.hpp @@ -3,6 +3,7 @@ #include #include +#include namespace sequant { @@ -12,7 +13,39 @@ class Tensor; /// One batchability policy shared by the single-term optimizer and the runtime /// batched evaluator (make_evaluator, Task A3). All predicates default empty. struct BatchPolicy { - std::function is_batchable_index = {}; + /// Spaces batchable in the CONTRACTED role: a mode of such a space is + /// batchable where it is summed. Companion to \ref + /// is_batchable_external_index (the EXTERNAL role). Splitting batchability by + /// role lets a caller admit a space only where batching it is meaningful -- + /// e.g. a space batchable only as an external spectator contributes none of + /// its contracted occurrences to the optimizer's 2^m search. Building block; + /// the derived "batchable in any role" query is \ref is_batchable_index(). + /// Defaults to decline every index; a caller opts spaces in explicitly. + std::function is_batchable_contracted_index = + [](Index const&) { return false; }; + /// Spaces batchable in the EXTERNAL role: a mode of such a space is batchable + /// where it is open on the term root (a spectator carried to the result), not + /// where it is contracted. Building block; declared adjacent to its + /// contracted companion. Defaults to decline every index; a caller that wants + /// external batching sets this predicate explicitly (there is no fallback to + /// the contracted role). + std::function is_batchable_external_index = + [](Index const&) { return false; }; + + /// Derived "batchable in ANY role": the union of the two building-block + /// predicates. This is NEVER a settable field -- it is computed from + /// \ref is_batchable_contracted_index and \ref is_batchable_external_index. + /// The runtime batched evaluator's accept predicate is this union (a mode is + /// accepted at runtime if it is batchable in either role); the factorizer's + /// role filters instead consume the individual building blocks. The building + /// blocks default-decline, so both are always callable here. + std::function is_batchable_index() const { + auto contracted = is_batchable_contracted_index; + auto external = is_batchable_external_index; + return [contracted, external](Index const& ix) { + return contracted(ix) || external(ix); + }; + } /// Per-index per-batch slice size (in elements) for a batchable index -- an /// UPPER BOUND, not a goal. Both the single-term optimizer and the runtime /// batched evaluator treat it as a ceiling: the realized whole-tile batch is @@ -21,6 +54,22 @@ struct BatchPolicy { std::function batch_target_size = {}; std::function is_volatile_leaf = {}; + /// If true, an external/spectator index -- open on the whole network's result + /// yet contracted at no node -- is eligible for batching; its per-slice size + /// comes from \c batch_target_size(ix) like any batchable index. Default + /// false = no spectator batching (byte-identical to non-spectator behavior). + /// Necessary but not sufficient: spectator axes are emitted only under a + /// TIME-FIRST objective (DenseTimeSpaceBatched) and only when the selected + /// root's modeled peak exceeds \c peak_threshold. Spectator batching is + /// therefore currently unavailable under the space-first objectives. + bool batch_spectator_indices = false; + + /// Enable the order-aware multilevel recompute cost model (resident-scan peak + /// + ordered-key flops recompute). false (default) => byte-identical + /// set-keyed DP. Consulted only by the batched objectives (threaded via + /// CostParams). + bool order_aware_recompute = false; + /// If true, restrict batching to persistent (amplitude-independent) subtrees, /// declining to batch any subtree that contains a volatile leaf. If false /// (the default), batch ACROSS THE BOARD: slicing the batch axis shrinks any @@ -40,6 +89,24 @@ struct BatchPolicy { /// accumulator + contribution co-residency of a node that contracts a /// batchable index. double accumulation_factor = 0.0; + + /// Peak-memory budget in BYTES for the batched objectives. Its meaning + /// DIFFERS between them: + /// + /// - SPACE-FIRST (DenseSpaceTimeBatched): a hard feasibility gate. The + /// single-term optimizer minimizes flops among schedules whose modeled peak + /// is <= peak_threshold, falling back to min-peak (best effort) when none + /// fit. Default +infinity => every schedule feasible => min flops => no + /// batching, i.e. here a finite value is the *enable* trigger for batching. + /// + /// - TIME-FIRST (DenseTimeSpaceBatched): NOT a feasibility gate. Root + /// selection ignores it entirely (peak breaks exact flop ties only), so it + /// can neither constrain the schedule's peak nor enable CONTRACTED-axis + /// batching (which is emitted regardless). Its ONLY effect is to trigger + /// EXTERNAL (spectator) axis emission, together with + /// \c batch_spectator_indices: axes are emitted iff the selected root's + /// modeled peak exceeds this threshold. + double peak_threshold = std::numeric_limits::infinity(); }; } // namespace sequant diff --git a/SeQuant/core/optimize/cost_model.hpp b/SeQuant/core/optimize/cost_model.hpp index a1f86fada1..275836849e 100644 --- a/SeQuant/core/optimize/cost_model.hpp +++ b/SeQuant/core/optimize/cost_model.hpp @@ -1,6 +1,7 @@ #ifndef SEQUANT_CORE_OPTIMIZE_COST_MODEL_HPP #define SEQUANT_CORE_OPTIMIZE_COST_MODEL_HPP +#include #include // helpers + EvalSequence + OptRes #include @@ -10,9 +11,12 @@ #include #include #include +#include #include +#include #include #include +#include #include #include @@ -66,6 +70,34 @@ EvalSequence run_single_term_opt(Model const& m, TensorNetwork const& network, return m.reconstruct(ctx, st); } +/// \brief Companion to \ref run_single_term_opt that also reports, for each +/// contraction (\c -1) node of the returned \c EvalSequence in emission order, +/// the sliced-set of batchable \c Index values realized at that node. Requires +/// \p Model to additionally expose \c reconstruct_batched_modes (currently only +/// \ref PeakBatchedModel does); see its doc comment for the precise per-node +/// convention (RPN / post-order, left-first, matching the shared \c build +/// recursion used by \ref reconstruct). +/// +/// \return The optimal EvalSequence, paired with one \c container::svector +/// of sliced \c Index per \c -1 token of that sequence, in the same +/// left-first post-order the sequence itself was emitted in. For the +/// nt==1 shortcut (no contractions) the modes vector is empty; for the +/// nt==2 shortcut (single contraction, no DP context is built) the +/// modes vector holds one empty entry (no batching info available). +template +std::pair> +run_single_term_opt_axes(Model const& m, TensorNetwork const& network, + TIdxs const& tidxs) { + auto const nt = network.tensors().size(); + if (nt == 1) return {EvalSequence{0}, {}}; + if (nt == 2) + return {EvalSequence{0, 1, -1}, + container::vector{NodeBatchAnnotation{}}}; + typename Model::Context ctx = m.build_context(network, tidxs); + auto st = solve_single_term(m, network, tidxs, ctx); + return m.reconstruct_batched_modes(ctx, st, network, tidxs); +} + /// \brief Additive single-term cost model (FLOPs or operand storage size). /// /// Implements the additive single-term DP, factored into the CostModel hooks @@ -293,8 +325,11 @@ inline double roofline_op_cost(double flops, double traffic, template struct PeakModel { IdxToSz idxsz; - /// Optional k-aware inner (CSV/PNO composite) extent; see footprint_counter. - std::function inner_pow = {}; + /// k-aware inner (CSV/PNO composite) extent; see footprint_counter. REQUIRED + /// whenever the network has composite indices (empty => inner_aware_volume + /// throws); pass an explicit no-op only for composite-free networks. No + /// default: omitting it silently mis-sized composites (4-PAO-integral bug). + std::function inner_pow; /// Predicate marking a leaf tensor as volatile (amplitude-dependent). Used /// ONLY to weight the secondary flop tie-break: a volatile contraction is /// replayed every iteration, so its flops are scaled by \c volatile_weight. @@ -315,6 +350,10 @@ struct PeakModel { /// peak increase for a potentially large flop reduction (e.g. forming a /// persistent 4-PNO integral instead of recomputing a ladder). double peak_flops_tolerance = 0.0; + /// Perf-first / peak-second selection: when true, `reconstruct` selects the + /// root-frontier point by (flops, then peak) instead of (peak, then flops), + /// bypassing `peak_flops_tolerance`. Default false = peak-first (unchanged). + bool perf_first = false; /// Prune disconnected (outer-product) subsets from the DP (see /// OptimizeOptions::prune_outer_products). Default true. bool prune_outer_products = true; @@ -435,6 +474,34 @@ struct PeakModel { EvalSequence reconstruct(Context const& /*ctx*/, container::vector const& st) const { size_t const full = st.size() - 1; + if (perf_first) { + // Perf-first / peak-second (non-batched): min flops, ties by lower peak, + // bypassing the peak_flops_tolerance epsilon band (a peak-first knob). + auto const& rootf = st[full]; + int pbest = 0; + for (int i = 1; i < static_cast(rootf.size()); ++i) + if (rootf[i].flops < rootf[pbest].flops || + (rootf[i].flops == rootf[pbest].flops && + rootf[i].peak < rootf[pbest].peak)) + pbest = i; + // Reuse the existing back-pointer walk with the chosen root index. + std::function pbuild = + [&](size_t n, int idx) -> EvalSequence { + if (std::popcount(n) == 1) + return EvalSequence{static_cast(std::countr_zero(n))}; + FrontPoint const& fp = st[n][idx]; + size_t const fs = fp.lp_first ? fp.lp : fp.rp; + int const fi = fp.lp_first ? fp.lp_idx : fp.rp_idx; + size_t const ss = fp.lp_first ? fp.rp : fp.lp; + int const si = fp.lp_first ? fp.rp_idx : fp.lp_idx; + EvalSequence s = pbuild(fs, fi); + EvalSequence b = pbuild(ss, si); + s.insert(s.end(), b.begin(), b.end()); + s.push_back(-1); + return s; + }; + return pbuild(full, pbest); + } // ε-tolerant selection: among frontier points within // (1 + peak_flops_tolerance) of the minimum peak, take the fewest flops // (ties broken by lower peak). tolerance == 0 recovers strict peak-min. @@ -483,11 +550,13 @@ struct PeakModel { template struct PeakBatchedModel { IdxToSz idxsz; - std::function is_batchable; std::function batch; std::function is_volatile_leaf; - /// Optional k-aware inner (CSV/PNO composite) extent; see footprint_counter. - std::function inner_pow = {}; + /// k-aware inner (CSV/PNO composite) extent; see footprint_counter. REQUIRED + /// whenever the network has composite indices (empty => inner_aware_volume + /// throws); pass an explicit no-op only for composite-free networks. No + /// default: omitting it silently mis-sized composites (4-PAO-integral bug). + std::function inner_pow; /// Replay weight applied to volatile contractions in the flop tie-break. double volatile_weight = 1.0; /// Roofline parameters for the secondary (tie-break) cost; see @@ -503,17 +572,111 @@ struct PeakBatchedModel { /// volatile leaf). Default false = batch across the board. See /// BatchPolicy::persistent_only. bool batch_persistent_only = false; - /// Relative peak tolerance for the final (root) selection; see - /// PeakModel::peak_flops_tolerance. 0 (default) = strict peak-min. + /// Unused by \ref reconstruct (superseded by the threshold-gated selection + /// below, driven by \ref peak_threshold / \ref numeric_size); retained for + /// source compatibility. See PeakModel::peak_flops_tolerance, which is + /// still consulted by the (unbatched) DensePeakSize model. double peak_flops_tolerance = 0.0; /// In-flight batch-contribution footprint multiplier; see /// BatchPolicy::accumulation_factor. Charged only on nodes that contract a /// batchable index (Ap != 0), into the all-co-resident peak term, to price /// the accumulator + contribution co-residency of K += contribution. double accumulation_factor = 0.0; + /// Peak-memory budget (BYTES) for threshold-gated selection; see + /// BatchPolicy::peak_threshold. +infinity (default) => min-flops (no + /// batching). + double peak_threshold = std::numeric_limits::infinity(); + /// Bytes per stored element, to compare the model's element-count peak to + /// peak_threshold (bytes). Default 8 (double / TensorD). + double numeric_size = 8.0; + /// Perf-first / peak-second selection: when true, `select_root` selects the + /// root-frontier point by (flops, then peak) and does NOT consult + /// `peak_threshold` as a feasibility gate (it can no longer force a + /// FLOPS-catastrophic factorization for its sliceability). Default false = + /// peak-first threshold-gated selection (unchanged). + bool perf_first = false; + + /// If true (the default), charge the batch RECOMPUTATION cost on the flops/ + /// exec mode. The batched evaluator re-executes each contraction per tile of + /// the ancestor batch modes its result does NOT carry (across-batch work is + /// recomputed; within-batch sharing is cached -- see eval.hpp "replays the + /// build of every compatible persistent final"). A node at ancestor-sliced- + /// set B is charged nbatches(b) for each b in B not open in the node, so a + /// schedule that slices many modes it must recompute across pays for it. The + /// alternative (false) assumes WORK PARITY (batching is free on flops), which + /// under-costs heavily-sliced families and does not reflect the true cost of + /// batching; kept only as an escape hatch for comparison. + bool charge_batch_recompute = true; + /// Opt-in refinement of \ref charge_batch_recompute: do not bill a node for + /// enclosing batch loops it can hoist above **for free**. + /// + /// The flat charge above is order-blind -- it bills nbatches(b) for every b + /// in B the node does not carry, with no notion of where the node sits + /// relative to those loops. That systematically OVER-charges the hoistable + /// case. This flag fixes the ORDER-INDEPENDENT half of that defect: when a + /// node carries NONE of the enclosing batched modes + /// (`B & open_modes[n] == 0`), slicing any of them cannot shrink it, so it + /// can be built once above all of them at unchanged footprint -- a pure flops + /// saving needing no peak representation. Its correct factor is rf = 1. + /// + /// The remaining half -- an escaped loop INNER to a node's placement when the + /// node does carry some enclosing mode -- is order-DEPENDENT and cannot be + /// expressed while the cell is keyed by a set: both nesting orders reach the + /// same `(n, B)` cell, and their correct costs differ. That needs an ordered + /// key and is deliberately NOT attempted here. + /// + /// Default false, so \ref charge_batch_recompute alone reproduces the + /// historical cost exactly. **\ref seeded_forest_peak must keep this false + /// on its probe**: its work-neutrality guard (seeded_flops != unseeded_flops + /// => decline) depends on the FLAT charge firing for a seeded external mode, + /// which by construction is carried by no node that does not carry it, i.e. + /// exactly the `B & open_modes[n] == 0` case this flag exempts. Enabling it + /// there would make the guard vacuously true and admit non-work-neutral + /// seeds. + bool order_aware_recompute = false; + /// Spaces batchable in the CONTRACTED role -- i.e. a mode of such a space is + /// batchable where it is summed at some node. Building block; companion to + /// \ref is_batchable_external_index. Declared here (outside the + /// positionally-initialized prefix) and ADJACENT to its external companion so + /// the two roles read together; set it by member assignment. Defaults to + /// decline every index => no mode is batchable in the contracted role. + std::function is_batchable_contracted_index = + [](Index const&) { return false; }; + /// Spaces batchable in the EXTERNAL role -- i.e. a mode of such a space is + /// batchable when it is open on the term root (a spectator carried to the + /// result), NOT when it is contracted. Companion to \ref + /// is_batchable_contracted_index, which admits spaces batchable in the + /// CONTRACTED role. Keeping the two roles as separate caller-supplied space + /// sets is what lets this layer stay domain-generic: the caller decides which + /// spaces are batchable in which role, and \ref build_context drops every + /// mode whose role's predicate rejects it (so a space batchable only as + /// external never bloats the 2^m search with its contracted occurrences). + /// Defaults to decline every index; a caller that wants external batching + /// sets it explicitly (there is no fallback to the contracted-role + /// predicate). Declared here (outside the positionally-initialized prefix) so + /// existing aggregate initializations are unaffected; set it by member + /// assignment. + std::function is_batchable_external_index = + [](Index const&) { return false; }; + + /// Derived "batchable in ANY role": true iff the index is batchable in the + /// contracted OR the external role. This is NOT a settable field; the DP's + /// role filters consume the individual building blocks, never this union. + /// The building blocks default-decline, so both are always callable here. + bool is_batchable(Index const& ix) const { + return is_batchable_contracted_index(ix) || is_batchable_external_index(ix); + } /// Prune disconnected (outer-product) subsets from the DP (see /// OptimizeOptions::prune_outer_products). Default true. bool prune_outer_products = true; + /// Term-level gate for \ref reconstruct_batched_modes emitting \c + /// BatchModeType::External entries (genuine external modes; see \ref + /// is_external_mode), threaded from \ref CostParams::batch_spectator_indices + /// / BatchPolicy::batch_spectator_indices. Default false so every OTHER + /// PeakBatchedModel construction (peak_cost_batched, compute_external_batch_ + /// axis's own model, existing tests) is unaffected and emits no External + /// entries -- byte-identical to before this member existed. + bool batch_spectator_indices = false; /// One non-dominated (peak, flops) trade-off for a (subset, sliced-set \c B) /// cell. \c aprime is the sliced-set chosen at this node; the children are @@ -538,9 +701,10 @@ struct PeakBatchedModel { /// Precomputed tables and per-(subset, sliced-set) lookup parameters built /// once by build_context. struct Context { - /// Ordered, deduplicated batchable indices (bit \c k maps to \c aux[k]). - container::vector aux; - /// Number of batchable indices (= aux.size()). + /// Ordered, deduplicated batchable indices (bit \c k maps to \c + /// batchable_modes[k]). + container::vector batchable_modes; + /// Number of batchable indices (= batchable_modes.size()). std::size_t m = 0; /// Number of sliced-sets (= 2^m). std::size_t nB = 1; @@ -548,8 +712,8 @@ struct PeakBatchedModel { std::size_t nt = 0; /// tables[B][n] = footprint of subset n under sliced-set B. container::vector> tables; - /// open_aux[n] = bitmask of batchable indices open in subset n. - container::vector open_aux; + /// open_modes[n] = bitmask of batchable indices open in subset n. + container::vector open_modes; /// Bitmask of volatile leaf tensors. std::size_t volatile_mask = 0; /// idx[n] = subset n's open (result) indices, for the flop tie-break. @@ -559,6 +723,12 @@ struct PeakBatchedModel { /// hot loop uses fast_flops (see below). std::function flops_of; + /// nbatches[k] = number of batch tiles of batchable_modes[k] = ceil(extent + /// / target), clamped to >= 1. Used to charge batch recomputation (see + /// charge_batch_recompute): a node inside an ancestor batch loop over + /// batchable_modes[k] that does not carry batchable_modes[k] is re-executed + /// nbatches[k] times. + container::vector nbatches; // --- fast per-subset flop precompute (relax tie-break hot path) --- // Per subset, sorted (FullLabelCompare-ordered) atom IDs: outer atoms are @@ -637,18 +807,154 @@ struct PeakBatchedModel { return mem == 1.0 ? 0.0 : mem; } - /// Context-restricted size of subset s under sliced-set ctx (the table is - /// indexed by the part of ctx actually open in s; mirrors the oracle). - double sz(std::size_t s, std::size_t ctx) const { - return tables[ctx & open_aux[s]][s]; + /// Footprint of subset s under an explicit sliced-set UNION mask (the + /// order-independent bitmask, as returned by \ref cell_union). Slicing is a + /// pure footprint change, so the size depends only on which modes are in + /// the union, never on the cell's nesting order -- this is the primitive + /// that + /// \ref sz and the external-placement re-price (\ref subtree_peak) share, + /// so an EXTERNAL mode can be injected into the union directly without a + /// precomputed ordered cell (external bits are excluded from build_cells). + double sz_u(std::size_t U, std::size_t s) const { + return tables[U & open_modes[s]][s]; } - /// Per-context leaf-sum of subset s (sum of singleton sizes under ctx). - double Lof(std::size_t s, std::size_t ctx) const { + /// Per-context leaf-sum of subset s under an explicit union mask. + double Lof_u(std::size_t U, std::size_t s) const { double r = 0.0; for (std::size_t b = 0; b < nt; ++b) - if (s & (std::size_t{1} << b)) r += sz(std::size_t{1} << b, ctx); + if (s & (std::size_t{1} << b)) r += sz_u(U, std::size_t{1} << b); return r; } + + /// Context-restricted size of subset s under sliced-set ctx (the table is + /// indexed by the part of ctx actually open in s; mirrors the oracle). + double sz(std::size_t s, std::size_t id) const { + return sz_u(cell_union(id), s); + } + /// Per-context leaf-sum of subset s (sum of singleton sizes under ctx). + double Lof(std::size_t s, std::size_t id) const { + return Lof_u(cell_union(id), s); + } + + // --- ordered-cell layer (order_aware_recompute) ----------------------- + // A DP cell is identified by `id`. When `ordered` is false, `id` IS the + // sliced-set bitmask B (identity; nCells == nB), so every helper here + // reduces to the historical bitmask ops and the DP is byte-identical. When + // true, `id` indexes an ordered SEQUENCE of batched modes (outer to inner); + // cell_union recovers the bitmask for the (order-independent) footprint + // tables, descend appends a contracted-here set as inner modes, and + // escaped_outer charges only enclosing modes OUTER to a node's innermost- + // carried placement. + bool ordered = false; + std::size_t cap = 3; // max sequence length when ordered + std::size_t nCells = 1; // number of DP cells (== nB when !ordered) + // Bitmask of EXTERNAL batchable modes (bit k set iff is_external_mode(k)): + // open on the root, contracted at no node. Only set when ordered (built + // from open_modes, which build_context now assigns before build_cells()). + // build_cells skips these bits when enumerating sequences, so an ordered + // cell's union/sequence never contains an external mode -- nesting order + // is meaningless for a mode that is never a contracted-here set at any + // node, and admitting it would blow up the enumeration for nothing. + std::size_t external_mask = 0; + container::vector cell_union_; // id -> union bitmask + container::vector> + cell_seq_; // id -> ordered mode-bit indices + container::vector + cell_descend_; // id*m + k -> child id / SIZE_MAX + + std::size_t cell_union(std::size_t id) const { + return ordered ? cell_union_[id] : id; + } + // Append the modes of `Ap` (ascending bit order = canonical co-contracted + // order) as inner positions. SIZE_MAX if that exceeds cap or repeats a + // mode. + std::size_t descend(std::size_t id, std::size_t Ap) const { + if (!ordered) return id | Ap; + for (std::size_t k = 0; k < m; ++k) + if (Ap & (std::size_t{1} << k)) { + id = cell_descend_[id * m + k]; + if (id == std::numeric_limits::max()) return id; + } + return id; + } + // Enclosing modes charged as recompute for a node carrying `carried`: + // those OUTER to the node's innermost-carried placement it does not carry. + // Carr == 0 (no carried enclosing mode) hoists above every loop => none. + std::size_t escaped_outer(std::size_t id, std::size_t carried) const { + if (!ordered) return id & ~carried; + auto const& seq = cell_seq_[id]; + int placement = -1; + for (int p = static_cast(seq.size()) - 1; p >= 0; --p) + if ((carried >> seq[p]) & 1u) { + placement = p; + break; + } + if (placement < 0) return 0; + std::size_t esc = 0; + for (int p = 0; p < placement; ++p) + if (!((carried >> seq[p]) & 1u)) esc |= (std::size_t{1} << seq[p]); + return esc; + } + + // Fill the cell tables. When !ordered, nCells == nB and the tables stay + // empty (helpers use the bitmask directly). When ordered, enumerate every + // ordered sequence of batched modes up to `cap` length (id 0 = the empty + // sequence = the term root); cell_union_/cell_seq_/cell_descend_ index + // them. + void build_cells() { + if (!ordered) { + nCells = nB; + return; + } + // Enumeration blowup guard: estimate Sum_{k<=cap} P(m,k) and fall back to + // set-keyed if it would be huge (a bounded-nest schedule is still + // correct; only optimality is lost). With cap==3 this is ~m^3/6, so it + // only trips for pathologically many batchable indices. + { + std::size_t const limc = std::min(m, cap); + std::size_t est = 1, term = 1; + for (std::size_t k = 1; k <= limc; ++k) { + term *= (m - k + 1); + est += term; + } + if (est > 100000) { + ordered = false; + nCells = nB; + return; + } + } + cell_seq_.assign(1, container::svector{}); + cell_union_.assign(1, std::size_t{0}); + std::map, std::size_t> seq_id; + seq_id.emplace(cell_seq_[0], std::size_t{0}); + std::size_t const lim = std::min(m, cap); + for (std::size_t id = 0; id < cell_seq_.size(); ++id) { + if (cell_seq_[id].size() >= lim) continue; + for (std::uint8_t k = 0; k < m; ++k) { + if ((external_mask >> k) & 1u) continue; // external: not nestable + if (cell_union_[id] & (std::size_t{1} << k)) continue; + auto ns = cell_seq_[id]; + ns.push_back(k); + if (seq_id.find(ns) == seq_id.end()) { + seq_id.emplace(ns, cell_seq_.size()); + cell_union_.push_back(cell_union_[id] | (std::size_t{1} << k)); + cell_seq_.push_back(std::move(ns)); + } + } + } + nCells = cell_seq_.size(); + cell_descend_.assign(nCells * m, std::numeric_limits::max()); + for (std::size_t id = 0; id < nCells; ++id) { + if (cell_seq_[id].size() >= lim) continue; + for (std::uint8_t k = 0; k < m; ++k) { + if ((external_mask >> k) & 1u) continue; // external: not nestable + if (cell_union_[id] & (std::size_t{1} << k)) continue; + auto ns = cell_seq_[id]; + ns.push_back(k); + cell_descend_[id * m + k] = seq_id.at(ns); + } + } + } }; template @@ -657,29 +963,91 @@ struct PeakBatchedModel { // CSE is not supported for DensePeakSizeBatched. Context ctx; ctx.nt = network.tensors().size(); - ctx.aux = batchable_index_list(network, is_batchable); - ctx.m = ctx.aux.size(); - // The accumulation_factor charge is per accumulation node (charged on each - // node that contracts a batchable index). Its semantics are only - // well-defined for a single batch axis; with multiple batchable indices the - // per-node, once-per-node charge would conflate independent accumulations. - SEQUANT_ASSERT( - (accumulation_factor == 0.0 || ctx.m <= 1) && - "DensePeakSizeBatched: accumulation_factor != 0 requires at most one " - "batchable index"); + // Candidates from BOTH batchability roles (contracted / external); the role + // filter below keeps each mode only if its actual role admits it. + ctx.batchable_modes = batchable_mode_list( + network, is_batchable_contracted_index, is_batchable_external_index); + // open_modes must be assigned BEFORE build_cells() so the ordered-cell + // enumeration can identify and exclude EXTERNAL modes (see + // Context::external_mask below). NOT pruned: is_external_mode scans it + // over the FULL subset lattice (including disconnected subsets) to verify + // a mode is never contracted, so every subset's open-mode bitmask must be + // real. Computed here first because the role filter needs the root open + // set, then recomputed if the filter shrinks the mode list. + ctx.open_modes = subset_open_aux(network, tidxs, ctx.batchable_modes); + { + // Role filter: a mode open on the root occurs in the EXTERNAL role and is + // batchable only if is_batchable_external_index admits it; otherwise it + // occurs CONTRACTED and is batchable only if + // is_batchable_contracted_index admits it. Dropping the rest keeps 2^m + // free of modes that can never be batched in the role they actually occur + // in (e.g. the contracted members of a space that is only + // external-batchable). Each role consults ONLY its own building block -- + // there is no fallback from the external role to the contracted one, so a + // space batchable only in the contracted role never admits its external + // occurrences. Both building blocks default-decline, hence are always + // callable here. + std::size_t const root = (std::size_t{1} << ctx.nt) - 1; + container::vector kept; + kept.reserve(ctx.batchable_modes.size()); + for (std::size_t k = 0; k < ctx.batchable_modes.size(); ++k) { + Index const& ix = ctx.batchable_modes[k]; + bool const ext = (ctx.open_modes[root] >> k) & 1u; + bool const keep = ext ? is_batchable_external_index(ix) + : is_batchable_contracted_index(ix); + if (keep) kept.push_back(ix); + } + if (kept.size() != ctx.batchable_modes.size()) { + ctx.batchable_modes = std::move(kept); + ctx.open_modes = subset_open_aux(network, tidxs, ctx.batchable_modes); + } + } + ctx.m = ctx.batchable_modes.size(); + // Per-mode batch-tile count for the recompute charge: ceil(extent/target), + // >= 1. batch() returns the target tile size; a 0/absent target => 1 tile + // (no batching of that mode => no recompute). + ctx.nbatches.assign(ctx.m, 1.0); + for (std::size_t k = 0; k < ctx.m; ++k) { + double const ext = static_cast(idxsz(ctx.batchable_modes[k])); + double const tgt = static_cast(batch(ctx.batchable_modes[k])); + ctx.nbatches[k] = (tgt > 0.0) ? std::max(1.0, std::ceil(ext / tgt)) : 1.0; + } + // accumulation_factor is charged per accumulation node (Ap != 0) and is + // valid for any number of batchable indices: with nested accumulation the + // per-node charges co-exist at the peak. Validated by the identity + // peak_cost_batched == reconstructed_batched_peak (test [batched-accum]). ctx.nB = std::size_t{1} << ctx.m; + // A mode open on the root and contracted at no node (is_external_mode) is + // carried unchanged from leaves to root: nesting order is meaningless for + // it (it never appears in any node's contracted-here set), so the ordered + // enumeration excludes it -- ordered cells contain contracted modes only. + for (std::size_t k = 0; k < ctx.m; ++k) + if (is_external_mode(ctx, k)) ctx.external_mask |= (std::size_t{1} << k); + // Order-aware cell layer: when order_aware_recompute is set, DP cells index + // ordered sequences of batched modes; otherwise cells are the bitmask B and + // this is a no-op (nCells == nB). cap == m => full enumeration (never worse + // than the set-keyed DP); build_cells falls back to set-keyed if m is + // large. + ctx.ordered = order_aware_recompute; + // Cap the ordered nest DEPTH (sequence length), NOT m: the number of modes + // that co-nest in a single term (m_B) is small (<=3 for C60 contracted), + // while m (all batchable indices) can be large -- cap==m gives Sum P(m,k) + // cells (13700 at m=7). Depth 3 => Sum_{k<=3} P(m,k) (260 at m=7), + // polynomial in m, so it engages on large-m terms too. A term needing a + // deeper nest loses only optimality (a bounded-nest schedule is still + // correct), never correctness. build_cells still guards a hard cell-count + // blowup. + ctx.cap = std::min(ctx.m, 3); + ctx.build_cells(); // Outer-product pruning: skip building tables for disconnected subsets the // DP will never form (solve_single_term also skips them). connected[n]==1 // for singletons/empty and for connected subsets; the (~2x connected) // needed-mask is derived internally where a complement lookup requires it. auto const connected = outer_product_connectivity(network, tidxs, prune_outer_products); - ctx.tables = sliced_footprints(network, tidxs, idxsz, is_batchable, batch, - ctx.aux, inner_pow, &connected); - // open_aux is NOT pruned: is_spectator_axis scans open_aux over the FULL - // subset lattice (including disconnected subsets) to verify an axis is - // never contracted, so every subset's open-axis bitmask must be real. - ctx.open_aux = subset_open_aux(network, tidxs, ctx.aux); + ctx.tables = + sliced_footprints(network, tidxs, idxsz, is_batchable_contracted_index, + batch, ctx.batchable_modes, inner_pow, &connected); ctx.volatile_mask = leaf_volatile_mask(network, is_volatile_leaf); // Per-subset open indices + a flop counter, for the lexicographic // (peak, then flops) tie-break (mirrors PeakModel). The flop tie-break uses @@ -761,14 +1129,14 @@ struct PeakBatchedModel { } State leaf(Context const& ctx, size_t n) const { - State s(ctx.nB); - for (std::size_t B = 0; B < ctx.nB; ++B) + State s(ctx.nCells); + for (std::size_t B = 0; B < ctx.nCells; ++B) s[B].push_back(BFrontPoint{ctx.sz(n, B), 0.0, 0, 0, true, 0, -1, -1}); return s; } State init(Context const& ctx, size_t /*n*/) const { - return State(ctx.nB); // nB empty frontiers; relax fills them + return State(ctx.nCells); // empty frontiers; relax fills them } void relax(Context& ctx, size_t n, size_t lp, size_t rp, State const& lp_st, @@ -776,7 +1144,7 @@ struct PeakBatchedModel { // Secondary (tie-break) cost: roofline wall-time proxy per replay, charged // volatile_weight times for volatile (replayed) contractions. Uses the full // (unsliced) operand+result footprint as the per-replay traffic; slicing - // reduces peak (primary axis), not total work. machine_balance==0 => flops. + // reduces peak (primary mode), not total work. machine_balance==0 => flops. double const w = (ctx.volatile_mask & n) ? volatile_weight : 1.0; double const cflops = w * roofline_op_cost( @@ -785,46 +1153,85 @@ struct PeakBatchedModel { : ctx.flops_of(ctx.idx[lp], ctx.idx[rp], ctx.idx[n]), ctx.sz(lp, 0) + ctx.sz(rp, 0) + ctx.sz(n, 0), machine_balance, fast_mem_elems, block_tiles, block_prefactor); - for (std::size_t B = 0; B < ctx.nB; ++B) { + for (std::size_t B = 0; B < ctx.nCells; ++B) { + // Batch recomputation charge: this node sits inside the ancestor batch + // loops over the modes in B. For each b in B whose mode this node's + // result does NOT carry (b not open in n), the node is re-executed + // nbatches(b) times (across-batch recompute); modes it carries are + // partitioned (x1). Default off => rf==1 => historical work-parity cost. + double rf = 1.0; + if (charge_batch_recompute) { + // Recompute charge over the enclosing modes this node is re-executed + // across. !ordered: B & ~carried (every escaped mode -- the historical + // set charge). ordered: only escaped modes OUTER to the node's + // innermost- carried placement (escaped modes inner to it hoist above + // for free; Carr == 0 hoists above the whole nest => none, subsuming + // A3a). escaped_outer collapses to the set charge when !ordered, + // byte-identical. + std::size_t const esc = ctx.escaped_outer(B, ctx.open_modes[n]); + for (std::size_t k = 0; k < ctx.m; ++k) + if (esc & (std::size_t{1} << k)) rf *= ctx.nbatches[k]; + } + double const cflops_B = cflops * rf; // Batchable indices contracted at THIS node: open at children but not at // the parent. By default batching is applied ACROSS THE BOARD: slicing - // the batch axis shrinks any intermediate carrying it regardless of + // the batch mode shrinks any intermediate carrying it regardless of // volatility (footprint objective) while leaving flops unchanged, so the // persistence gate would only ever raise the modelled peak. Set // batch_persistent_only to restore the persistent-only gate (decline to // slice subsets that contain a volatile leaf). - std::size_t const Acand = + std::size_t const contracted_here = (batch_persistent_only && (ctx.volatile_mask & n)) ? std::size_t{0} - : ((ctx.open_aux[lp] | ctx.open_aux[rp]) & ~ctx.open_aux[n]); - // Enumerate every subset A' of Acand (including the empty set). - std::size_t Ap = Acand; + : ((ctx.open_modes[lp] | ctx.open_modes[rp]) & + ~ctx.open_modes[n]); + // Enumerate every subset A' of contracted_here (including the empty set). + std::size_t Ap = contracted_here; while (true) { - std::size_t const C = B | Ap; - double const szlp = ctx.sz(lp, C), szrp = ctx.sz(rp, C), - szn = ctx.sz(n, B); - // A node that contracts a batchable index (Ap != 0) is accumulated over - // the aux batches (K += contribution); the in-flight contribution (same - // index set as the result, size szn) co-resides with the accumulator. - // Charge it once, on the all-co-resident moment only -- the pre-result - // staged terms (Lrp+pl, szlp+prr) exclude it since szn is not yet - // built. - double const contrib = (Ap != 0) ? accumulation_factor * szn : 0.0; - double const both = szlp + szrp + szn + contrib; - double const Lrp = ctx.Lof(rp, C), Llp = ctx.Lof(lp, C); - // Cross every (peak,flops) trade-off of the two children at context C. - for (int li = 0; li < static_cast(lp_st[C].size()); ++li) - for (int ri = 0; ri < static_cast(rp_st[C].size()); ++ri) { - double const pl = lp_st[C][li].peak, prr = rp_st[C][ri].peak; - double const lpf = std::max({Lrp + pl, szlp + prr, both}); - double const rpf = std::max({Llp + prr, szrp + pl, both}); - pareto_insert(acc[B], BFrontPoint{std::min(lpf, rpf), - lp_st[C][li].flops + - rp_st[C][ri].flops + cflops, - lp, rp, lpf <= rpf, Ap, li, ri}); - } + std::size_t const C = ctx.descend(B, Ap); + // ordered: skip an over-cap / repeat descent (SIZE_MAX). !ordered: + // descend == B|Ap and never returns SIZE_MAX, so this is + // byte-identical. + if (C != std::numeric_limits::max()) { + double const szlp = ctx.sz(lp, C), szrp = ctx.sz(rp, C), + szn = ctx.sz(n, B); + // A node that contracts a batchable index (Ap != 0) is accumulated + // over the batches of that mode (K += contribution); the in-flight + // contribution (same index set as the result, size szn) co-resides + // with the accumulator. Charge it once, on the all-co-resident moment + // only + // -- the pre-result staged terms (Lrp+pl, szlp+prr) exclude it since + // szn is not yet built. + double const contrib = (Ap != 0) ? accumulation_factor * szn : 0.0; + double const both = szlp + szrp + szn + contrib; + double const Lrp = ctx.Lof(rp, C), Llp = ctx.Lof(lp, C); + // Resident-scan (order_aware_recompute): a node that batches (Ap != + // 0) allocates its accumulator (szn) up front and holds it across the + // batch loop, so it co-resides with the children as they evaluate per + // batch -- the pre-result staged terms (which exclude szn as "not yet + // built") must include it. Composed recursively, each ancestor + // batching node's szn stacks onto every descendant's peak: the + // Sum-of-enclosing-residents scan. Ap == 0 (unbatched, built once) + // keeps szn out, byte-identical. + double const res = (order_aware_recompute && Ap != 0) ? szn : 0.0; + // Cross every (peak,flops) trade-off of the two children at context + // C. + for (int li = 0; li < static_cast(lp_st[C].size()); ++li) + for (int ri = 0; ri < static_cast(rp_st[C].size()); ++ri) { + double const pl = lp_st[C][li].peak, prr = rp_st[C][ri].peak; + double const lpf = + std::max({Lrp + pl + res, szlp + prr + res, both}); + double const rpf = + std::max({Llp + prr + res, szrp + pl + res, both}); + pareto_insert( + acc[B], BFrontPoint{std::min(lpf, rpf), + lp_st[C][li].flops + rp_st[C][ri].flops + + cflops_B, + lp, rp, lpf <= rpf, Ap, li, ri}); + } + } // C != SIZE_MAX if (Ap == 0) break; - Ap = (Ap - 1) & Acand; + Ap = (Ap - 1) & contracted_here; } } } @@ -832,23 +1239,142 @@ struct PeakBatchedModel { void finalize(Context& /*ctx*/, size_t /*n*/, container::vector& /*st*/) const {} - EvalSequence reconstruct(Context const& ctx, - container::vector const& st) const { + /// \brief True iff batchable mode bit \p k is a genuine external mode of + /// this network: it is OPEN on the root result AND is contracted at NO node. + /// + /// An external mode is carried unchanged from the leaves up to the root: + /// whenever any tensor in a subset carries the mode, the mode stays OPEN in + /// that subset, so it never appears in any node's contracted-at-node set + /// (\c contracted_here = (open_modes[lp]|open_modes[rp]) & ~open_modes[n]) + /// and slicing it is purely a footprint change with identical work. The + /// forest-batching seed path asserts this before honoring a seed mode: + /// seeding a mode that is actually contracted somewhere would mis-size the + /// nodes that contract it. + bool is_external_mode(Context const& ctx, std::size_t k) const { std::size_t const root = (std::size_t{1} << ctx.nt) - 1; - // ε-tolerant selection on the root's B=0 frontier: among points within - // (1 + peak_flops_tolerance) of the minimum peak, fewest flops (ties broken - // by lower peak). tolerance == 0 recovers strict peak-min. - auto const& rootf = st[root][0]; - double minpeak = std::numeric_limits::max(); - for (auto const& fp : rootf) minpeak = std::min(minpeak, fp.peak); - double const thresh = minpeak * (1.0 + peak_flops_tolerance); + // Must be carried on the root result (a genuine external index). + if (!((ctx.open_modes[root] >> k) & 1u)) return false; + // Leaves (single-tensor subsets) that carry mode k. + std::size_t leafmask = 0; + for (std::size_t b = 0; b < ctx.nt; ++b) + if ((ctx.open_modes[std::size_t{1} << b] >> k) & 1u) + leafmask |= (std::size_t{1} << b); + // Whenever the mode is AVAILABLE in a subset (some carrying leaf is in it) + // it must be OPEN in that subset -- else it is contracted at the node that + // forms that subset, i.e. not a pure external mode. + for (std::size_t n = 1; n <= root; ++n) + if ((leafmask & n) && !((ctx.open_modes[n] >> k) & 1u)) return false; + return true; + } + + /// Threshold-gated root-frontier selection shared by \ref reconstruct and + /// \ref reconstruct_batched_modes: among points whose peak (bytes) fits + /// peak_threshold, pick fewest flops (ties by lower peak). If none fit, pick + /// min peak (best effort). peak_threshold == +inf => all feasible => min + /// flops => the non-batched schedule. Returns the chosen index into + /// \c st[root][root_B]. + /// + /// \param root_B The batch context read at the ROOT. The default 0 (the + /// empty sliced-set) is the historical behavior, kept byte-identical. + /// The forest-batching seed path passes a non-zero mask to SEED an + /// external mode into the root frontier, so the whole tree + /// is sized (and its peak reported) with that mode sliced -- + /// \c st[root][root_B] is already computed by the DP's B-loop; only + /// which B the root selection reads changes. + int select_root(Context const& ctx, container::vector const& st, + std::size_t root_B = 0) const { + std::size_t const root = (std::size_t{1} << ctx.nt) - 1; + auto const& rootf = st[root][root_B]; + auto peak_bytes = [this](double peak_elems) { + return peak_elems * numeric_size; + }; + if (perf_first) { + // Perf-first / peak-second: min flops, ties by lower peak. The frontier + // keeps one min-peak point per distinct flops value (pareto_insert prunes + // equal-flops higher-peak points), so this both picks the cheapest + // factorization and takes its fully-sliced (min-peak) realization. + // + // peak_threshold acts as a CEILING, not a min-peak objective: among the + // points whose byte peak fits the budget, take the fewest flops (ties by + // lower peak). A term is therefore batched only when its cheapest + // (unbatched) schedule would exceed the budget -- flops are never traded + // for peak BELOW the ceiling, so this cannot force the flops-catastrophic + // factorization that peak-first (the !perf_first branch) can. If no point + // fits (the budget is below even the min-peak schedule) keep the + // perf-first character and fall back to GLOBAL min flops (best effort, + // accepting the overage) rather than peak-first's min-peak fallback. + // peak_threshold == +inf (the default when no budget is set) makes every + // point feasible, reducing this to the pure min-flops selection -- + // byte-identical to before this ceiling existed. + auto better = [&](int i, int j) { + return rootf[i].flops < rootf[j].flops || + (rootf[i].flops == rootf[j].flops && + rootf[i].peak < rootf[j].peak); + }; + int pbest = -1; + for (int i = 0; i < static_cast(rootf.size()); ++i) + if (peak_bytes(rootf[i].peak) <= peak_threshold && + (pbest < 0 || better(i, pbest))) + pbest = i; + bool const fit = pbest >= 0; // a schedule met the ceiling + if (!fit) // nothing fits the budget: perf-first best effort + for (int i = 0; i < static_cast(rootf.size()); ++i) + if (pbest < 0 || better(i, pbest)) pbest = i; + // DIAGNOSTIC (gated): same fields as the peak-first branch below, plus + // `fit` = whether the ceiling was met (0 => the budget was below even the + // min-peak schedule, so this fell back to global min-flops). Use the + // per-term chosen_peak_gb to calibrate peak_threshold under a perf-first + // (dense_time_space) objective. + if (pbest >= 0 && std::getenv("SEQUANT_SELROOT_DEBUG")) { + double gmin = std::numeric_limits::max(); + for (auto const& p : rootf) gmin = std::min(gmin, p.flops); + std::cerr << "[selroot] chosen_flops=" << rootf[pbest].flops + << " chosen_peak_gb=" << (peak_bytes(rootf[pbest].peak) / 1e9) + << " fit=" << (fit ? 1 : 0) << " nfront=" << rootf.size() + << " global_min_flops=" << gmin << "\n"; + } + return pbest; + } int best = -1; + bool any_feasible = false; for (int i = 0; i < static_cast(rootf.size()); ++i) - if (rootf[i].peak <= thresh && - (best < 0 || rootf[i].flops < rootf[best].flops || - (rootf[i].flops == rootf[best].flops && - rootf[i].peak < rootf[best].peak))) - best = i; + if (peak_bytes(rootf[i].peak) <= peak_threshold) { + any_feasible = true; + if (best < 0 || rootf[i].flops < rootf[best].flops || + (rootf[i].flops == rootf[best].flops && + rootf[i].peak < rootf[best].peak)) + best = i; + } + if (!any_feasible) { + // Infeasible: no schedule fits the budget. Fall back to min peak. + double minpeak = std::numeric_limits::max(); + for (int i = 0; i < static_cast(rootf.size()); ++i) + if (rootf[i].peak < minpeak) { + minpeak = rootf[i].peak; + best = i; + } + } + // DIAGNOSTIC (gated): the DP's OWN chosen frontier cost (.flops is the + // roofline exec cost, volatile-weighted -- the true objective), plus the + // global-min .flops over the WHOLE frontier (feasible or not) so a + // non-min-feasible selection is visible. Sum the printed chosen_flops + // across terms to test threshold monotonicity of the DP's real objective. + if (best >= 0 && std::getenv("SEQUANT_SELROOT_DEBUG")) { + double gmin = std::numeric_limits::max(); + for (auto const& p : rootf) gmin = std::min(gmin, p.flops); + std::cerr << "[selroot] chosen_flops=" << rootf[best].flops + << " chosen_peak_gb=" << (peak_bytes(rootf[best].peak) / 1e9) + << " feasible=" << (any_feasible ? 1 : 0) + << " nfront=" << rootf.size() << " global_min_flops=" << gmin + << "\n"; + } + return best; + } + + EvalSequence reconstruct(Context const& ctx, + container::vector const& st) const { + std::size_t const root = (std::size_t{1} << ctx.nt) - 1; + int const best = select_root(ctx, st); // Recursive back-pointer walk: at (n, B, idx) read the chosen front point, // form child context C = B | aprime, recurse in lp_first order. std::function build = @@ -856,7 +1382,7 @@ struct PeakBatchedModel { if (std::popcount(n) == 1) return EvalSequence{static_cast(std::countr_zero(n))}; BFrontPoint const& r = st[n][B][idx]; - std::size_t const C = B | r.aprime; + std::size_t const C = ctx.descend(B, r.aprime); std::size_t const fs = r.lp_first ? r.lp : r.rp; int const fi = r.lp_first ? r.lp_idx : r.rp_idx; std::size_t const ss = r.lp_first ? r.rp : r.lp; @@ -869,6 +1395,516 @@ struct PeakBatchedModel { }; return build(root, 0, best); } + + /// \brief Forest-level external seed: size the min-flops factorization with a + /// single external mode \p seed_axis sliced into the ROOT batch + /// context, and confirm the slice is work-neutral. + /// + /// An external mode is contracted at NO node, so batching it is a + /// single outer block-loop over the whole tree: work-neutral (no flops + /// change), scaling every carrying node's footprint by ~block/extent. This + /// re-solves the DP with \p seed_axis admitted as batchable and given a + /// block-sized batch target (\ref sliced_footprints only shrinks a mode for + /// which \c is_batchable is true, so the base table cannot be swap-read -- + /// see the 2026-07-20 external-mode-batching design note), at an INFINITE + /// budget so the seeded and unseeded root selections both resolve to the SAME + /// global-min-flops factorization. That isolates the slice: the returned peak + /// is that factorization's footprint WITH \p seed_axis sliced, and + /// work-neutrality is verified as \c seeded_flops == \c unseeded_flops. + /// + /// Returns false (and leaves \p out_seeded_peak_bytes untouched) if the mode + /// is not recognized as a batchable external mode, if the selection fails, + /// or if the slice is NOT work-neutral (\c seeded_flops != \c + /// unseeded_flops -- the seed trips \ref charge_batch_recompute on a + /// subtree that does not carry it, i.e. the external mode is not carried on + /// every node; declining rather than emitting keeps the emitted schedule + /// honest, since a non-carrying subtree would be recomputed per block). + /// Genuine forest external modes of the CSV/PNO residual giants (the + /// external occ carried on every composite PNO leg up to the root) ARE + /// carried on every node and slice free. + /// \brief Forest-level external seed, SINGLE or JOINT: size the min-flops + /// factorization with EVERY external mode in \p seed_axes sliced + /// simultaneously into the ROOT batch context, and confirm the joint slice + /// is work-neutral. Slicing several external modes at once is a NESTED outer + /// block-loop over the whole tree, still contracted at no node, so the + /// carrying nodes' footprint scales by the PRODUCT of each mode's + /// block/extent -- e.g. both external occ i,j of a doubles residual give + /// ~(block_i/ext_i)*(block_j/ext_j), a strictly bigger drop than either + /// alone. \p block_of gives the per-mode block (batch target) size. + /// + /// The single-mode case is \p seed_axes of size 1. Work-neutrality is checked + /// for the JOINT mask (\c seeded_flops == \c unseeded_flops): every mode must + /// be carried on every node it enters, else the DP charges batch recompute + /// and the joint seed is declined. When declined, \p out_*_flops (if + /// non-null) still report the measured flops so a caller can see WHY. + /// + /// \p out_seeded_flops / \p out_unseeded_flops (optional) receive the two + /// root-frontier flop counts (equal iff work-neutral); used by the D1.3 + /// selection-policy experiment to tabulate work-neutrality directly. + /// + /// \note RETIRED-IN-PLACE (S3.4). Superseded by the node-level external-mode + /// placement in \ref reconstruct_batched_modes (the `node_level_placement` + /// branch), which generalizes this root-only, work-neutral-or-decline seed to + /// per-node placement with hoist-vs-recompute for enclosed non-carriers. Kept + /// only for the order_aware_recompute==OFF regime and the direct diagnostic + /// probes in test_eval_dryrun.cpp; slated for deletion in the clean master + /// reimplementation of the external batching path. + template + bool seeded_forest_peak( + TensorNetwork const& network, TIdxs const& tidxs, + container::svector const& seed_axes, + std::function const& block_of, + double& out_seeded_peak_bytes, double* out_seeded_flops = nullptr, + double* out_unseeded_flops = nullptr) const { + if (seed_axes.empty()) return false; + PeakBatchedModel probe = *this; + // Isolate the slice from the perf-first feasibility filter: at +inf every + // point is feasible so both root selections take the global-min-flops + // factorization, and any flops difference is attributable to the slice + // alone (a genuine external mode has none). + probe.peak_threshold = std::numeric_limits::infinity(); + // The work-neutrality guard below compares seeded vs unseeded flops and + // DEPENDS on the flat charge firing for a subtree that does not carry the + // seed -- which is exactly the case order_aware_recompute exempts. Keep the + // probe on the flat rule, or the guard goes vacuously true and admits seeds + // whose invariant subtrees the runtime would rebuild per block. + probe.order_aware_recompute = false; + container::svector seed_labels; + for (auto const& ax : seed_axes) + seed_labels.push_back(std::wstring(ax.full_label())); + auto is_seed = [seed_labels](Index const& ix) { + std::wstring const l(ix.full_label()); + for (auto const& s : seed_labels) + if (s == l) return true; + return false; + }; + auto base_batchable = probe.is_batchable_contracted_index; + auto base_batch = probe.batch; + probe.is_batchable_contracted_index = [base_batchable, + is_seed](Index const& ix) { + return (base_batchable && base_batchable(ix)) || is_seed(ix); + }; + probe.batch = [base_batch, is_seed, block_of](Index const& ix) { + if (is_seed(ix)) return block_of(ix); + return base_batch ? base_batch(ix) : std::size_t{0}; + }; + auto pctx = probe.build_context(network, tidxs); + std::size_t seed = 0; // JOINT root-batch mask over all seed modes + for (auto const& lbl : seed_labels) { + std::size_t k_seed = pctx.m; // sentinel == not found + for (std::size_t k = 0; k < pctx.m; ++k) + if (std::wstring(pctx.batchable_modes[k].full_label()) == lbl) { + k_seed = k; + break; + } + if (k_seed >= pctx.m || !probe.is_external_mode(pctx, k_seed)) + return false; + seed |= (std::size_t{1} << k_seed); + } + auto pst = solve_single_term(probe, network, tidxs, pctx); + std::size_t const proot = (std::size_t{1} << pctx.nt) - 1; + int const best0 = probe.select_root(pctx, pst, 0); + int const bestS = probe.select_root(pctx, pst, seed); + if (best0 < 0 || bestS < 0) return false; + double const unseeded_flops = pst[proot][0][best0].flops; + double const seeded_flops = pst[proot][seed][bestS].flops; + if (out_unseeded_flops) *out_unseeded_flops = unseeded_flops; + if (out_seeded_flops) *out_seeded_flops = seeded_flops; + // Work-neutrality guard (load-bearing): if any seed enters the root batch + // context of a subtree that does NOT carry it, the DP charges that subtree + // nbatches(seed) recompute, so seeded_flops > unseeded_flops. Decline in + // that case rather than mis-report a "free" slice. + if (seeded_flops != unseeded_flops) return false; + out_seeded_peak_bytes = pst[proot][seed][bestS].peak * probe.numeric_size; + return true; + } + + /// Modeled peak (in elements) of the chosen back-pointer subtree rooted at + /// subset \p n, read at DP schedule cell \p Bsched and sized under the + /// explicit union mask \p Usize. This is the reusable, node-level re-price of + /// the \c sim recursion in \ref reconstructed_batched_peak: it follows the + /// SAME back-pointer walk (children/order/aprime chosen by the DP) and the + /// SAME staged-peak algebra (\c stage_first / \c stage_second / \c stage_form + /// with the resident-scan \c res and accumulation \c contrib terms), but + /// sizes every subset from \p Usize instead of from the schedule cell's + /// union. With + /// \p Usize == \c cell_union(Bsched) it reproduces \c sim exactly; the + /// phase-2 external-placement pass calls it with EXTERNAL mode bits OR-ed + /// into \p Usize to price a node's subtree as if those modes were sliced + /// there. + /// + /// The schedule cell \p Bsched (used only to index \c st for the + /// back-pointer) never carries an external bit -- an external mode is + /// contracted at no node and so never enters any \c aprime -- so \c + /// ctx.descend on \p Bsched stays a valid precomputed cell. Only \p Usize + /// carries injected externals, and it flows into the footprint tables by + /// bitmask (\ref Context::sz_u), which need no ordered cell. + double subtree_peak(Context const& ctx, container::vector const& st, + std::size_t n, std::size_t Bsched, std::size_t Usize, + int idx) const { + if (std::popcount(n) == 1) return ctx.sz_u(Usize, n); + auto const& r = st[n][Bsched][idx]; + std::size_t const C = ctx.descend(Bsched, r.aprime); + std::size_t const Uc = Usize | r.aprime; + std::size_t const f = r.lp_first ? r.lp : r.rp; + int const fi = r.lp_first ? r.lp_idx : r.rp_idx; + std::size_t const s = r.lp_first ? r.rp : r.lp; + int const si = r.lp_first ? r.rp_idx : r.lp_idx; + double const peak_f = subtree_peak(ctx, st, f, C, Uc, fi); + double const peak_s = subtree_peak(ctx, st, s, C, Uc, si); + double const res = + (order_aware_recompute && r.aprime != 0) ? ctx.sz_u(Usize, n) : 0.0; + double const stage_first = ctx.Lof_u(Uc, s) + peak_f + res; + double const stage_second = ctx.sz_u(Uc, f) + peak_s + res; + double const contrib = + (r.aprime != 0) ? accumulation_factor * ctx.sz_u(Usize, n) : 0.0; + double const stage_form = + ctx.sz_u(Uc, f) + ctx.sz_u(Uc, s) + ctx.sz_u(Usize, n) + contrib; + return std::max({stage_first, stage_second, stage_form}); + } + + /// Companion to \ref reconstruct that additionally reports, for each \c -1 + /// (contraction) entry emitted in the returned \c EvalSequence in emission + /// order, the vector of \c Index sliced at that node (\c + /// ctx.batchable_modes[bit] for each set bit of that node's \c aprime). Leaf + /// entries contribute nothing. Does not change \ref reconstruct's own output; + /// the two walks are kept in lock-step so the RPN order and the per-node + /// modes line up. + /// + /// \p out_root_peak_bytes (when non-null) receives the term's REPORTED root + /// peak in bytes: the unseeded footprint normally, or -- for a perf-first + /// over-budget term with \ref batch_spectator_indices on -- the forest-seeded + /// (external-sliced) footprint of the SAME min-flops factorization, when a + /// work-neutral external seed is adopted (see \ref + /// seeded_forest_peak). The contracted min-flops DP / back-pointer walk is + /// untouched; only the reported peak and the emitted \c External tags reflect + /// the seed. + template + std::pair> + reconstruct_batched_modes(Context const& ctx, + container::vector const& st, + TensorNetwork const& network, TIdxs const& tidxs, + double* out_root_peak_bytes = nullptr) const { + std::size_t const root = (std::size_t{1} << ctx.nt) - 1; + int const best = select_root(ctx, st); // shared helper (see below) + // Term-level "should batch external modes" gate: emit only when the + // perf-first objective is selected AND the selected root's unseeded byte + // peak exceeds peak_threshold, on top of batch_spectator_indices already + // being required below. select_root(ctx, st) above defaults to root_B=0 + // (unseeded); st[root][0][best].peak * numeric_size is the byte-peak + // compared to peak_threshold here. When over budget, select_root has fallen + // back to the GLOBAL min-flops factorization (nothing fit the ceiling), so + // `best` IS the min-flops factorization the forest seed re-sizes. + double const unseeded_root_peak_bytes = + st[root][0][best].peak * numeric_size; + bool const over_budget = batch_spectator_indices && perf_first && + (unseeded_root_peak_bytes > peak_threshold); + + // Forest-level external seeding (D1, hole 1; policy settled by the D1.3 + // experiment): for an over-budget term, JOINTLY seed EVERY work-neutral + // external mode into the root batch context and re-size. An + // external mode is contracted at no node, so slicing it is free (identical + // flops); slicing several at once is a nested outer block-loop whose + // footprint scales by the PRODUCT of each mode's block/extent -- a strictly + // bigger, still-work-neutral peak drop than any single seed (D1.3 evidence: + // both C60 external occ i,j slice free, 1874 -> 1124 GB single -> 674 GB + // joint at block/extent = 72/120). Greedy over ctx.batchable_modes order: + // try to add each external mode to the joint seed set, KEEP it only if the + // resulting joint slice stays work-neutral (seeded_forest_peak succeeds) -- + // so a declined external mode is skipped rather than aborting the search + // (the "first ADOPTABLE" fix), and the adopted set is the maximal + // work-neutral joint. Emit `External` for exactly the adopted modes (a + // selection outcome, not a fixed post-hoc stamp over all external modes). + // The contracted min-flops DP and the back-pointer walk below are untouched + // -- the seed only re-sizes and gates the emit. block size comes from + // batch(seed_axis) (batch_target_size); sweeping it under budget is a + // caller knob (D1.3). + double reported_root_peak_bytes = unseeded_root_peak_bytes; + std::size_t chosen_seed_mask = + 0; // ctx.batchable_modes bits stamped External + // Per-node external placement (S3.3): subset n -> mask of external modes + // sliced AT n. Populated by the phase-2 pass below; the build walk stamps + // `External` per node from this (node-level placement is the whole point), + // instead of the OLD root-global chosen_seed_mask stamp. + std::map placed_at_node; + // Phase-2 replaces the OLD root-level forest seed when the order-aware + // model is on; both flags off (the default) => byte-identical (the OLD + // `else if` path runs exactly as before). + bool const node_level_placement = + order_aware_recompute && batch_spectator_indices; + // Shared child-extraction used by the three back-pointer walks below + // (place, stamp_carrying_descendants, and the build emit walk): given a + // node subset `n`, its enclosing cell `B`, and the chosen frontier index + // `idx`, fetch the frontier point, descend to the children's cell `C`, and + // return the two child subsets/indices in the canonical `lp_first` order. + // Keeping the `(lp_first ? lp : rp)` descent in ONE place removes the + // staleness risk of the three otherwise-duplicated copies drifting apart. + struct ChildFrontier { + std::size_t f; + int fi; + std::size_t s; + int si; + std::size_t C; + }; + auto child_frontier = [&](std::size_t n, std::size_t B, + int idx) -> ChildFrontier { + auto const& r = st[n][B][idx]; + std::size_t const C = ctx.descend(B, r.aprime); + return ChildFrontier{ + r.lp_first ? r.lp : r.rp, r.lp_first ? r.lp_idx : r.rp_idx, + r.lp_first ? r.rp : r.lp, r.lp_first ? r.rp_idx : r.lp_idx, C}; + }; + if (node_level_placement) { + // Node-level external-mode placement (S3.1 phase 2). Walk the chosen + // schedule tree; at each node whose modeled peak exceeds peak_threshold, + // greedily slice a carried EXTERNAL mode when doing so lowers that node's + // subtree peak. An external mode is contracted at no node, so slicing it + // is work-neutral and its loop is LOCAL -- injecting it at node n only + // re-sizes n's subtree. Sizing is threaded as a union mask `Usize` (the + // schedule cell `Bsched` stays a valid precomputed DP cell -- externals + // never enter a contracted-here set, so they need no ordered cell; they + // enter the footprint tables by bitmask via Context::sz_u). The pass does + // NOT touch the contracted DP tables or the min-flops back-pointer tree; + // it only re-sizes and records placements for the emit. + // + // D1 fix (i) -- propagate an adopted external placement DOWN to the + // carrying-subtree descendants. The greedy cascade below adopts an + // external mode at the OUTERMOST over-budget node and threads the slice + // down as `Ueff`; by the time `place` recurses into the (giant) + // descendants their subtree_peak is already <= threshold, so the greedy + // loop never fires for them and `placed_at_node` stays empty there. But + // the runtime slices a node ONLY from that node's OWN `batched_here()` + // External stamp, so a descendant that carries the external mode FREE is + // never sliced and is materialized/cached at full extent. This records + // the adopted bit on every descendant of `n` (walking the chosen + // back-pointer tree, same descent as `place`) whose result carries the + // mode, so the emit walk stamps `External` on them too. Placement ONLY -- + // it does NOT re-size (the slice is already threaded via `Ueff`) and does + // NOT touch the min-flops back-pointer tree. + std::function + stamp_carrying_descendants = [&](std::size_t n, std::size_t Bsched, + int idx, std::size_t bit) -> void { + if (std::popcount(n) == 1) return; + auto const [f, fi, s, si, C] = child_frontier(n, Bsched, idx); + if (std::popcount(f) > 1) { + if (ctx.open_modes[f] & bit) placed_at_node[f] |= bit; + stamp_carrying_descendants(f, C, fi, bit); + } + if (std::popcount(s) > 1) { + if (ctx.open_modes[s] & bit) placed_at_node[s] |= bit; + stamp_carrying_descendants(s, C, si, bit); + } + }; + std::function place = + [&](std::size_t n, std::size_t Bsched, std::size_t Usize, + int idx) -> double { + if (std::popcount(n) == 1) return ctx.sz_u(Usize, n); + std::size_t Ueff = Usize; + // Greedy cascade: while this node's modeled peak is over budget, adopt + // the first carried external mode whose slice lowers it; re-scan until + // under budget or no external helps. + for (;;) { + double const cur = + subtree_peak(ctx, st, n, Bsched, Ueff, idx) * numeric_size; + if (cur <= peak_threshold) break; + bool improved = false; + for (std::size_t k = 0; k < ctx.m; ++k) { + std::size_t const bit = std::size_t{1} << k; + if (!((ctx.external_mask >> k) & 1u)) continue; // external only + if (!((ctx.open_modes[n] >> k) & 1u)) continue; // node carries it + if (Ueff & bit) continue; // not yet sliced + double const trial = + subtree_peak(ctx, st, n, Bsched, Ueff | bit, idx) * + numeric_size; + if (trial < cur) { + Ueff |= bit; + placed_at_node[n] |= bit; + // D1 fix (i): also stamp the carrying-subtree descendants so the + // emit walk annotates (and the runtime slices) the giants that + // carry this external mode FREE below `n`. + stamp_carrying_descendants(n, Bsched, idx, bit); + // NB: the node-level emit reads placed_at_node[n] (per-node), NOT + // chosen_seed_mask, so we do NOT touch chosen_seed_mask here -- + // it stays 0 on this path and emit_external is never consulted + // (that field gates only the OLD `else if` root-global stamp). + improved = true; + break; + } + } + if (!improved) break; + } + auto const& r = st[n][Bsched][idx]; + std::size_t const Uc = Ueff | r.aprime; + auto const [f, fi, s, si, C] = child_frontier(n, Bsched, idx); + double const peak_f = place(f, C, Uc, fi); + double const peak_s = place(s, C, Uc, si); + double const res = + (order_aware_recompute && r.aprime != 0) ? ctx.sz_u(Ueff, n) : 0.0; + double const stage_first = ctx.Lof_u(Uc, s) + peak_f + res; + double const stage_second = ctx.sz_u(Uc, f) + peak_s + res; + double const contrib = + (r.aprime != 0) ? accumulation_factor * ctx.sz_u(Ueff, n) : 0.0; + double const stage_form = + ctx.sz_u(Uc, f) + ctx.sz_u(Uc, s) + ctx.sz_u(Ueff, n) + contrib; + return std::max({stage_first, stage_second, stage_form}); + }; + double const root_peak = place(root, 0, ctx.cell_union(0), best); + reported_root_peak_bytes = root_peak * numeric_size; + } else if (over_budget) { + // LEGACY external path (retired-in-place, S3.4). This root-global forest + // seed is SUBSUMED by the node_level_placement branch above whenever the + // order-aware model is on (order_aware_recompute && batch_spectator_ + // indices): node-level placement generalizes root seeding (the root is + // just the outermost node) AND, unlike this work-neutral-or-decline seed, + // it allows enclosed non-carriers via the hoist-vs-recompute trade. This + // branch survives only for the order_aware_recompute==OFF regime (so the + // pre-order-aware behavior stays byte-identical) and for the direct + // seeded_forest_peak diagnostic probes in test_eval_dryrun.cpp. It, and + // seeded_forest_peak itself, are to be DELETED in the clean master + // reimplementation once the order-aware model is the sole external path. + container::svector adopted; + auto block_of = [this](Index const& ix) { return batch(ix); }; + for (std::size_t k = 0; k < ctx.m; ++k) { + if (!is_external_mode(ctx, k)) continue; + container::svector trial = adopted; + trial.push_back(ctx.batchable_modes[k]); + double trial_peak_bytes = 0.0; + if (seeded_forest_peak(network, tidxs, trial, block_of, + trial_peak_bytes)) { + adopted = std::move(trial); + chosen_seed_mask |= (std::size_t{1} << k); + reported_root_peak_bytes = + std::min(reported_root_peak_bytes, trial_peak_bytes); + } + } + } + bool const emit_external = chosen_seed_mask != 0; + if (out_root_peak_bytes) *out_root_peak_bytes = reported_root_peak_bytes; + container::vector node_axes; + std::function build = + [&](std::size_t n, std::size_t B, int idx) -> EvalSequence { + if (std::popcount(n) == 1) + return EvalSequence{static_cast(std::countr_zero(n))}; + BFrontPoint const& r = st[n][B][idx]; + auto const [fs, fi, ss, si, C] = child_frontier(n, B, idx); + // DP-side recompute accounting for the CHOSEN schedule: for this node, + // which batchable modes does it carry, which ancestor batch loops (B) is + // it inside, and which of those does it NOT carry (esc -> re-executed + // nbatches[k] times, the charge_batch_recompute term)? A node with a + // non-empty escaped set and rf>1 IS charged recompute by the DP; if the + // expensive gC-class nodes show rf==1, the DP is not pricing the runtime + // recompute. Gated; prints via wcerr to render Index labels directly. + if (std::getenv("SEQUANT_DP_RECOMPUTE_DEBUG")) { + std::size_t const esc = B & ~ctx.open_modes[n]; + double rf = 1.0; + std::wstring carried, inside, escaped; + for (std::size_t k = 0; k < ctx.m; ++k) { + std::wstring const lbl = + std::wstring(ctx.batchable_modes[k].full_label()); + if ((ctx.open_modes[n] >> k) & 1u) carried += lbl + L" "; + if ((B >> k) & 1u) inside += lbl + L" "; + if ((esc >> k) & 1u) { + escaped += lbl + L"(x" + std::to_wstring(ctx.nbatches[k]) + L") "; + rf *= ctx.nbatches[k]; + } + } + std::wcerr << L"[dp-recompute] ntensors=" << std::popcount(n) + << L" size=" << ctx.sz(n, B) << L" carries={" << carried + << L"} inside_batch={" << inside << L"} escaped={" << escaped + << L"} rf=" << rf << L"\n"; + } + EvalSequence s = build(fs, C, fi); + EvalSequence b = build(ss, C, si); + s.insert(s.end(), b.begin(), b.end()); + container::svector> modes; + // External modes FIRST (outer), then Contracted (D1 fix (ii)): `modes` + // becomes `ann.axes`, the node's own realized-loop order == the runtime + // pick order (eval.hpp), so emitting External before Contracted realizes + // a co-carried external (occ) mode OUTER of the contracted (aux) loop, + // preventing the scatter from widening the external mode to full extent + // per contracted block. External entries, unlike Contracted ones, do not + // depend on the chosen frontier point's aprime -- an external mode is + // never contracted anywhere, so annotating it is purely informational + // (this node's result happens to carry it), independent of which + // sliced-set the DP chose to slice at this node. Emit follows selection: + // stamp ONLY the mode/modes the adopted feasible schedule actually relied + // on (chosen_seed_mask -- the forest seed that made this over-budget term + // fit / shrink), intersected with the node's carried set, NOT every + // external mode. chosen_seed_mask == 0 (no seed adopted, e.g. + // batch_spectator_indices off, or not over budget, or the slice was not + // work-neutral) => emit_mask_n == 0 => no External entries => the + // push-order flip is a true no-op for `modes` in that case. This is + // NOT the same condition as "node_level_placement off": the LEGACY + // `else if (emit_external)` branch just below can set chosen_seed_mask + // (hence emit_mask_n) != 0 even with node_level_placement off + // (batch_spectator_indices on, order_aware_recompute off, over budget -- + // the regime MPQC ships today). There, External-before-Contracted + // intentionally realizes the external mode OUTER of a co-carried + // contracted one (same D1 fix (ii) rationale as the node-level-placement + // path): the computed tensor is unchanged, only realization/recompute + // order at co-carrying nodes differs from before this fix. + std::size_t emit_mask_n = 0; + if (node_level_placement) { + // Node-level placement (S3.3): stamp `External` at node n only for the + // modes the phase-2 pass injected AT n. + auto it = placed_at_node.find(n); + if (it != placed_at_node.end()) emit_mask_n = it->second; + } else if (emit_external) { + // OLD root-global forest seed: stamp every carried adopted mode. + emit_mask_n = chosen_seed_mask & ctx.open_modes[n]; + } + for (std::size_t k = 0; k < ctx.m; ++k) + if (emit_mask_n & (std::size_t{1} << k)) + modes.push_back({ctx.batchable_modes[k], BatchModeType::External}); + for (std::size_t k = 0; k < ctx.m; ++k) + if (r.aprime & (std::size_t{1} << k)) + modes.push_back({ctx.batchable_modes[k], BatchModeType::Contracted}); + NodeBatchAnnotation ann; + ann.axes = std::move(modes); + // Order-aware placement/lifetime bridge (A4): emit this node's residency + // (contracted_modes) and effective use count for a later runtime hoist + // pass. Populated ONLY on the ordered (order_aware_recompute) path; + // otherwise the defaults (effective_count == 1, contracted_modes empty) + // keep the OFF path byte-identical. `B` is the node's ordered enclosing + // cell, so all the needed pieces (open_modes, escaped_outer, nbatches) + // are in hand at this frame -- no second structure or pass. + if (order_aware_recompute) { + // Order-aware gate for per-level placement: this node WAS emitted by + // the order-aware model, so it participates in hoist placement even + // when its residency union is empty (a whole-nest invariant -> root + // scope). + ann.order_aware = true; + // contracted_modes = the enclosing CONTRACTED batchable modes this node + // carries open (variant to). The external-only lifetime mask + // (EvalExpr::sliced_modes) cannot express this; per-level placement + // unions the two to reproduce the combined-nest residency. Emitted + // per-occurrence (a function of open_modes[n], hence consistent across + // canonically-equal occurrences), realized loops only (nbatches > 1), + // contracted modes only (excludes external via ctx.external_mask). + // Empty on the OFF path (this whole block is gated on + // order_aware_recompute) and empty for a whole-nest-invariant node. + ann.contracted_modes.clear(); + for (std::size_t k = 0; k < ctx.m; ++k) + if (((ctx.open_modes[n] >> k) & 1u) && + !((ctx.external_mask >> k) & 1u) && ctx.nbatches[k] > 1.0) + ann.contracted_modes.push_back(ctx.batchable_modes[k]); + // effective_count = rf = prod of nbatches[k] over the escaped-outer set + // (enclosing loops the node does NOT carry). With no CSE on this path + // the back-pointer object is a strict tree (single consumer), so this + // per-node rf IS the effective use count. + std::size_t const esc = ctx.escaped_outer(B, ctx.open_modes[n]); + double rf = 1.0; + for (std::size_t k = 0; k < ctx.m; ++k) + if (esc & (std::size_t{1} << k)) rf *= ctx.nbatches[k]; + ann.effective_count = static_cast(std::llround(rf)); + } + node_axes.push_back(std::move(ann)); // one entry per -1, in RPN order + s.push_back(-1); + return s; + }; + auto seq = build(root, 0, best); + return {std::move(seq), std::move(node_axes)}; + } }; /// \brief Achieved minimum peak memory (the DensePeakSize objective value) for @@ -896,10 +1932,23 @@ double peak_cost_batched( TensorNetwork const& network, TIdxs const& tidxs, IdxToSz&& idxsz, std::function const& is_batchable, std::function const& batch_target_size, - std::function const& is_volatile_leaf) { - PeakBatchedModel> model{std::forward(idxsz), - is_batchable, batch_target_size, - is_volatile_leaf}; + std::function const& is_volatile_leaf, + double accumulation_factor = 0.0, bool order_aware_recompute = false) { + PeakBatchedModel> model{ + std::forward(idxsz), + batch_target_size, + is_volatile_leaf, + /* inner_pow */ {}, + /* volatile_weight */ 1.0, + /* machine_balance */ 0.0, + /* fast_mem_elems */ 0.0, + /* block_tiles */ 3.0, + /* block_prefactor */ 1.0, + /* batch_persistent_only */ false, + /* peak_flops_tolerance */ 0.0, + accumulation_factor}; + model.is_batchable_contracted_index = is_batchable; + model.order_aware_recompute = order_aware_recompute; auto ctx = model.build_context(network, tidxs); auto st = solve_single_term(model, network, tidxs, ctx); // root subset's B=0 frontier; its smallest peak is the achieved minimum. @@ -908,6 +1957,99 @@ double peak_cost_batched( return mn; } +/// \brief Result of seeding an external mode into the batched DP's +/// ROOT frontier (forest batching, P1). The seed mode is contracted at NO node +/// (a pure external mode carried on the root and every subtree holding a +/// composite leg that bears it), so slicing it does NOT change total work +/// (flops) but +/// shrinks every intermediate that carries it -- in particular the perf-first +/// PPL W giant -- by ~occ_block/occ. Carries both the unseeded (B=0) and seeded +/// (B={k_seed}) root selections so a caller can report the drop, plus the +/// chosen mode + block size for Task 4 to thread through the public optimize +/// API. +struct SeededBatchedResult { + double unseeded_peak_bytes = 0; + double seeded_peak_bytes = 0; + double unseeded_flops = 0; + double seeded_flops = 0; + std::optional seeded_axis; + std::size_t occ_block = 0; + bool spectator_ok = false; +}; + +/// \brief Runs the batched peak DP once and reports the ROOT-frontier selection +/// at B=0 (unseeded) and at B={k_seed} (the seed mode sliced into the root +/// batch context), under the model's own (perf-first or peak-first) selector. +/// +/// The external mode (e.g. the residual's own external occupied +/// index, carried only as a composite protoindex) is admitted to \c +/// ctx.batchable_modes by \ref batchable_mode_list's proto pass regardless of +/// the base \c is_batchable predicate; but \ref sliced_footprints only shrinks +/// a mode for which \c is_batchable is true. This routine therefore extends the +/// predicate (and gives the seed mode a block-sized batch target \p occ_block) +/// so the DP tables actually slice the seed mode when its bit is set. Every +/// other mode is untouched, and B=0 never sets the seed bit (the external mode +/// is contracted nowhere, so it never enters any node's \c contracted_here and +/// hence never enters any child context at B=0), so the unseeded (B=0) +/// selection is byte-identical to \p base_model's. +/// +/// \p seed_axis MUST be a genuine external mode (asserted via +/// \ref PeakBatchedModel::is_external_mode): carried on the root and +/// contracted at no node. Seeding a contracted mode would mis-size the nodes +/// that contract it, so this asserts rather than silently mis-report. +template +SeededBatchedResult seeded_root_peak_batched( + PeakBatchedModel base_model, TensorNetwork const& network, + TIdxs const& tidxs, Index const& seed_axis, std::size_t occ_block) { + SeededBatchedResult out; + out.occ_block = occ_block; + std::wstring const seed_label(seed_axis.full_label()); + + auto base_batchable = base_model.is_batchable_contracted_index; + auto base_batch = base_model.batch; + base_model.is_batchable_contracted_index = [base_batchable, + seed_label](Index const& ix) { + return (base_batchable && base_batchable(ix)) || + std::wstring(ix.full_label()) == seed_label; + }; + base_model.batch = [base_batch, seed_label, occ_block](Index const& ix) { + if (std::wstring(ix.full_label()) == seed_label) return occ_block; + return base_batch ? base_batch(ix) : std::size_t{0}; + }; + + auto ctx = base_model.build_context(network, tidxs); + // Locate the seed mode's bit in ctx.batchable_modes. + std::size_t k_seed = ctx.m; // sentinel == not found + for (std::size_t k = 0; k < ctx.m; ++k) + if (std::wstring(ctx.batchable_modes[k].full_label()) == seed_label) { + k_seed = k; + break; + } + SEQUANT_ASSERT(k_seed < ctx.m && + "seed mode not recognized as a batchable mode"); + // SAFEGUARD: only honor a genuine external mode. + out.spectator_ok = base_model.is_external_mode(ctx, k_seed); + SEQUANT_ASSERT(out.spectator_ok && + "seed mode is not a pure external mode " + "(open on the root AND contracted at no node)"); + out.seeded_axis = ctx.batchable_modes[k_seed]; + + auto st = solve_single_term(base_model, network, tidxs, ctx); + std::size_t const root = (std::size_t{1} << ctx.nt) - 1; + std::size_t const seed = std::size_t{1} << k_seed; + + int const best0 = base_model.select_root(ctx, st, 0); + int const bestS = base_model.select_root(ctx, st, seed); + SEQUANT_ASSERT(best0 >= 0 && bestS >= 0); + auto const& p0 = st[root][0][best0]; + auto const& pS = st[root][seed][bestS]; + out.unseeded_peak_bytes = p0.peak * base_model.numeric_size; + out.seeded_peak_bytes = pS.peak * base_model.numeric_size; + out.unseeded_flops = p0.flops; + out.seeded_flops = pS.flops; + return out; +} + /// \brief Independent memory-simulation recomputation of the chosen batched /// reconstruction's model-A peak. Builds \ref PeakBatchedModel, runs /// \ref solve_single_term for the back-pointer table, and recomputes the @@ -918,50 +2060,41 @@ double reconstructed_batched_peak( TensorNetwork const& network, TIdxs const& tidxs, IdxToSz&& idxsz, std::function const& is_batchable, std::function const& batch_target_size, - std::function const& is_volatile_leaf) { - PeakBatchedModel> model{std::forward(idxsz), - is_batchable, batch_target_size, - is_volatile_leaf}; + std::function const& is_volatile_leaf, + double accumulation_factor = 0.0, bool order_aware_recompute = false) { + PeakBatchedModel> model{ + std::forward(idxsz), + batch_target_size, + is_volatile_leaf, + /* inner_pow */ {}, + /* volatile_weight */ 1.0, + /* machine_balance */ 0.0, + /* fast_mem_elems */ 0.0, + /* block_tiles */ 3.0, + /* block_prefactor */ 1.0, + /* batch_persistent_only */ false, + /* peak_flops_tolerance */ 0.0, + accumulation_factor}; + model.is_batchable_contracted_index = is_batchable; + model.order_aware_recompute = order_aware_recompute; auto ctx = model.build_context(network, tidxs); auto st = solve_single_term(model, network, tidxs, ctx); auto const nt = network.tensors().size(); - // Simulate the peak of evaluating subtree n at ancestor context B by walking - // the chosen back-pointers. A leaf is resident at its own size. For an - // internal node, with child context C = B | aprime, evaluate the lp_first - // child fully (peak: its own simulated peak), then hold its result (sized at - // C) while evaluating the second child (whose inputs co-reside at Lof), then - // both results co-reside while the parent result (sized at B) is formed. - // Re-derives the chosen reconstruction's peak by following the back-pointer - // tree (contexts/orders chosen by the DP) and recomputing each child's peak - // via recursion, rather than reading the DP's minimized st[*].peak table. - // The per-node combination (stage_first/stage_second/stage_form) uses the - // same staged-peak formula as the DP's lpf. What this validates independently - // is the back-pointer walk itself (which children, order, context). The - // Task-2/Task-3 batched oracle is the independent guard on the staged-peak - // algebra. - auto sim = [&](auto&& self, std::size_t n, std::size_t B, int idx) -> double { - if (std::popcount(n) == 1) return ctx.sz(n, B); - auto const& r = st[n][B][idx]; - std::size_t const C = B | r.aprime; - std::size_t const f = r.lp_first ? r.lp : r.rp; // evaluated first - int const fi = r.lp_first ? r.lp_idx : r.rp_idx; - std::size_t const s = r.lp_first ? r.rp : r.lp; // evaluated second - int const si = r.lp_first ? r.rp_idx : r.lp_idx; - double const peak_f = self(self, f, C, fi); - double const peak_s = self(self, s, C, si); - // While the first child evaluates, the second child's leaf inputs sit - // resident (Lof(s, C)). While the second child evaluates, the first - // child's result sits resident (sz(f, C)). When both results exist, the - // parent result (sz(n, B)) is materialized alongside them. - double const stage_first = ctx.Lof(s, C) + peak_f; - double const stage_second = ctx.sz(f, C) + peak_s; - double const stage_form = ctx.sz(f, C) + ctx.sz(s, C) + ctx.sz(n, B); - return std::max({stage_first, stage_second, stage_form}); - }; - + // Simulate the peak of evaluating the chosen root subtree by walking the + // chosen back-pointers (which children, order, context chosen by the DP) and + // recomputing each node's peak via the staged-peak algebra + // (stage_first/stage_second/stage_form with the resident-scan res term and + // the accumulation contrib term), rather than reading the DP's minimized + // st[*].peak table. This is exactly PeakBatchedModel::subtree_peak sized at + // the schedule cell's own union (no external injection): sizing under + // cell_union(B) reproduces the historical `sim` byte-for-byte. What this + // validates independently is the back-pointer walk itself; the batched oracle + // is the independent guard on the staged-peak algebra. std::size_t const root = (std::size_t{1} << nt) - 1; - return sim(sim, root, 0, pareto_best(st[root][0])); + int const best = pareto_best(st[root][0]); + return model.subtree_peak(ctx, st, root, /*Bsched=*/0, + /*Usize=*/ctx.cell_union(0), best); } /// \brief Compile-time concept for a single-term-DP cost model. diff --git a/SeQuant/core/optimize/optimize.cpp b/SeQuant/core/optimize/optimize.cpp index fb64d00bc5..8cbefd81bc 100644 --- a/SeQuant/core/optimize/optimize.cpp +++ b/SeQuant/core/optimize/optimize.cpp @@ -1,14 +1,18 @@ #include #include #include +#include #include #include #include #include +#include +#include #include #include #include #include +#include #include #include @@ -18,10 +22,12 @@ #include #include +#include #include #include #include #include +#include #include #include #include @@ -87,38 +93,58 @@ void log_chosen_factorization(ExprPtr const& result, /// Optimize a Product that contains only Tensor and scalar factors. ExprPtr opt_pure_product(Product const& prod, OptimizeOptions const& opts) { bool const subnet_cse = opts.CSE.subnet; - CostParams const cost{opts.batch_policy.is_volatile_leaf, - opts.volatile_weight, - opts.footprint_weight, - opts.peak_flops_tolerance, - opts.roofline, - opts.batch_policy.accumulation_factor, - opts.prune_outer_products}; + // Build the cost knobs field-by-field from OptimizeOptions / its BatchPolicy. + // Batching config (both role predicates, batch_target_size, inner_pow, + // batch_persistent_only) now travels on CostParams rather than as loose args. + CostParams cost; + cost.is_volatile_leaf = opts.batch_policy.is_volatile_leaf; + cost.volatile_weight = opts.volatile_weight; + cost.footprint_weight = opts.footprint_weight; + cost.peak_flops_tolerance = opts.peak_flops_tolerance; + cost.roofline = opts.roofline; + cost.accumulation_factor = opts.batch_policy.accumulation_factor; + cost.peak_threshold = opts.batch_policy.peak_threshold; + cost.prune_outer_products = opts.prune_outer_products; + cost.batch_spectator_indices = opts.batch_policy.batch_spectator_indices; + cost.order_aware_recompute = opts.batch_policy.order_aware_recompute; + cost.is_batchable_contracted_index = + opts.batch_policy.is_batchable_contracted_index; + cost.is_batchable_external_index = + opts.batch_policy.is_batchable_external_index; + cost.batch_target_size = opts.batch_policy.batch_target_size; + cost.inner_pow = opts.inner_pow; + cost.batch_persistent_only = opts.batch_policy.persistent_only; + // Filled only by the DensePeakSizeBatched arm below (via out_axes); every + // other objective leaves it empty, so the term_batch_axes insertion at the + // end is then a no-op-shaped empty-vector entry (harmless: Task 3.3 only + // consumes entries for summands the batched objective actually annotated). + container::vector node_axes; auto run = [&]() -> ExprPtr { if (opts.objective_function == ObjectiveFunction::DenseFLOPs) return opt::single_term_opt( - prod, opts.idx_to_extent, subnet_cse, cost, - opts.batch_policy.is_batchable_index, - opts.batch_policy.batch_target_size, opts.inner_pow); + prod, opts.idx_to_extent, subnet_cse, cost); if (opts.objective_function == ObjectiveFunction::DenseSize) return opt::single_term_opt( + prod, opts.idx_to_extent, subnet_cse, cost); + if (opts.objective_function == ObjectiveFunction::DenseSpaceTime) + return opt::single_term_opt( + prod, opts.idx_to_extent, subnet_cse, cost); + if (opts.objective_function == ObjectiveFunction::DenseTimeSpace) + return opt::single_term_opt( + prod, opts.idx_to_extent, subnet_cse, cost); + if (opts.objective_function == ObjectiveFunction::DenseSpaceTimeBatched) + return opt::single_term_opt( prod, opts.idx_to_extent, subnet_cse, cost, - opts.batch_policy.is_batchable_index, - opts.batch_policy.batch_target_size, opts.inner_pow); - if (opts.objective_function == ObjectiveFunction::DensePeakSize) - return opt::single_term_opt( - prod, opts.idx_to_extent, subnet_cse, cost, - opts.batch_policy.is_batchable_index, - opts.batch_policy.batch_target_size, opts.inner_pow); + opts.term_batch_axes ? &node_axes : nullptr); SEQUANT_ASSERT(opts.objective_function == - ObjectiveFunction::DensePeakSizeBatched); - return opt::single_term_opt( + ObjectiveFunction::DenseTimeSpaceBatched); + return opt::single_term_opt( prod, opts.idx_to_extent, subnet_cse, cost, - opts.batch_policy.is_batchable_index, - opts.batch_policy.batch_target_size, opts.inner_pow, - opts.batch_policy.persistent_only); + opts.term_batch_axes ? &node_axes : nullptr); }; ExprPtr result = run(); + if (opts.term_batch_axes) + (*opts.term_batch_axes)[result.get()] = std::move(node_axes); if (std::getenv("SEQUANT_FACTORIZER_DEBUG")) log_chosen_factorization(result, opts); return result; @@ -248,6 +274,14 @@ ExprPtr optimize(ExprPtr const& expr, OptimizeOptions opts) { /*parallel_outer=*/true); } +OptimizeResult optimize_result(ExprPtr const& expr, OptimizeOptions opts) { + if (!opts.idx_to_extent) opts.idx_to_extent = default_idx_to_size(); + OptimizeResult res; + res.expr = optimize_impl(expr, opts, opts.reorder == ReorderSum::Reorder, + /*parallel_outer=*/true); + return res; +} + ResultExpr& optimize(ResultExpr& expr, OptimizeOptions opts) { expr.expression() = optimize(expr.expression(), std::move(opts)); return expr; diff --git a/SeQuant/core/optimize/optimize.hpp b/SeQuant/core/optimize/optimize.hpp index caa6952d44..e40da552c2 100644 --- a/SeQuant/core/optimize/optimize.hpp +++ b/SeQuant/core/optimize/optimize.hpp @@ -2,10 +2,19 @@ #define SEQUANT_OPTIMIZE_OPTIMIZE_HPP #include +#include #include +#include +#include + namespace sequant { +/// Result of \ref optimize_result: the optimized expression. +struct OptimizeResult { + ExprPtr expr; +}; + /// Optimize the expression for lower evaluation cost. /// /// \param expr Expression to be optimized. @@ -16,6 +25,11 @@ namespace sequant { /// \return Optimized expression. ExprPtr optimize(ExprPtr const& expr, OptimizeOptions opts = {}); +/// \copydoc optimize(ExprPtr const&, OptimizeOptions) +/// +/// \return The optimized expression, wrapped in \ref OptimizeResult. +OptimizeResult optimize_result(ExprPtr const& expr, OptimizeOptions opts = {}); + /// \copydoc optimize(ExprPtr const&, OptimizeOptions) ResultExpr& optimize(ResultExpr& expr, OptimizeOptions opts = {}); diff --git a/SeQuant/core/optimize/options.hpp b/SeQuant/core/optimize/options.hpp index b5813717db..8610ca8296 100644 --- a/SeQuant/core/optimize/options.hpp +++ b/SeQuant/core/optimize/options.hpp @@ -2,47 +2,91 @@ #define SEQUANT_CORE_OPTIMIZE_OPTIONS_HPP #include +#include +#include #include #include +#include +#include +#include +#include namespace sequant { class Index; class Tensor; +class Expr; /// Objective function to minimize in single-term and top-level optimize /// routines. The `Dense*` models assume dense tensors: /// - `DenseFLOPs` counts floating-point operations. /// - `DenseSize` counts result-tensor storage elements (summed over /// intermediates) -- a gross-traffic proxy, not a peak. -/// - `DensePeakSize` minimizes peak memory: the maximum over the evaluation -/// schedule of the combined size of all simultaneously-live tensors -/// (intermediates AND resident input leaves, the all-co-resident model), -/// and it chooses the evaluation order that minimizes that peak. Unlike the -/// order-independent `DenseFLOPs`/`DenseSize`, the contraction order is a -/// real lever here. NOTE: `DensePeakSize` does not yet support +/// - `DenseSpaceTime` (formerly `DensePeakSize`; that name is kept as a +/// deprecated alias) is PEAK-FIRST, perf-second: it minimizes peak memory -- +/// the maximum over the evaluation schedule of the combined size of all +/// simultaneously-live tensors (intermediates AND resident input leaves, the +/// all-co-resident model) -- and breaks ties by the roofline perf cost. +/// Unlike the order-independent `DenseFLOPs`/`DenseSize`, the contraction +/// order is a real lever here. NOTE: does not yet support /// common-subexpression elimination (`CSEOptions::subnet` must be false). -/// - `DensePeakSizeBatched` extends `DensePeakSize` with a per-index -/// batchability model: each index satisfying -/// `OptimizeOptions::batch_policy.is_batchable_index` is treated as -/// independently sliced to -/// `min(extent, batch_policy.batch_target_size(ix))` elements per index -- +/// - `DenseSpaceTimeBatched` (formerly `DensePeakSizeBatched`) extends +/// `DenseSpaceTime` with a per-index batchability model: each index +/// satisfying `OptimizeOptions::batch_policy.is_batchable_index()` (the +/// derived union of the contracted- and external-role building blocks) is +/// treated as independently sliced to `min(extent, +/// batch_policy.batch_target_size(ix))` elements per index -- /// `batch_target_size` is an upper bound, so this is a conservative /// (over-)estimate of the realized whole-tile batch, which the backend rounds /// *down* to a tile multiple (never above the target; see /// `mode_batches_of_trange1`). The /// DP minimises peak over the worst-case sliced configuration. Only consulted /// by the batched oracle and DP; requires -/// `batch_policy.is_batchable_index` and `batch_policy.batch_target_size` -/// to be set. +/// a batchability role predicate +/// (`batch_policy.is_batchable_contracted_index` and/or +/// `is_batchable_external_index`) and `batch_policy.batch_target_size` to be +/// set. Final selection is threshold-gated by `peak_threshold`. +/// - `DenseTimeSpace` / `DenseTimeSpaceBatched` are the PERF-FIRST, peak-second +/// duals of `DenseSpaceTime` / `DenseSpaceTimeBatched`: they select a +/// factorization by roofline perf first and peak second (same Pareto-frontier +/// and roofline machinery, opposite lexicographic order). Because slicing is +/// perf-neutral, a perf-first primary never prefers a FLOPS-catastrophic +/// factorization merely for its sliceability. Naming: +/// `Dense{Primary}{Secondary}`, `Space` = peak/size, `Time` = perf. +/// +/// `peak_threshold` and these objectives: it is NOT a feasibility gate for +/// ROOT SELECTION (`select_root` does not consult it on the perf-first +/// branch, so it cannot force a flops-catastrophic choice; peak breaks exact +/// ties only). It IS consulted elsewhere, and only by these objectives: +/// `DenseTimeSpaceBatched` emits EXTERNAL batch modes iff +/// `batch_policy.batch_spectator_indices` is set AND the selected root's +/// modeled peak exceeds `peak_threshold` (see +/// `PeakBatchedModel::reconstruct_batched_modes`). CONTRACTED batch modes are +/// emitted regardless of the threshold. So under a time-first objective the +/// threshold is an external-batching trigger, not a space constraint -- and +/// external batching is currently available ONLY under a time-first +/// objective. /// /// Leaves room for `Sparse*` models later. enum class ObjectiveFunction { DenseFLOPs, DenseSize, - DensePeakSize, - DensePeakSizeBatched + /// Peak-first, perf-second. (Formerly `DensePeakSize`.) + DenseSpaceTime, + /// Batched variant of `DenseSpaceTime`. (Formerly `DensePeakSizeBatched`.) + DenseSpaceTimeBatched, + /// Perf-first, peak-second: never prefers a FLOPS-catastrophic factorization + /// for its sliceability. + DenseTimeSpace, + /// Batched variant of `DenseTimeSpace`. + DenseTimeSpaceBatched, + /// Deprecated aliases (same underlying values as the renamed constants above, + /// placed AFTER them so `DenseSpaceTime` keeps `DensePeakSize`'s old value). + /// Kept so existing code and JSON inputs ("dense_peak_size") keep compiling; + /// every existing `== DensePeakSize` guard still catches `DenseSpaceTime`. + DensePeakSize = DenseSpaceTime, + DensePeakSizeBatched = DenseSpaceTimeBatched }; /// Whether to reorder summands so terms with shared intermediates appear @@ -95,8 +139,9 @@ struct CostParams { /// Per-intermediate storage-footprint penalty (DenseFLOPs/DenseSize only; see /// OptimizeOptions::footprint_weight). Not used by the peak objectives. double footprint_weight = 0.0; - /// Relative peak tolerance for the peak objectives' final selection; see - /// OptimizeOptions::peak_flops_tolerance. + /// Relative peak tolerance for DensePeakSize's final selection; see + /// OptimizeOptions::peak_flops_tolerance. Unused by DensePeakSizeBatched, + /// whose final selection is instead threshold-gated by \c peak_threshold. double peak_flops_tolerance = 0.10; /// Roofline parameters for the peak objectives' secondary cost; see /// \ref RooflineParams. machine_balance == 0 => pure-flop tie-break. @@ -105,9 +150,53 @@ struct CostParams { /// DensePeakSizeBatched; see BatchPolicy::accumulation_factor. 0 (default) = /// no penalty. double accumulation_factor = 0.0; + /// Peak-memory budget in BYTES for the batched objectives; see + /// BatchPolicy::peak_threshold. Space-first (DenseSpaceTimeBatched): a hard + /// feasibility gate on root selection; +infinity (default) => every schedule + /// feasible => no batching. Time-first (DenseTimeSpaceBatched): NOT a gate on + /// root selection -- it only triggers EXTERNAL mode emission. + double peak_threshold = std::numeric_limits::infinity(); /// Prune disconnected (outer-product) subsets from the single-term DP; see /// OptimizeOptions::prune_outer_products. true (default) = prune. bool prune_outer_products = true; + /// Gate for external-index batching; see + /// BatchPolicy::batch_spectator_indices. false (default) => \ref + /// opt::detail::PeakBatchedModel::reconstruct_batched_modes emits no + /// \c BatchModeType::External entries, so every existing (non-external-aware) + /// caller stays byte-identical. NOTE this is necessary but not sufficient: + /// reconstruct_batched_modes also requires a TIME-FIRST objective + /// (DenseTimeSpaceBatched) and a selected-root peak above \c peak_threshold. + bool batch_spectator_indices = false; + + /// Enable the order-aware multilevel recompute cost model (resident-scan peak + /// + ordered-key flops recompute). false (default) => byte-identical + /// set-keyed DP. Only the batched objectives consult it. + bool order_aware_recompute = false; + + /// Spaces batchable in the CONTRACTED role, threaded from + /// BatchPolicy::is_batchable_contracted_index. Building block consumed by the + /// single-term optimizer's DP contracted-role filter (a mode of such a space + /// is sliced where it is summed). Declared adjacent to its external + /// companion. Defaults to decline every index. + std::function is_batchable_contracted_index = + [](Index const&) { return false; }; + /// Spaces batchable in the EXTERNAL role, threaded from + /// BatchPolicy::is_batchable_external_index. Defaults to decline every index; + /// a caller that wants external batching sets it explicitly (there is no + /// fallback to the contracted-role predicate). + std::function is_batchable_external_index = + [](Index const&) { return false; }; + /// Per-index per-batch slice size (an upper bound) for a batchable index; + /// threaded from BatchPolicy::batch_target_size. Consulted only by the + /// batched objectives. Empty (default) => no target => no slicing. + std::function batch_target_size = {}; + /// k-aware inner (CSV/PNO composite) extent applied by every cost counter; + /// threaded from OptimizeOptions::inner_pow. Empty => composites sized by the + /// idxsz provider (k=1). REQUIRED whenever the network has composite indices. + std::function inner_pow = {}; + /// When true, only persistent (volatile-leaf-free) subnetworks are batched; + /// threaded from BatchPolicy::persistent_only. Batched objectives only. + bool batch_persistent_only = false; }; /// A type-erased provider mapping an Index to its extent. Used by the public @@ -141,15 +230,21 @@ struct OptimizeOptions { /// should return the k-th power mean of the per-pair domain /// (\c (Sum_pairs d^k / nocc^N)^(1/k)), so that the product over the group /// times the outer \c nocc^N equals the true block-sparse volume - /// \c Sum_pairs d^k. If empty, composites fall back to \ref idx_to_extent - /// (k=1), which grossly under-sizes multi-composite tensors. - std::function inner_pow = {}; + /// \c Sum_pairs d^k. REQUIRED whenever the network has composite indices: + /// leaving it empty no longer silently falls back to \ref idx_to_extent + /// (which grossly mis-sized multi-composite tensors and inverted + /// factorization choices, e.g. a 4-PAO integral) -- the sizing code throws + /// instead. No default: pass an explicit no-op only for composite-free work. + std::function inner_pow; - /// Batchability policy: bundles the three per-index and per-leaf predicates - /// that govern batched evaluation. All three fields default to empty (no - /// batchable indices, no volatile leaves). The three sub-fields are: - /// - `is_batchable_index`: marks an Index as living in a batchable space - /// (e.g. DF/RI aux; = the eval cache's accept_aux). + /// Batchability policy: bundles the per-index and per-leaf predicates + /// that govern batched evaluation. All predicate fields default to empty (no + /// batchable indices, no volatile leaves). The key sub-fields are: + /// - `is_batchable_contracted_index` / `is_batchable_external_index`: the + /// two role building blocks marking an Index as living in a batchable + /// space in the contracted / external role. The derived union + /// `is_batchable_index()` (= the eval cache's accept_aux) is the runtime + /// accept. /// - `batch_target_size`: per-index slice size (an upper bound); a sliced /// batchable index ix contributes min(extent, batch_target_size(ix)), a /// conservative over-estimate of the realized tile-floored batch. Only @@ -164,9 +259,13 @@ struct OptimizeOptions { /// Real-valued weight on the cost of each volatile contraction (re-evaluated /// on every replay of the network), while persistent (volatile-independent) /// contractions are counted once. Conceptually the expected number of - /// replays. Default 1.0 (no change). Only consulted when - /// batch_policy.is_volatile_leaf is non-empty and objective_function == - /// ObjectiveFunction::DenseFLOPs. + /// replays. Default 1.0 (no change). Consulted whenever + /// batch_policy.is_volatile_leaf is non-empty, by DenseFLOPs AND by the peak + /// objectives: it scales the per-contraction cost in all of them (see + /// AdditiveModel and PeakModel/PeakBatchedModel's `cflops`). That cost is the + /// PRIMARY mode for the time-first objectives (so volatile_weight materially + /// changes which factorization DenseTimeSpace* picks) but only the tie-break + /// among equal-peak schedules for the space-first ones. double volatile_weight = 1.0; /// Relative peak tolerance for the peak objectives' final selection: among @@ -175,8 +274,9 @@ struct OptimizeOptions { /// 0 = strict peak-min (flop tie-break only on exact peak ties). The default /// 0.10 trades up to a 10% peak increase for a (often much larger) flop /// reduction -- e.g. forming a persistent 4-PNO integral instead of - /// recomputing a particle-ladder. Only consulted by DensePeakSize / - /// DensePeakSizeBatched. + /// recomputing a particle-ladder. Only consulted by DensePeakSize. + /// DensePeakSizeBatched's final selection is instead threshold-gated by + /// \ref BatchPolicy::peak_threshold. double peak_flops_tolerance = 0.10; /// Per-intermediate memory-footprint penalty added to the single-term @@ -210,6 +310,18 @@ struct OptimizeOptions { /// DensePeakSizeBatched. RooflineParams roofline = {}; + /// Optional out-channel: when non-null, optimize() records for each + /// top-level summand's optimized ExprPtr its per-contraction-node + /// sliced-sets (RPN / post-order, left-first, matching single_term_opt's + /// Product construction). Consumed by binarize via BinarizationOptions to + /// annotate eval nodes. Default null => no behavior change. Only populated + /// when objective_function == ObjectiveFunction::DensePeakSizeBatched; + /// every other objective leaves any pre-existing map entry for that + /// summand untouched (it is simply never written). + std::shared_ptr>> + term_batch_axes = {}; + /// Prune disconnected (outer-product) subsets from the single-term /// contraction DP: a subset whose induced subgraph is disconnected under the /// "share a contractible (non-target) index" relation is an outer product the diff --git a/SeQuant/core/optimize/single_term.hpp b/SeQuant/core/optimize/single_term.hpp index a793f41885..cbda1830ba 100644 --- a/SeQuant/core/optimize/single_term.hpp +++ b/SeQuant/core/optimize/single_term.hpp @@ -6,6 +6,7 @@ // objective models live in cost_model.hpp (which includes the detail header). // Including cost_model.hpp here brings both, then the public // single_term_opt entry points below route through run_single_term_opt. +#include #include #include #include @@ -37,19 +38,26 @@ namespace detail { /// structure, NOT on anonymous index identity — so that two subnetworks /// deemed equivalent by the subnet-CSE canonicalization also agree on /// volatility (the CSE path stores one cost per canonical subnet). -/// volatile_weight/footprint_weight apply to DenseFLOPs only; -/// peak_flops_tolerance/roofline apply to the peak objectives only. -/// \param is_batchable_index Predicate marking an index as batchable (sliced); -/// ObjectiveFunction::DensePeakSizeBatched only. -/// \param batch_target_size Per-index per-batch slice size (an upper bound) for -/// batchable indices; ObjectiveFunction::DensePeakSizeBatched only. -/// \param inner_pow Optional k-aware CSV/PNO composite extent applied by every -/// cost counter; see \ref inner_aware_volume. Empty (default) sizes -/// composites by \p idxsz (k=1), which under-counts multi-composite -/// tensors. -/// \param batch_persistent_only When true, only persistent (volatile-leaf-free) -/// subnetworks are batched; ObjectiveFunction::DensePeakSizeBatched -/// only. +/// footprint_weight applies to DenseFLOPs only; volatile_weight applies +/// to DenseFLOPs AND the peak objectives (primary mode for the +/// time-first ones, tie-break for the space-first ones); roofline +/// applies to the peak objectives only; peak_flops_tolerance applies to +/// DensePeakSize only (DensePeakSizeBatched's final selection is instead +/// threshold-gated by peak_threshold, which under the TIME-first batched +/// objective gates only external-mode emission, not root selection). +/// All batching config lives on \p cost (\ref CostParams): +/// is_batchable_contracted_index / is_batchable_external_index mark an +/// index as batchable (sliced) in the contracted / external role, +/// batch_target_size is the per-index per-batch slice-size upper bound, +/// inner_pow is the optional k-aware CSV/PNO composite extent, and +/// batch_persistent_only restricts batching to persistent subnetworks -- +/// all ObjectiveFunction::DensePeakSizeBatched only. +/// \param out_axes When non-null AND \p Metric == +/// ObjectiveFunction::DensePeakSizeBatched, filled with the per-node +/// sliced-sets of the returned sequence's contraction (\c -1) nodes, in +/// the same left-first post-order the sequence itself was built in +/// (see \ref opt::detail::run_single_term_opt_axes). Cleared (left +/// empty) otherwise. Ignored (must stay null) for every other metric. /// \return Optimal evaluation sequence under the chosen cost metric. If there /// are equivalent optimal sequences then the result is the one that /// keeps the order of tensors in the network as original as possible. @@ -58,10 +66,8 @@ template EvalSequence single_term_opt( TensorNetwork const& network, IdxToSz&& idxsz, bool subnet_cse, CostParams const& cost = {}, - std::function const& is_batchable_index = {}, - std::function batch_target_size = {}, - std::function const& inner_pow = {}, - bool batch_persistent_only = false) { + container::vector* out_axes = nullptr) { + if (out_axes) out_axes->clear(); decltype(OptRes::indices) tidxs{}; // Unpack the cost knobs into the names the recurrence arms below use, so the // arms are unchanged. (void)-cast all, since each if-constexpr arm uses only @@ -73,6 +79,12 @@ EvalSequence single_term_opt( double const accumulation_factor = cost.accumulation_factor; RooflineParams const& roofline = cost.roofline; bool const prune_outer_products = cost.prune_outer_products; + // Batching config now lives on CostParams (was loose positional args). + auto const& is_batchable_contracted_index = + cost.is_batchable_contracted_index; + auto const& batch_target_size = cost.batch_target_size; + auto const& inner_pow = cost.inner_pow; + bool const batch_persistent_only = cost.batch_persistent_only; (void)is_volatile_leaf; (void)volatile_weight; (void)footprint_weight; @@ -86,10 +98,13 @@ EvalSequence single_term_opt( // DP's subset bits. size_t volatile_mask = 0; double nr = 1.0; - if constexpr (Metric == ObjectiveFunction::DensePeakSize) { + if constexpr (Metric == ObjectiveFunction::DenseSpaceTime || + Metric == ObjectiveFunction::DenseTimeSpace) { SEQUANT_ASSERT(!subnet_cse && - "subnet_cse not supported with DensePeakSize (Phase 1)"); - (void)is_batchable_index; + "subnet_cse not supported with DenseSpaceTime (Phase 1)"); + SEQUANT_ASSERT(!out_axes && + "out_axes only supported with the batched peak objectives"); + (void)is_batchable_contracted_index; (void)batch_target_size; (void)batch_persistent_only; (void)footprint_weight; // peak objectives use the roofline tie-break @@ -104,9 +119,11 @@ EvalSequence single_term_opt( roofline.block_tiles, roofline.block_prefactor, peak_flops_tolerance}; + model.perf_first = (Metric == ObjectiveFunction::DenseTimeSpace); model.prune_outer_products = prune_outer_products; return run_single_term_opt(model, network, tidxs); - } else if constexpr (Metric == ObjectiveFunction::DensePeakSizeBatched) { + } else if constexpr (Metric == ObjectiveFunction::DenseSpaceTimeBatched || + Metric == ObjectiveFunction::DenseTimeSpaceBatched) { SEQUANT_ASSERT( !subnet_cse && "subnet_cse not supported with DensePeakSizeBatched (Phase 2)"); @@ -114,7 +131,6 @@ EvalSequence single_term_opt( // is_volatile_leaf gates batching; volatile_weight / roofline feed the // secondary tie-break among equal-peak schedules. PeakBatchedModel model{idxsz, - is_batchable_index, batch_target_size, is_volatile_leaf, inner_pow, @@ -125,10 +141,29 @@ EvalSequence single_term_opt( roofline.block_prefactor, batch_persistent_only, peak_flops_tolerance, - accumulation_factor}; + accumulation_factor, + cost.peak_threshold}; + model.perf_first = (Metric == ObjectiveFunction::DenseTimeSpaceBatched); model.prune_outer_products = prune_outer_products; + model.batch_spectator_indices = cost.batch_spectator_indices; + model.order_aware_recompute = cost.order_aware_recompute; + // Building blocks: the contracted-role predicate feeds the DP's contracted + // filter; the external-role predicate feeds the external filter. Each + // defaults to decline (returns false); there is no cross-role fallback, so + // a mode is batchable in a role only if that role's predicate admits it. + model.is_batchable_contracted_index = is_batchable_contracted_index; + model.is_batchable_external_index = cost.is_batchable_external_index; + // charge_batch_recompute defaults to true (see PeakBatchedModel): the batch + // recomputation cost is always reflected on the flops/exec mode. + if (out_axes) { + auto [seq, modes] = run_single_term_opt_axes(model, network, tidxs); + *out_axes = std::move(modes); + return seq; + } return run_single_term_opt(model, network, tidxs); } else if constexpr (Metric == ObjectiveFunction::DenseFLOPs) { + SEQUANT_ASSERT(!out_axes && + "out_axes only supported with DensePeakSizeBatched"); if (is_volatile_leaf && volatile_weight > 1.0) { size_t i = 0; for (auto&& t : network.tensors()) { @@ -138,7 +173,7 @@ EvalSequence single_term_opt( } nr = volatile_weight; } - (void)is_batchable_index; + (void)is_batchable_contracted_index; (void)batch_target_size; (void)batch_persistent_only; (void)peak_flops_tolerance; @@ -152,9 +187,12 @@ EvalSequence single_term_opt( return run_single_term_opt(model, network, tidxs); } else { static_assert(Metric == ObjectiveFunction::DenseSize, - "Only DenseFLOPs, DenseSize, DensePeakSize, and " - "DensePeakSizeBatched ObjectiveFunction supported."); - (void)is_batchable_index; + "Only DenseFLOPs, DenseSize, DenseSpaceTime, " + "DenseSpaceTimeBatched, DenseTimeSpace, and " + "DenseTimeSpaceBatched ObjectiveFunction supported."); + SEQUANT_ASSERT(!out_axes && + "out_axes only supported with DensePeakSizeBatched"); + (void)is_batchable_contracted_index; (void)batch_target_size; (void)batch_persistent_only; (void)peak_flops_tolerance; @@ -181,31 +219,38 @@ EvalSequence single_term_opt( /// \return Parenthesized product expression. /// /// @note @c prod is assumed to consist of only Tensor expressions -/// @note The remaining parameters (\c subnet_cse, \c cost, -/// \c is_batchable_index, \c batch_target_size, \c inner_pow, -/// \c batch_persistent_only) are forwarded verbatim to the detail -/// \ref single_term_opt overload; see it for their semantics. +/// @note The remaining parameters (\c subnet_cse, \c cost) are forwarded +/// verbatim to the detail \ref single_term_opt overload; see it for +/// their semantics. All batching config (contracted/external role +/// predicates, \c batch_target_size, \c inner_pow, \c +/// batch_persistent_only) now lives on \ref CostParams. +/// \param out_axes When non-null AND \p Metric == +/// ObjectiveFunction::DensePeakSizeBatched, filled with the per-node +/// sliced-sets of the returned Product tree's contraction nodes, in +/// the same left-first post-order the nested Product below is built +/// in (so \c (*out_axes)[j] annotates the j-th Product node formed by +/// the \c -1-handling arm of the loop below). Left empty if \p prod +/// has fewer than 3 factors (no factorization is performed) or if +/// \p Metric != DensePeakSizeBatched. /// template ExprPtr single_term_opt( Product const& prod, IdxToSz&& idxsz, bool subnet_cse = false, CostParams const& cost = {}, - std::function const& is_batchable_index = {}, - std::function batch_target_size = {}, - std::function const& inner_pow = {}, - bool batch_persistent_only = false) { + container::vector* out_axes = nullptr) { using ranges::views::filter; using ranges::views::reverse; + if (out_axes) out_axes->clear(); if (prod.factors().size() < 3) return ex(Product{prod.scalar(), prod.factors().begin(), prod.factors().end(), Product::Flatten::No}); auto const tensors = prod | filter(&ExprPtr::template is) | ranges::to_vector; - auto seq = detail::single_term_opt( - TensorNetwork{tensors}, std::forward(idxsz), subnet_cse, cost, - is_batchable_index, batch_target_size, inner_pow, batch_persistent_only); + auto seq = detail::single_term_opt(TensorNetwork{tensors}, + std::forward(idxsz), + subnet_cse, cost, out_axes); auto result = container::svector{}; for (auto i : seq) if (i == -1) { diff --git a/SeQuant/core/optimize/single_term_detail.hpp b/SeQuant/core/optimize/single_term_detail.hpp index 5e4c5f63f0..c8e25d3529 100644 --- a/SeQuant/core/optimize/single_term_detail.hpp +++ b/SeQuant/core/optimize/single_term_detail.hpp @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -93,6 +94,17 @@ double inner_aware_volume(Tot const& tot_idxs, Ixex const& ixex, mem *= inner_pow(c, k); } } else { + // No inner_pow, but this tensor HAS composite (CSV/PNO tensor-of-tensor) + // indices: sizing them by the base extent silently mis-sizes the tensor + // (each composite counted at its full base-space extent instead of its + // per-proto domain), which has repeatedly inverted factorization choices + // (e.g. picking a 4-PAO integral). An empty inner_pow is only valid for a + // network with NO composites; refuse to guess here. + if (!ranges::empty(tot_idxs.inner)) + throw std::invalid_argument( + "inner_aware_volume: composite (CSV/PNO) indices present but no " + "inner_pow provided -- sizing composites by base extent is a bug. " + "Pass a real inner_pow (e.g. SizeRegime::inner_pow_fn())."); mem = ranges::accumulate(tot_idxs.inner, mem, std::multiplies{}, ixex); } return mem; @@ -250,20 +262,68 @@ container::vector subset_footprints( /// it. The returned list assigns each index a stable bit position: index at /// position \c k is bit \c k of a sliced-set bitmask \c B. /// +/// In addition to the top-level slots, two more passes admit pure-occupied +/// indices that \p is_batchable never sees: +/// +/// - the pure-occupied protoindices of composite (CSV/PNO/OSV +/// tensor-of-tensor) legs. A composite leg carries its external occupied +/// indices ONLY as protoindices -- they never appear as a top-level +/// bra/ket/aux slot -- so the slot scan alone drops them. +/// - an explicit pure-occupied index that is open (external) on the network +/// root, i.e. a member of \c network.ext_indices(). Such an index is a +/// genuine top-level slot, but \p is_batchable is typically scoped to a +/// non-occupied space (e.g. DF/RI aux), so it would otherwise never be +/// admitted as a batching candidate. Contracted (internal) occupied +/// indices -- those that connect two or more tensors -- are NOT open on +/// the root and so are never admitted by this pass. +/// +/// Admitting either lets the batched DP slice that external-occ external +/// mode (mirrored in \ref subset_open_aux). Both passes are guarded by index +/// space (only pure-occupied indices are admitted) so PAO/aux/internal-occ +/// recognition is unchanged. +/// /// \param network The TensorNetwork to scan. /// \param is_batchable Predicate returning true for indices in a batchable /// space (e.g. a DF/RI auxiliary space). /// \return Ordered, deduplicated list of batchable indices. -inline container::vector batchable_index_list( +/// \brief Candidate batchable modes: every index (and protoindex) whose space +/// is batchable in EITHER role. +/// +/// Batchability is role-based and caller-defined, keeping this layer +/// domain-generic (no index-space kind is named here): +/// - \p is_batchable admits a space batchable when the mode is CONTRACTED +/// (summed at some node); +/// - \p is_batchable_external admits a space batchable when the mode is +/// EXTERNAL (open on the term root -- a spectator carried to the result). +/// +/// This returns the UNION of both roles. Each mode's actual role is resolved by +/// \ref PeakBatchedModel::build_context from the root open set, which then +/// drops any mode its role's predicate rejects -- e.g. a mode admitted only as +/// external but appearing contracted is not batchable, which keeps the 2^m +/// search space free of modes that can never be batched in the role they occur +/// in. Protoindices are candidates too: they become plain outer modes in the +/// array view. +inline container::vector batchable_mode_list( TensorNetwork const& network, - std::function const& is_batchable) { + std::function const& is_batchable, + std::function const& is_batchable_external = {}) { container::vector aux; - if (!is_batchable) return aux; + // This free function keeps empty-defaulted predicate params for direct + // single-arg callers, so each is null-guarded below; an all-declining (or + // all-empty) input yields an empty result naturally -- no early return. + auto match = [&](Index const& ix) { + return (is_batchable && is_batchable(ix)) || + (is_batchable_external && is_batchable_external(ix)); + }; for (auto&& t : network.tensors()) { auto tp = std::dynamic_pointer_cast(t); - for (auto&& ix : ranges::views::concat(tp->bra(), tp->ket(), tp->aux())) - if (is_batchable(ix) && ranges::find(aux, ix) == ranges::end(aux)) + for (auto&& ix : ranges::views::concat(tp->bra(), tp->ket(), tp->aux())) { + if (match(ix) && ranges::find(aux, ix) == ranges::end(aux)) aux.push_back(ix); + for (auto&& p : ix.proto_indices()) + if (match(p) && ranges::find(aux, p) == ranges::end(aux)) + aux.push_back(p); + } } return aux; } @@ -286,7 +346,7 @@ inline container::vector batchable_index_list( /// the backend rounds *down* to a tile multiple (never above the /// target). /// \param aux_list Ordered list of distinct batchable indices (as returned -/// by \ref batchable_index_list). +/// by \ref batchable_mode_list). /// \param inner_pow Optional k-aware CSV/PNO composite extent forwarded to each /// per-\c B \ref subset_footprints call; see \ref inner_aware_volume. /// Orthogonal to slicing (composites are not the batchable aux indices). @@ -299,19 +359,25 @@ container::vector> sliced_footprints( container::vector const& aux_list, std::function const& inner_pow = {}, container::vector const* connected = nullptr) { + // Retained for API compatibility; shrinkability is decided by aux_list + // membership (which spans all batchability roles), not by this predicate. + (void)is_batchable; std::size_t const m = aux_list.size(); container::vector> tables(std::size_t{1} << m); for (std::size_t B = 0; B < tables.size(); ++B) { auto extent = [&, B](Index const& ix) -> std::size_t { std::size_t e = idxsz(ix); - if (is_batchable && is_batchable(ix)) { - auto it = ranges::find(aux_list, ix); - if (it != ranges::end(aux_list)) { - std::size_t k = - static_cast(it - ranges::begin(aux_list)); - if (B & (std::size_t{1} << k)) - return std::min(e, batch_target_size(ix)); - } + // Membership in aux_list IS the authoritative "this is a batchable mode" + // test: that list spans ALL batchability roles (contracted and external). + // Gating additionally on the contracted-role predicate silently makes a + // slice of any mode admitted outside it a NO-OP -- the mode sits in the + // sliced set B yet keeps its full extent, so the DP sees no benefit and + // never batches it. + auto it = ranges::find(aux_list, ix); + if (it != ranges::end(aux_list)) { + std::size_t k = static_cast(it - ranges::begin(aux_list)); + if (B & (std::size_t{1} << k)) + return std::min(e, std::max(batch_target_size(ix), 1)); } return e; }; @@ -543,8 +609,8 @@ inline SubnetMetadata build_subnet_metadata( /// \brief Per-subset bitmask of batchable indices that are OPEN in that subset. /// -/// For each subset \c n of the input tensors, bit \c k of \c open_aux[n] is set -/// iff \c aux_list[k] is among the open (external) indices of subset \c n +/// For each subset \c n of the input tensors, bit \c k of \c open_modes[n] is +/// set iff \c aux_list[k] is among the open (external) indices of subset \c n /// (those that remain after contracting \c n's tensors, with \c tidxs as the /// final targets). Used by the multi-mode batched DP and oracle to restrict the /// sliced-set context to indices actually open in a sized subset, so that table @@ -553,26 +619,41 @@ inline SubnetMetadata build_subnet_metadata( /// \param network The TensorNetwork. /// \param tidxs Target (open) indices of the network. /// \param aux_list Ordered list of distinct batchable indices (as returned by -/// \ref batchable_index_list); index \c k maps to bit \c k. -/// \return \c open_aux[n] for every subset \c n. +/// \ref batchable_mode_list); index \c k maps to bit \c k. +/// \return \c open_modes[n] for every subset \c n. template container::vector subset_open_aux( TensorNetwork const& network, TIdxs const& tidxs, container::vector const& aux_list) { container::vector results( (std::size_t{1} << network.tensors().size())); - // NOT pruned: is_spectator_axis consumes open_aux over the full subset + // NOT pruned: is_external_mode consumes open_modes over the full subset // lattice (including disconnected subsets), so every entry must be real. init_results(network, tidxs, results); - container::vector open_aux(results.size(), 0); + // A batchable mode may be open either DIRECTLY (a top-level open index) or as + // a PROTOINDEX of an open index: protoindices become plain outer modes in the + // array view, so a mode carried as a proto of an open index is open too. Both + // are checked structurally -- no index-space kind is consulted, keeping this + // layer domain-generic. + container::vector open_modes(results.size(), 0); for (std::size_t n = 0; n < results.size(); ++n) { for (std::size_t k = 0; k < aux_list.size(); ++k) { - if (ranges::find(results[n].indices, aux_list[k]) != - ranges::end(results[n].indices)) - open_aux[n] |= (std::size_t{1} << k); + Index const& ax = aux_list[k]; + bool open = ranges::find(results[n].indices, ax) != + ranges::end(results[n].indices); + if (!open) { + for (auto const& ix : results[n].indices) { + auto const& pr = ix.proto_indices(); + if (ranges::find(pr, ax) != ranges::end(pr)) { + open = true; + break; + } + } + } + if (open) open_modes[n] |= (std::size_t{1} << k); } } - return open_aux; + return open_modes; } /// Per-(subset, sliced-set) state for the multi-mode batched peak DP. From f130582f64c96397ff6130af19c41e3b54a96060 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Wed, 29 Jul 2026 12:55:29 -0400 Subject: [PATCH 003/213] eval: multimode batched evaluator External-mode scatter + contracted accumulate; cache scope chain with fall-through; slice-on-use; per-level placement driven by a per-canonical lifetime mask (cross-occurrence meet) unioned with contracted residency; iterative (stack-safe) tree traversal. --- CMakeLists.txt | 1 + .../eval/backends/tiledarray/eval_context.hpp | 3 + .../core/eval/backends/tiledarray/result.hpp | 209 ++++- SeQuant/core/eval/cache_manager.hpp | 204 ++++- SeQuant/core/eval/eval.hpp | 850 +++++++++++++++--- SeQuant/core/eval/eval_expr.cpp | 92 +- SeQuant/core/eval/eval_expr.hpp | 147 ++- SeQuant/core/eval/fwd.hpp | 12 + SeQuant/core/eval/lifetime_mask.hpp | 138 +++ SeQuant/core/eval/node_batch_annotation.hpp | 61 ++ SeQuant/core/eval/result.hpp | 48 + 11 files changed, 1544 insertions(+), 221 deletions(-) create mode 100644 SeQuant/core/eval/lifetime_mask.hpp create mode 100644 SeQuant/core/eval/node_batch_annotation.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 29643264fc..0e67235020 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -438,6 +438,7 @@ set(SeQuant_eval_src SeQuant/core/eval/eval_expr.hpp SeQuant/core/eval/eval_node.hpp SeQuant/core/eval/eval_node_compare.hpp + SeQuant/core/eval/node_batch_annotation.hpp SeQuant/core/eval/result.cpp SeQuant/core/eval/result.hpp SeQuant/core/eval/fwd.hpp diff --git a/SeQuant/core/eval/backends/tiledarray/eval_context.hpp b/SeQuant/core/eval/backends/tiledarray/eval_context.hpp index fd5310a978..ef9b1440a9 100644 --- a/SeQuant/core/eval/backends/tiledarray/eval_context.hpp +++ b/SeQuant/core/eval/backends/tiledarray/eval_context.hpp @@ -6,10 +6,13 @@ #include #include #include +#include #include +#include #include +#include #include #include diff --git a/SeQuant/core/eval/backends/tiledarray/result.hpp b/SeQuant/core/eval/backends/tiledarray/result.hpp index 6c640cf1cc..de62af4b1f 100644 --- a/SeQuant/core/eval/backends/tiledarray/result.hpp +++ b/SeQuant/core/eval/backends/tiledarray/result.hpp @@ -329,6 +329,27 @@ template return TA::TiledRange(dims.begin(), dims.end()); } +/// Map a contiguous element range `[elem_lo, elem_hi)` on a mode's TiledRange1 +/// to the tile range `[tile_lo, tile_hi)` it must coincide with. A tiled +/// backend can only cut or scatter whole tiles, so the element bounds must be +/// in-range and fall on tile boundaries; this asserts both (mode_batches() +/// yields exactly such tile-aligned ranges). Shared by slice_mode() (GATHER a +/// block out) and write_into_slice() (SCATTER a block in) so both agree on the +/// element-to-tile contract and its alignment preconditions. +[[nodiscard]] inline std::pair slice_bounds_to_tiles( + TA::TiledRange1 const& tr1, std::size_t elem_lo, std::size_t elem_hi) { + SEQUANT_ASSERT(elem_lo >= tr1.elements_range().first && elem_lo < elem_hi && + elem_hi <= tr1.elements_range().second); + std::size_t const tile_lo = tr1.element_to_tile(elem_lo); + SEQUANT_ASSERT(tr1.tile(tile_lo).first == elem_lo); // lo on a tile boundary + std::size_t const tile_hi = (elem_hi >= tr1.elements_range().second) + ? tr1.tile_extent() + : tr1.element_to_tile(elem_hi); + SEQUANT_ASSERT(elem_hi >= tr1.elements_range().second || + tr1.tile(tile_hi).first == elem_hi); // hi on a tile boundary + return {tile_lo, tile_hi}; +} + } // namespace detail /// TA::Tensor memory use logger @@ -353,6 +374,14 @@ template TA::DistArray const& arr, std::size_t mode, std::size_t tile_lo, std::size_t tile_hi); +// defined below; declared here so the result classes' write_into_slice() +// overrides can call it. The scatter inverse of slice_array_over_mode(). +template +void write_array_into_mode(TA::DistArray& dest, + TA::DistArray const& block, + std::size_t mode, std::size_t tile_lo, + std::size_t tile_hi); + /// Partition a TiledRange1 into contiguous, tile-aligned element-range batches, /// each covering at most \p target_batch_size elements: whole tiles are /// appended to a batch until the next tile would push it over the target, so \p @@ -434,22 +463,8 @@ class ResultTensorTA final : public Result { [[nodiscard]] ResultPtr slice_mode(std::size_t mode, std::size_t elem_lo, std::size_t elem_hi) const override { - auto const& tr1 = get().trange().dim(mode); - // slice_mode takes element bounds, but a tiled backend can only cut on tile - // boundaries; mode_batches() returns exactly such (tile-aligned, in-range) - // bounds. Assert the precondition so misuse is caught rather than silently - // producing an over- or under-sized slice (which would break batched sums). - SEQUANT_ASSERT(elem_lo >= tr1.elements_range().first && elem_lo < elem_hi && - elem_hi <= tr1.elements_range().second); - std::size_t const tile_lo = tr1.element_to_tile(elem_lo); - SEQUANT_ASSERT(tr1.tile(tile_lo).first == - elem_lo); // lo on a tile boundary - std::size_t const tile_hi = (elem_hi >= tr1.elements_range().second) - ? tr1.tile_extent() - : tr1.element_to_tile(elem_hi); - SEQUANT_ASSERT(elem_hi >= tr1.elements_range().second || - tr1.tile(tile_hi).first == - elem_hi); // hi on a tile boundary + auto const [tile_lo, tile_hi] = detail::slice_bounds_to_tiles( + get().trange().dim(mode), elem_lo, elem_hi); return eval_result( slice_array_over_mode(get(), mode, tile_lo, tile_hi)); } @@ -460,6 +475,38 @@ class ResultTensorTA final : public Result { target_batch_size); } + void write_into_slice(Result const& block, std::size_t mode, + std::size_t block_lo, std::size_t block_hi) override { + SEQUANT_ASSERT(block.is()); + auto& dest = get(); + auto const [tile_lo, tile_hi] = detail::slice_bounds_to_tiles( + dest.trange().dim(mode), block_lo, block_hi); + write_array_into_mode(dest, block.get(), mode, tile_lo, tile_hi); + } + + [[nodiscard]] ResultPtr pre_sized_zeros_over_mode( + std::size_t mode, Result const& axis_src, + std::size_t axis_src_mode) const override { + SEQUANT_ASSERT(axis_src.is()); + auto const& self = get(); + auto const& src = axis_src.get(); + auto const rank = self.trange().rank(); + SEQUANT_ASSERT(mode < rank); + SEQUANT_ASSERT(axis_src_mode < src.trange().rank()); + // Take *this's outer trange but swap in the external axis's FULL tiling + // (from axis_src's mode axis_src_mode). Every other mode of a block partial + // is already full extent, so only the sliced axis needs widening. + std::vector dims; + dims.reserve(rank); + for (std::size_t d = 0; d < rank; ++d) dims.push_back(self.trange().dim(d)); + dims[mode] = src.trange().dim(axis_src_mode); + ArrayT dest(self.world(), TA::TiledRange(dims.begin(), dims.end())); + dest.fill_local(numeric_type(0)); + dest.world().gop.fence(); + log_ta_tensor_host_memory_use(); + return eval_result(std::move(dest)); + } + [[nodiscard]] ResultPtr prod(Result const& other, std::array const& annot, DeNest DeNestFlag) const override { @@ -639,22 +686,8 @@ class ResultTensorOfTensorTA final : public Result { [[nodiscard]] ResultPtr slice_mode(std::size_t mode, std::size_t elem_lo, std::size_t elem_hi) const override { - auto const& tr1 = get().trange().dim(mode); - // slice_mode takes element bounds, but a tiled backend can only cut on tile - // boundaries; mode_batches() returns exactly such (tile-aligned, in-range) - // bounds. Assert the precondition so misuse is caught rather than silently - // producing an over- or under-sized slice (which would break batched sums). - SEQUANT_ASSERT(elem_lo >= tr1.elements_range().first && elem_lo < elem_hi && - elem_hi <= tr1.elements_range().second); - std::size_t const tile_lo = tr1.element_to_tile(elem_lo); - SEQUANT_ASSERT(tr1.tile(tile_lo).first == - elem_lo); // lo on a tile boundary - std::size_t const tile_hi = (elem_hi >= tr1.elements_range().second) - ? tr1.tile_extent() - : tr1.element_to_tile(elem_hi); - SEQUANT_ASSERT(elem_hi >= tr1.elements_range().second || - tr1.tile(tile_hi).first == - elem_hi); // hi on a tile boundary + auto const [tile_lo, tile_hi] = detail::slice_bounds_to_tiles( + get().trange().dim(mode), elem_lo, elem_hi); return eval_result( slice_array_over_mode(get(), mode, tile_lo, tile_hi)); } @@ -665,6 +698,57 @@ class ResultTensorOfTensorTA final : public Result { target_batch_size); } + void write_into_slice(Result const& block, std::size_t mode, + std::size_t block_lo, std::size_t block_hi) override { + SEQUANT_ASSERT(block.is()); + auto& dest = get(); + auto const [tile_lo, tile_hi] = detail::slice_bounds_to_tiles( + dest.trange().dim(mode), block_lo, block_hi); + write_array_into_mode(dest, block.get(), mode, tile_lo, tile_hi); + } + + [[nodiscard]] ResultPtr pre_sized_zeros_over_mode( + std::size_t mode, Result const& axis_src, + std::size_t axis_src_mode) const override { + auto const& self = get(); + auto const rank = self.trange().rank(); + SEQUANT_ASSERT(mode < rank); + // The axis-carrying leaf supplying K's FULL tiling for mode `mode` may be + // nested (this_type) or flat (that_type, e.g. an integral over the external + // occ index): read the widened axis TiledRange1 from whichever kind. Only + // this one OUTER TiledRange1 is needed; every other mode of a block partial + // is already at full extent, so *this's own outer tiling supplies them. + TA::TiledRange1 const axis_dim = [&]() -> TA::TiledRange1 { + if (axis_src.is()) { + auto const& src = axis_src.get(); + SEQUANT_ASSERT(axis_src_mode < src.trange().rank()); + return src.trange().dim(axis_src_mode); + } + SEQUANT_ASSERT(axis_src.is()); + auto const& src = axis_src.get(); + SEQUANT_ASSERT(axis_src_mode < src.trange().rank()); + return src.trange().dim(axis_src_mode); + }(); + std::vector dims; + dims.reserve(rank); + for (std::size_t d = 0; d < rank; ++d) dims.push_back(self.trange().dim(d)); + dims[mode] = axis_dim; + // A zero ToT is represented with empty inner tiles (tot_inner_rank() == 0): + // build the widened OUTER trange, then give every local outer tile a + // well-formed (empty-inner) outer tile over its range -- exactly the zero + // ToT that slice_array_over_mode() emits, and a valid destination that the + // ToT write_array_into_mode() block-assignment overwrites per scatter. The + // batches tile the widened `mode` axis with no gaps, so every outer tile is + // subsequently overwritten by some block's real inner tensors. + using value_type = typename ArrayT::value_type; + ArrayT dest(self.world(), TA::TiledRange(dims.begin(), dims.end())); + for (auto it = dest.begin(); it != dest.end(); ++it) + if (dest.is_local(it.index())) *it = value_type{it.make_range()}; + dest.world().gop.fence(); + log_ta_tensor_host_memory_use(); + return eval_result(std::move(dest)); + } + [[nodiscard]] ResultPtr prod(Result const& other, std::array const& annot, DeNest DeNestFlag) const override { @@ -889,6 +973,67 @@ template return out; } +/// \brief Scatter a per-block DistArray into a contiguous tile range of one +/// mode of a pre-sized destination -- the inverse of +/// slice_array_over_mode(). +/// +/// Writes \p block into tiles `[tile_lo, tile_hi)` of \p dest's mode \p mode, +/// leaving every other tile of \p dest untouched. \p dest must already be +/// allocated over its full TiledRange (the caller sizes the whole shape), and +/// \p block's TiledRange must equal \p dest's sub-block over `[tile_lo, +/// tile_hi)` (as produced by slice_array_over_mode() for the same mode/range). +/// Implemented with TA's block() on the assignment LHS, so only the addressed +/// sub-block is written and block-sparse shape is preserved. Every mode's +/// element lobound is preserved (via TA's `preserve_lobound`), exactly as the +/// lobound-preserving GATHER in slice_array_over_mode(): the destination +/// sub-block and the source share element coordinates, so a spectator index +/// carrying a nonzero lobound (e.g. a frozen-core offset) lands at its true +/// offset rather than being rebased to 0. Reconstructs a whole result from a +/// disjoint, gap-free tiling of one mode: scattering each block of a partition +/// reproduces the array `slice_array_over_mode()` would gather back out. +template +void write_array_into_mode(TA::DistArray& dest, + TA::DistArray const& block, + std::size_t mode, std::size_t tile_lo, + std::size_t tile_hi) { + using ranges::views::iota; + auto const rank = dest.trange().rank(); + SEQUANT_ASSERT(mode < rank); + SEQUANT_ASSERT(tile_lo < tile_hi && + tile_hi <= dest.trange().dim(mode).tile_extent()); + container::svector lo(rank, 0), hi(rank); + for (std::size_t d = 0; d < rank; ++d) + hi[d] = dest.trange().dim(d).tile_extent(); + lo[mode] = tile_lo; + hi[mode] = tile_hi; + // For a tensor-of-tensor array the annotation must label an inner block + // ("outer;inner"); a flat annotation trips DistArray's is_tot_index() check. + // The block() is over outer modes only, so both sides share one annotation. + using value_type = typename TA::DistArray::value_type; + std::string annot; + if constexpr (TA::detail::is_tensor_of_tensor_v) { + auto const inner_rank = detail::tot_inner_rank(block); + if (inner_rank == 0) { + // block has all-empty inner tiles (tot_inner_rank() == 0): it represents + // zero and there is no inner rank to form the ToT annotation block() + // needs. A zero contribution leaves the pre-sized destination slice as + // it was, so skip the scatter entirely -- mirroring the zero-ToT early + // return in slice_array_over_mode(). + return; + } + annot = TA::detail::dummy_annotation(static_cast(rank), + static_cast(inner_rank)); + } else { + annot = detail::ords_to_annot(iota(std::size_t{0}, rank)); + } + // preserve_lobound: address the destination sub-block in its original element + // coordinates (keeping every mode's lobound) so it matches the source block, + // which slice_array_over_mode() also gathered with preserve_lobound. Plain + // block() would rebase the sub-block to 0 and mismatch the source trange. + dest(annot).block(lo, hi, TA::preserve_lobound) = block(annot); + TA::DistArray::wait_for_lazy_cleanup(dest.world()); +} + /// \brief Compute the result's OUTER TiledRange for a binary product from the /// type-erased operands and the [left, right, result] annotations. /// diff --git a/SeQuant/core/eval/cache_manager.hpp b/SeQuant/core/eval/cache_manager.hpp index 92fcdb8052..2899a1710d 100644 --- a/SeQuant/core/eval/cache_manager.hpp +++ b/SeQuant/core/eval/cache_manager.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -18,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -70,6 +72,25 @@ class CacheManager { std::any const& node, Result const& left, Result const& right, std::array const& annot)>; + /// The batch context: an ordered stack (outermost-first) of the enclosing + /// realized batch loops, one entry per loop, `{axis K, {block_lo, block_hi}}` + /// (element range). Set on the per-block scratch by the batched evaluator + /// before it re-enters evaluate(); read by the Enter-stage slice-on-use so a + /// cached intermediate fetched from an ancestor scope is sliced to the modes + /// of the loops the fetch crossed (see eval.hpp). Empty (default) => no + /// enclosing batch loop, so slice-on-use is inert and behavior is + /// byte-identical to the pre-slice-on-use path. + using BatchContext = + container::svector>>; + + /// Result of access_at(): the fetched pointer plus the hop distance (number + /// of parent links crossed) to the scope that held it. hops == 0 means a + /// local hit; a null ptr carries hops == 0. + struct AccessResult { + ResultPtr ptr; + std::size_t hops; + }; + private: using hasher_type = TreeNodeHasher; using comparator_type = TreeNodeEqualityComparator; @@ -157,6 +178,14 @@ class CacheManager { std::unordered_map cache_map_; + /// Parent cache for the scope chain (loop-nest visibility). A batch scratch + /// sets this to the cache one level up; access() delegates on a local miss + /// so a loop-invariant node stored once at an ancestor level is found by + /// every inner body without copy-down. Null (default) => standalone cache, + /// byte-identical to pre-scope-chain behavior. Non-owning; the parent must + /// outlive this cache. + CacheManager* parent_ = nullptr; + /// Running high-water mark (bytes) of the eval engine's live working set, /// updated by note_working_set() and cleared by reset(). Held here rather /// than in the recursive evaluate() so it persists across the whole @@ -169,6 +198,12 @@ class CacheManager { shaped_product_hook_type shaped_product_hook_{}; + /// Enclosing realized batch loops for slice-on-use (see BatchContext). Empty + /// (default) => no enclosing batch loop; the batched evaluator sets it on the + /// per-block scratch before each re-entry. Not cleared by reset() (it is + /// per-loop-iteration structural, re-set each block by the evaluator). + BatchContext batch_context_{}; + public: /// Sets the custom evaluator (see custom_evaluator_type). Pass an empty /// std::function to clear it. @@ -193,6 +228,41 @@ class CacheManager { return shaped_product_hook_; } + /// Sets the batch context (see batch_context_). Pass an empty context to + /// clear it. + void set_batch_context(BatchContext c) noexcept { + batch_context_ = std::move(c); + } + + /// \return the batch context (empty if none is set). + [[nodiscard]] BatchContext const& batch_context() const noexcept { + return batch_context_; + } + + /// Sets the scope-chain parent (see parent_). Pass nullptr to detach. + void set_parent(CacheManager* p) noexcept { parent_ = p; } + + /// \return the scope-chain parent (see parent_), or nullptr if this is a + /// standalone / chain-root cache. Used by the batched evaluator to + /// walk up to a target ancestor level when hoisting an invariant. + [[nodiscard]] CacheManager* parent() const noexcept { return parent_; } + + /// Ensure a scope-hoist slot exists for @p key so a loop-invariant + /// intermediate can be stored here (store() is a no-op for an unregistered + /// key). The slot is NON-persistent with an effectively unbounded life, so it + /// is never drained by access() and lives until the next reset() -- per-batch + /// for a batch scratch (rebuilt for the next batch of the loop it is scoped + /// to), per-term for the real cache (rebuilt for the next term). Idempotent: + /// an existing entry (with any stored data) is left untouched. The unbounded + /// life -- rather than the emitted effective_count -- is deliberate: a + /// whole-nest invariant's escaped-outer set is empty, so its emitted + /// effective_count is 1, which as a life would drain the entry on first use; + /// reset() is the correct lifetime boundary for a hoisted invariant. + void ensure_hoist_slot(key_type const& key) { + cache_map_.try_emplace( + key, entry{std::numeric_limits::max(), /*persistent=*/false}); + } + /// Default persistence classifier: every entry is non-persistent (NP). struct all_non_persistent { bool operator()(key_type const&) const noexcept { return false; } @@ -239,13 +309,30 @@ class CacheManager { /// @brief Access cached data. /// /// @param key The key that identifies the cached data. - /// @return ResultPtr to Result - ResultPtr access(key_type const& key) noexcept { + /// @return the fetched pointer plus the hop distance (number of parent links + /// crossed) to the scope that held it; {nullptr, 0} on a total miss. + /// + /// A local entry only "hits" if it is currently holding data; a key + /// registered locally but never (yet) stored here -- e.g. a hoisted + /// loop-invariant node whose value lives only at an ancestor level -- is a + /// local miss just like an unregistered key, and must fall through the + /// same way. Standalone (parent_ == nullptr) behavior is unchanged: a total + /// miss returns {nullptr, 0}. The hop distance surfaces the value's lifetime + /// scope so the caller (Enter-stage slice-on-use) can slice it to exactly the + /// batch loops the fetch crossed. + [[nodiscard]] AccessResult access_at(key_type const& key) noexcept { if (auto found = cache_map_.find(key); found != cache_map_.end()) - return found->second.access(); - return nullptr; + if (auto data = found->second.access(); data) return {data, 0}; + if (!parent_) return {nullptr, 0}; + auto up = parent_->access_at(key); + return {up.ptr, up.hops + 1}; // count the link we just crossed } + /// @param key The key that identifies the cached data. + /// @return ResultPtr to Result. Thin forwarder to access_at() that drops the + /// hop distance, for the non-batched callers that do not slice. + ResultPtr access(key_type const& key) noexcept { return access_at(key).ptr; } + /// /// @param key The key to identify the cached data. /// @param data The data to be cached. @@ -437,7 +524,7 @@ struct zero_footprint { }; /// Default batchability predicate for cache_manager: no index is batchable, so -/// the free-batchable-axis caching veto is inert (preserves the pre-batch +/// the free-batchable-mode caching veto is inert (preserves the pre-batch /// behavior for callers that do not pass a predicate). struct never_batchable { bool operator()(auto const&) const noexcept { return false; } @@ -458,18 +545,37 @@ struct never_batchable { /// of huge intermediates that carry a free large-space index (e.g. a /// half-transformed DF integral with a free projected-AO index), at the /// cost of recomputation. 0 (default) disables the gate. -/// \param is_batchable_index `bool(Index const&)`: an index the runtime batched -/// evaluator slices over (typically the DF/RI auxiliary). A node whose -/// *result* (canonical) indices contain such an index carries a -/// batchable axis FREE: the evaluator slices it per batch and the -/// single-term optimizer prices it sliced, so caching it whole would -/// hold an intermediate both other components mean to slice. Such nodes -/// are NOT cached (neither NP repeat nor P frontier) -- recomputed -/// (sliced under each consumer's batch trigger) instead of materialized -/// whole and held. This is the structural counterpart of \p -/// max_footprint: the batch axis, not a byte threshold, identifies the -/// free-large-index intermediates. The default never_batchable accepts -/// nothing, leaving the veto inert. +/// \param is_batchable_contracted_index `bool(Index const&)`: an index sliced +/// in the CONTRACTED role (typically the DF/RI auxiliary). This is the +/// contracted-role building block, NEVER the derived role union: the +/// veto is contracted-stamp-only. A node whose +/// own \c batched_here() carries such an index tagged \c +/// BatchModeType::Contracted -- i.e. a mode actually sliced AT this node +/// -- FREE in its *result* (canonical) indices is, by construction, a +/// free-large-index intermediate the evaluator builds one batch-slice +/// at a time and the single-term optimizer prices sliced. Caching it +/// whole would hold an intermediate the runtime means to slice. Such +/// nodes are NOT cached (neither NP repeat nor P frontier) -- +/// recomputed (sliced under each consumer's batch trigger) instead of +/// materialized whole and held. A node whose \c batched_here() carries +/// only \c BatchModeType::External entries (an external index the node +/// is merely invariant under, not one sliced at this node -- e.g. a +/// loop-invariant intermediate like \c gC) is NOT vetoed: it is +/// genuinely batch-invariant and stays cacheable. This is the +/// structural counterpart of \p max_footprint: the sliced batch mode, +/// not a byte threshold, identifies the free-large-index +/// intermediates. The default never_batchable accepts nothing, leaving +/// the veto inert. Independently of \p is_batchable_contracted_index, a +/// node whose cross-occurrence lifetime mask is non-empty (\c +/// !EvalExpr::mask_all_full(); this builder itself calls \c +/// stamp_lifetime_masks over \p nodes before the DAG walk below, so +/// the mask is always current here regardless of caller -- see \c +/// lifetime_mask.hpp) is likewise batch-variant -- some enclosing +/// External batch mode slices it in every occurrence, so its value +/// differs per batch of that mode -- and is refused run-scope +/// residence even with an empty \c batched_here(); only an all-full +/// node (empty mask, including every node on the OFF path) is +/// admitted. /// \see CacheManager, cache_manager template > const& n) { @@ -485,6 +591,17 @@ auto cache_manager(meta::eval_node_range auto const& nodes, auto&& is_volatile, { footprint_of(n) } -> std::convertible_to; } { + // Stamp the cross-occurrence lifetime mask on this SAME forest before the + // DAG walk / veto below reads it (part (b) of the batch-variant veto reads + // EvalExpr::mask_all_full()). Doing this here -- rather than leaving it to + // each caller -- makes "mask is current for the veto" an invariant of this + // builder instead of a per-caller obligation: every caller of this overload + // (SeQuant's build_dryrun_cache, mpqc's build_cache_manager) is covered + // uniformly. Unconditional and idempotent; a no-op when the forest carries + // no External batched_here() stamps (every mask stays empty/all-full), so + // this never changes behavior on the OFF path. + sequant::stamp_lifetime_masks(nodes); + using TreeNode = std::ranges::range_value_t>; using Hasher = TreeNodeHasher; @@ -523,23 +640,50 @@ auto cache_manager(meta::eval_node_range auto const& nodes, auto&& is_volatile, // Footprint gate: a node whose result is larger than max_footprint is never // cached (so it is recomputed by each consumer rather than materialized whole // and held), bounding the footprint of huge free-large-index intermediates. - // Free-batchable-axis veto: a node whose result carries an index the runtime - // slices over (is_batchable_index) is, by construction, a free-large-index - // intermediate the evaluator builds one batch-slice at a time and the - // optimizer prices sliced. Caching it -- as an NP repeat or an NV/V-frontier - // P node -- would materialize and hold it whole, contradicting both. Veto its - // caching (the structural form of the max_footprint gate) so each consumer - // recomputes it sliced under its own batch trigger. + // Batch-variant veto ("a batched node cannot be run-scope"): this builder + // populates the outermost / persistent (run-scope) cache, so it must refuse + // any node that is batch-VARIANT -- one whose cached value would depend on + // which batch is live. Such a node is refused for two reasons: caching it + // whole contradicts the runtime slicing it (and the optimizer pricing it + // sliced), and -- the F1 safety invariant -- a child batch scratch that + // misses locally falls through to this cache for ANY key, so a batch-variant + // final left here could be served full (wrong-batch) to an inner body. A node + // is batch-variant iff either: + // (a) its own batched_here() carries a mode actually sliced AT this node + // (BatchModeType::Contracted, is_batchable_contracted_index) FREE in + // its result -- + // a free-large-index intermediate the evaluator builds one slice at a + // time; or + // (b) its cross-occurrence lifetime mask is non-empty (\c + // !n->mask_all_full(), \c lifetime_mask.hpp) -- some External batch + // mode (of this node or an enclosing ancestor, over ALL its + // occurrences under the canonical meet) slices it, so its value + // differs per batch of that mode even if it slices nothing itself. + // A node that is invariant to every batched mode is NOT vetoed and stays + // cacheable at run scope -- this is where a hoisted loop-invariant + // intermediate (all-full mask; or an External-only / no batched_here entry, + // e.g. gC) lands. OFF path (no order-aware annotations, hence no \c + // stamp_lifetime_masks External stamps): every mask is empty (all-full, + // \c EvalExpr::sliced_modes_ default-constructed), so neither disjunct + // fires and the veto admits exactly what it did before -- byte-identical. std::unordered_map filtered; for (auto&& [n, c] : counts) { if (!(c >= min_repeats || persistent.contains(n))) continue; - bool free_batchable_axis = false; - for (auto const& ix : n->canon_indices()) - if (is_batchable_index(ix)) { - free_batchable_axis = true; + auto const& canon_ix = n->canon_indices(); + bool sliced_batch_axis = false; + for (auto const& [ix, kind] : n->batched_here()) + if (kind == BatchModeType::Contracted && + is_batchable_contracted_index(ix) && + std::find(canon_ix.begin(), canon_ix.end(), ix) != canon_ix.end()) { + sliced_batch_axis = true; break; } - if (free_batchable_axis || + // (b): a node whose cross-occurrence mask is non-empty is sliced by some + // enclosing external mode in every occurrence => batch-variant => refused + // run-scope residence. all-full (empty mask; incl. the OFF path) is + // admitted. + bool const batch_variant = sliced_batch_axis || !n->mask_all_full(); + if (batch_variant || (max_footprint > 0. && footprint_of(n) > max_footprint)) { persistent.erase(n); // keep is_persistent consistent with what is cached continue; diff --git a/SeQuant/core/eval/eval.hpp b/SeQuant/core/eval/eval.hpp index 96e3ffbea9..3761ef53ea 100644 --- a/SeQuant/core/eval/eval.hpp +++ b/SeQuant/core/eval/eval.hpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -19,6 +20,7 @@ #include #include +#include #include #include #include @@ -199,7 +201,7 @@ enum struct TermMode { Begin, End }; /// One log record per eval op. Line format: /// // clang-format off -/// Eval | |