diff --git a/SeQuant/core/batch_policy.hpp b/SeQuant/core/batch_policy.hpp index 48c3c5cd45..c5dfe26e16 100644 --- a/SeQuant/core/batch_policy.hpp +++ b/SeQuant/core/batch_policy.hpp @@ -3,6 +3,7 @@ #include #include +#include namespace sequant { @@ -21,6 +22,12 @@ 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). + bool batch_spectator_indices = 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 +47,13 @@ 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 objective. 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. + /// This is the *enable* trigger for batching (a finite value turns it on). + double peak_threshold = std::numeric_limits::infinity(); }; } // namespace sequant 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..a19dbf4112 --- /dev/null +++ b/SeQuant/core/eval/backends/dryrun/cost_profile.hpp @@ -0,0 +1,315 @@ +#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 + +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_index before building the cache, so + /// the cache veto and the replay evaluator can never batch on different + /// predicates. Only \c build_dryrun_cache() called directly honors it 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; }); + + 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_index is COPIED over \c cfg.is_batchable_index +/// internally so the cache veto and the evaluator batch on the same +/// axis predicate (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 -------------------------- + // Force the cache veto and the replay evaluator to slice on the SAME axis + // predicate: a mismatch would let the cache retain a giant the evaluator + // batches (or vice versa), silently diverging the modeled peak from the run. + CacheConfig local_cfg = cfg; + local_cfg.is_batchable_index = policy.is_batchable_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..f48a3f50cc --- /dev/null +++ b/SeQuant/core/eval/backends/dryrun/result.hpp @@ -0,0 +1,512 @@ +#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 + +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)); + 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; + } + + [[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); + } + + 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); + } + + 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 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/eval.hpp b/SeQuant/core/eval/eval.hpp index 96e3ffbea9..8e865be9cb 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 | |