mbpt: closed-shell Kramers (time-reversal) CC tracer (round 2)

.pre-commit-config.yaml

@@ -14,8 +14,9 @@ repos:
         hooks:
             -   id: remove-crlf
             -   id: forbid-tabs
-                # MPQC output and JS/XML/CSS/CMake can contain tabs
-                exclude: \.(out|cmake|js|xml|css)$
+                # MPQC output, JS/XML/CSS/CMake, and generated ITF code
+                # (tab-separated by format) can contain tabs
+                exclude: \.(out|cmake|js|xml|css)$|\.itfaa(\.expected)?$
     # see https://github.com/Lucas-C/pre-commit-hooks#forbid--remove-some-unicode-characters
     -   repo: local
         hooks:

CMakeLists.txt

@@ -424,6 +424,8 @@ set(SeQuant_mbpt_src
         SeQuant/domain/mbpt/space_qns.hpp
         SeQuant/domain/mbpt/spin.cpp
         SeQuant/domain/mbpt/spin.hpp
+        SeQuant/domain/mbpt/spinor.cpp
+        SeQuant/domain/mbpt/spinor.hpp
         SeQuant/domain/mbpt/vac_av.hpp
         SeQuant/domain/mbpt/vac_av.cpp
         SeQuant/domain/mbpt/utils.hpp
@@ -439,6 +441,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

SeQuant/core/batch_policy.hpp

@@ -5,6 +5,7 @@
 
 #include <cstddef>
 #include <functional>
+#include <limits>
 
 namespace sequant {
 
@@ -14,8 +15,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 {
-  SEQUANT_DESIGNATED_INIT_ONLY;
-  std::function<bool(Index const&)> 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<bool(Index const&)> 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<bool(Index const&)> 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<bool(Index const&)> 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
@@ -24,6 +56,22 @@ struct BatchPolicy {
   std::function<std::size_t(Index const&)> batch_target_size = {};
   std::function<bool(Tensor const&)> 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
@@ -43,6 +91,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<double>::infinity();
 };
 
 }  // namespace sequant

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 <SeQuant/core/container.hpp>
+#include <SeQuant/core/eval/backends/dryrun/size_regime.hpp>
+#include <SeQuant/core/index.hpp>
+#include <SeQuant/core/optimize/cost_model.hpp>
+#include <SeQuant/core/optimize/options.hpp>
+#include <SeQuant/core/optimize/single_term_detail.hpp>
+
+#include <cstddef>
+#include <functional>
+#include <utility>
+
+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<Index, std::size_t>;
+
+///
+/// \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<Index> 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<Index>{}, container::svector<Index>{});
+    return static_cast<std::size_t>(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<Index> const& out,
+                             container::svector<Index> 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<Index>{});
+  }
+
+  ///
+  /// \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<double>(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<std::size_t(Index const&)> 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