Template Struct PeakBatchedModel

• Defined in File cost_model.hpp

Nested Relationships

Nested Types

• Struct PeakBatchedModel::BFrontPoint

• Struct PeakBatchedModel::Context

Struct Documentation

template<typename IdxToSz>
struct PeakBatchedModel

Multi-mode batched peak-memory single-term cost model (DenseSpaceTimeBatched objective).

Implements the per-batchable-index all-co-resident pebble-game DP, factored into the CostModel hooks driven by run_single_term_opt. Follows the batch-aware cost model design (section 6.2, model A): each DP cell is indexed by both a subset n and a sliced-set context B over the batchable indices, so a model State is the [B]-vector of per-context BatchedRes. No CSE; persistence-gated batching.

Template Parameters:

IdxToSz – A callable mapping an Index to its extent.

Public Types

using State = container::vector<container::vector<BFrontPoint>>

Per-subset DP cell: a [B]-vector (size nB = 2^m) of Pareto frontiers, one non-dominated (peak, flops) set per sliced-set context B.

Public Functions

inline bool is_batchable(Index const &ix) const

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.

template<typename TIdxs>
inline Context build_context(TensorNetwork const &network, TIdxs const &tidxs) const

inline State leaf(Context const &ctx, size_t n) const

inline State init(Context const &ctx, size_t) const

inline void relax(Context &ctx, size_t n, size_t lp, size_t rp, State const &lp_st, State const &rp_st, State &acc) const

inline void finalize(Context&, size_t, container::vector<State>&) const

inline bool is_external_mode(Context const &ctx, std::size_t k) const

True iff batchable mode bit 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 (contracted_here = (open_modes[lp]|open_modes[rp]) & ~open_modes[n]) and slicing it is purely a footprint change with identical work. This is what Context::external_mask is built from: such a mode is a nestable cell mode only under batch_spectator_indices, and only then may the DP open its loop at a node (BFrontPoint::eopen).

inline int select_root(Context const &ctx, container::vector<State> const &st) const

Threshold-gated root-frontier selection shared by reconstruct and 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 st[root][0] &#8212; the term root is always read at cell 0 (the empty sequence): every loop the schedule uses, contracted or external, is opened inside the tree by the DP, so the root has no enclosing nest.

inline EvalSequence reconstruct(Context const &ctx, container::vector<State> const &st) const

inline double subtree_peak(Context const &ctx, container::vector<State> const &st, std::size_t n, std::size_t Bsched, std::size_t Usize, int idx) const

Modeled peak (in elements) of the chosen back-pointer subtree rooted at subset n, read at DP schedule cell Bsched and sized under the explicit union mask Usize. This is the reusable, node-level re-price of the peak reconstruction in reconstructed_batched_peak: it follows the same back-pointer walk (children/order/aprime chosen by the DP) and the same staged-peak algebra (stage_first / stage_second / stage_form with the resident-scan res and accumulation contrib terms), but sizes every subset from Usize instead of from the schedule cell’s union. With Usize == cell_union(Bsched) the two agree exactly.

inline std::pair<EvalSequence, container::vector<NodeBatchAnnotation>> reconstruct_batched_modes(Context const &ctx, container::vector<State> const &st, double *out_root_peak_bytes = nullptr) const

Companion to reconstruct that additionally reports, for each -1 (contraction) entry emitted in the returned EvalSequence in emission order, the vector of Index sliced at that node (ctx.batchable_modes[bit] for each set bit of that node’s aprime). Leaf entries contribute nothing. Does not change reconstruct’s own output; the two walks are kept in lock-step so the RPN order and the per-node modes line up.

out_root_peak_bytes (when non-null) receives the term’s reported root peak in bytes: the selected root frontier point’s footprint, which already reflects every loop &#8212; contracted and external alike &#8212; the DP chose to open, since external opens are priced inside the DP (BFrontPoint::eopen) rather than stamped on afterwards.

Public Members

IdxToSz idxsz

std::function<std::size_t(Index const&)> batch

std::function<bool(Tensor const&)> is_volatile_leaf

std::function<double(Index const&, std::size_t)> 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).

double volatile_weight = 1.0

Replay weight applied to volatile contractions in the flop tie-break.

double machine_balance = 0.0

Roofline parameters for the secondary (tie-break) cost; see RooflineParams and roofline_op_cost. machine_balance == 0 (default) => pure-flop tie-break (no behavior change). Uses full (unsliced) operand+result footprints (total per-replay traffic; slicing reduces peak, not total work).

double fast_mem_elems = 0.0

double block_tiles = 3.0

double block_prefactor = 1.0

bool batch_persistent_only = false

If true, batch only persistent subtrees (decline any subset containing a volatile leaf). Default false = batch across the board. See BatchPolicy::persistent_only.

double peak_flops_tolerance = 0.0

Unused by reconstruct (superseded by the threshold-gated selection below, driven by peak_threshold / numeric_size); retained for source compatibility. See PeakModel::peak_flops_tolerance, which is still consulted by the (unbatched) DenseSpaceTime model.

double accumulation_factor = 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 peak_threshold = std::numeric_limits<double>::infinity()

Peak-memory budget (bytes) for threshold-gated selection; see BatchPolicy::peak_threshold. +infinity (default) => min-flops (no batching).

double numeric_size = 8.0

Bytes per stored element, to compare the model’s element-count peak to peak_threshold (bytes). Default 8 (double / TensorD).

bool perf_first = false

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, so the threshold cannot force a flops-catastrophic factorization for its sliceability. Default false = peak-first, threshold-gated selection.

bool charge_batch_recompute = true

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 &#8212; 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.

The charge is order-aware: DP cells are ordered nests (see Context::ordered), so a node is billed only for the enclosing loops outer to its innermost-carried placement (Context::escaped_outer). Loops inner to that placement, and every loop when the node carries none of the enclosing modes, hoist above for free (rf == 1). An order-blind set charge would systematically over-bill the hoistable case.

bool charge_bound_persistence = true

If true (the default), a value bound to a batch loop is charged as often as the loop runs. A loop runs once per production of the value that closes it (the node contracting the mode in batches), so a non-volatile subtree sliced by a loop that a volatile node closes is rebuilt every replay &#8212; it cannot persist across evaluations (the runtime’s persistence rule, explicit-cells design section 12) &#8212; and takes volatile_weight exactly as an amplitude-dependent value does. A non-volatile nest closed by a non-volatile node (a persistent intermediate summed over the DF index from two persistent factors) runs once and is charged once, as before. Without this the model assumed every amplitude-independent subtree is computed once, ever, regardless of slicing, and priced slicing a persistent operand of an amplitude contraction as free &#8212; part of the water-20 occupied-batching cost gap (3x measured vs 1x modelled). Applied at the closing node, where its volatility is known: for each child subtree that is non-volatile and carries a mode sliced here, the child’s whole cost is scaled (an over-charge only for descendants of that child which do not carry the mode and hoist above the loop). An external loop the DP opens here (BFrontPoint::eopen) closes here too &#8212; the full result is assembled outside it, per production of this node &#8212; so it is charged the same way: the scale keys on the whole sliced set S (the contracted aprime or the external opens), not on aprime alone.

std::function<bool(Index const&)> is_batchable_contracted_index = [](Index const&) { return false; }

Spaces batchable in the contracted role &#8212; i.e. a mode of such a space is batchable where it is summed at some node. Building block; companion to 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<bool(Index const&)> is_batchable_external_index = [](Index const&) { return false; }

Spaces batchable in the external role &#8212; 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 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 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.

bool prune_outer_products = true

Prune disconnected (outer-product) subsets from the DP (see OptimizeOptions::prune_outer_products). Default true.

bool batch_spectator_indices = false

Term-level gate for reconstruct_batched_modes emitting BatchModeType::External entries (genuine external modes; see is_external_mode), threaded from 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. With it on, an external batch loop is a per-node DP choice (BFrontPoint::eopen): external bits become nestable cell modes and reconstruct_batched_modes stamps External exactly where the DP opened or inherited them. There is no post-DP placement pass.

struct BFrontPoint

One non-dominated (peak, flops) trade-off for a (subset, sliced-set B) cell. aprime is the sliced-set chosen at this node; the children are read at context C = B | aprime, at frontier indices lp_idx / rp_idx. See PeakModel::FrontPoint for why a frontier (not a single peak-min cell) is needed.

Public Members

double peak = std::numeric_limits<double>::max()

double flops = std::numeric_limits<double>::max()

size_t lp = 0

size_t rp = 0

bool lp_first = true

std::size_t aprime = 0

int lp_idx = -1

int rp_idx = -1

std::size_t nsl = 0

Cumulative count of batchable modes sliced anywhere in this realization (this node’s aprime popcount plus both children’s nsl). Used only by the perf-first ceiling’s threshold-aware frontier domination (pareto_insert_ceiling) to break peak-below-budget ties toward the least-sliced realization, so the unsliced schedule survives whenever it fits the budget. Zero for leaves and for the unbatched / peak-first paths, where it is never consulted.

std::size_t eopen = 0

External modes whose batch loop this node opens: carried by the node, absent from its enclosing set, nested outer of aprime. The node is produced per batch and scattered into its full result outside the loop; the children see the mode in their enclosing set, non-carriers under it pay the recompute charge, and the opener is the loop’s closer for charge_bound_persistence. Children are read at descend(descend(B, eopen), aprime).

struct Context

Precomputed tables and per-(subset, sliced-set) lookup parameters built once by build_context.

Public Functions

inline double fast_flops(std::size_t lp, std::size_t rp) const

Flop count of the binary contraction (subset lp) x (subset rp), equivalent to flops_of(idx[lp], idx[rp], idx[n]) with n = lp|rp, but reading precomputed atom-ID lists instead of rebuilding index sets per call. (idx[n] is a subset of idx[lp] U idx[rp], so it adds nothing.)

inline double sz_u(std::size_t U, std::size_t s) const

Footprint of subset s under an explicit sliced-set union mask (the order-independent bitmask, as returned by 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 &#8212; this is the primitive that sz and the external-placement re-price (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).

inline double Lof_u(std::size_t U, std::size_t s) const

Per-context leaf-sum of subset s under an explicit union mask.

inline double sz(std::size_t s, std::size_t id) const

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).

inline double Lof(std::size_t s, std::size_t id) const

Per-context leaf-sum of subset s (sum of singleton sizes under ctx).

inline std::size_t cell_union(std::size_t id) const

inline std::size_t descend(std::size_t id, std::size_t Ap) const

inline std::size_t descend_pt(std::size_t id, std::size_t eopen, std::size_t Ap) const

inline std::size_t escaped_outer(std::size_t id, std::size_t carried) const

inline void build_cells(bool ext_nestable)

Public Members

container::vector<Index> batchable_modes

Ordered, deduplicated batchable indices (bit k maps to batchable_modes[k]).

std::size_t m = 0

Number of batchable indices (= batchable_modes.size()).

std::size_t nB = 1

Number of sliced-sets (= 2^m).

std::size_t nt = 0

Number of tensors in the network.

container::vector<container::vector<double>> tables

tables[B][n] = footprint of subset n under sliced-set B.

container::vector<std::size_t> open_modes

open_modes[n] = bitmask of batchable indices open in subset n.

std::size_t volatile_mask = 0

Bitmask of volatile leaf tensors.

container::vector<IndexSet> idx

idx[n] = subset n’s open (result) indices, for the flop tie-break.

std::function<double(IndexSet const&, IndexSet const&, IndexSet const&)> flops_of

flops_of(lhs, rhs, result) = flop count of one binary contraction. Retained as the reference for the fast_flops parity test; the relax hot loop uses fast_flops (see below).

container::vector<double> nbatches

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.

bool use_fast_flops = true

container::vector<container::svector<std::uint32_t>> f_outer

container::vector<container::svector<std::uint32_t>> f_inner

container::vector<double> fo_ext

container::vector<double> fi_ext

container::vector<Index> fi_index

container::vector<std::uint32_t> fi_grp

bool f_inner_engaged = false

std::function<double(Index const&, std::size_t)> f_inner_pow

mutable container::vector<std::uint32_t> f_grpcount

mutable container::svector<std::uint32_t> f_union_in

bool ordered = true

std::size_t cap = 3

std::size_t nCells = 1

std::size_t external_mask = 0

container::vector<std::size_t> cell_union_

container::vector<container::svector<std::uint8_t>> cell_seq_

container::vector<std::size_t> cell_descend_