Batched (Memory-Bounded) Evaluation

A tensor contraction such as a density-fitted 2-electron integral, \(g_{ij}^{ab} \approx \sum_K B^a_{iK} B^b_{jK}\), can involve one mode (here \(K\), an auxiliary fitting-basis index — SeQuant’s convention labels it Κ, see sequant::mbpt::add_df_spaces()) that is far larger than the others. Materializing the full, unsliced intermediate for such a contraction can dominate peak memory even though the final result is comparatively small. Batching addresses this: SeQuant can slice a large mode into blocks and accumulate (or scatter) the contraction block by block, trading a controlled amount of extra bookkeeping/recompute for a lower peak-memory footprint. Which modes get batched, and how aggressively, is decided automatically by a cost-model search — the same optimize() machinery that chooses contraction order — rather than hand-coded.

Vocabulary

A few terms are used precisely throughout this page and the API:

  • mode: one dimension of a tensor (not to be confused with the sequant::Index that labels it).

  • batch vs. slice: a batch is the loop that partitions a mode’s work; a slice is one block of that mode’s data.

  • role: whether a mode is contracted (summed away at some node) or external (present in the final result). A contracted mode can be batched locally, right at the node that sums it. An external mode remains in the final result; batching it assembles disjoint result blocks. The optimizer chooses the contraction node that opens the external loop. External batching requires a separate opt-in (see batch_spectator_indices below).

  • external includes both a free mode (present in one operand and the result) and a Hadamard mode (present in both operands and the result), following the index roles of TAPP (Brandejs et al.). Here, batch means memory blocking; it does not mean a Hadamard mode, as in batched GEMM terminology.

  • combine: partials over a contracted mode are summed; partials over an external mode occupy disjoint result blocks and are scattered into a destination.

  • persistent vs. volatile: a persistent intermediate depends only on data that stays fixed across a calculation (e.g. two-electron integrals) and can be cached and reused; a volatile one depends on data that changes every iteration (e.g. coupled-cluster amplitudes) and must be recomputed each time.

  • peak memory: the maximum, over the whole evaluation, of the combined size of every tensor simultaneously resident — the quantity the batched cost models minimize or bound.

BatchPolicy

sequant::BatchPolicy is the single configuration object consulted both when sequant::optimize() decides which modes to batch and when the runtime evaluator actually executes that decision. Its full field list is API-reference material, but seven fields matter for everyday use:

  • peak_threshold: the master on/off switch — a peak-memory budget in bytes. The default, \(+\infty\), means no batching at all; setting a finite value turns batching on and sets the budget the optimizer searches under (see below for what happens when no schedule fits).

  • is_batchable_contracted_index / is_batchable_external_index: predicates selecting which index spaces may be sliced, split by role (see above). Both default to “nothing is batchable”; a caller opts specific spaces in explicitly.

  • batch_target_size: a per-index upper bound on block size, in elements; the runtime rounds down to whole backend tiles. A single tile larger than the bound forms its own batch, so the bound cannot override the backend’s tile geometry.

  • batch_spectator_indices: opts external modes into batching as well (in addition to contracted ones); off by default.

  • is_volatile_leaf / persistent_only: identify amplitude-like leaves and, optionally, restrict batching over contracted modes to subtrees that contain none of them.

mbpt::load(mbpt::Convention::SR);

// register a density-fitting auxiliary space ("Κ") and make it much larger
// than the occupied/virtual spaces -- as in forming a density-fitted
// 2-electron integral g[i,j;a,b] ≈ Σ_Κ B[i,a,Κ] B[j,b,Κ], where Κ ranges
// over a large auxiliary fitting basis
auto isr = std::make_shared<IndexBasisRegistry>(
    *get_default_context().index_basis_registry());
mbpt::add_df_spaces(isr);
isr->extent(L"i", 4);
isr->extent(L"a", 8);
isr->extent(L"Κ", 100'000);
set_default_context(get_default_context_snapshot().set(isr));

auto expr = ex<Tensor>(L"B", bra{L"i_1"}, ket{L"a_1", L"Κ_1"}) *
            ex<Tensor>(L"B", bra{L"i_2"}, ket{L"a_2", L"Κ_1"}) *
            ex<Tensor>(L"S", bra{L"i_1"}, ket{L"i_2"});

Deciding what to batch

Passing a BatchPolicy to sequant::optimize() via OptimizeOptions::batch_policy, together with the batched objective function ObjectiveFunction::DenseTimeSpaceBatched (performance first, peak memory second), makes the cost-based search consider slicing as one more way to reduce peak memory, in addition to choosing contraction order:

// a BatchPolicy declares which index spaces may be sliced (split by
// contracted vs. external role), an upper bound on the per-slice block
// size, and -- via peak_threshold -- the peak-memory budget that actually
// turns batching on; the default peak_threshold of +infinity means "never
// batch"
BatchPolicy policy{
    .is_batchable_contracted_index =
        [](Index const& ix) { return ix.space().base_key() == L"Κ"; },
    .batch_target_size = [](Index const&) -> std::size_t { return 100; },
    .peak_threshold = 5'000'000.0};

// term_batch_axes is an optional out-channel: when set, optimize() records
// which indices it decided to slice at each contraction node
auto batch_axes = std::make_shared<std::unordered_map<
    Expr const*, container::vector<NodeBatchAnnotation>>>();

auto optimized = optimize(
    expr, {.objective_function = ObjectiveFunction::DenseTimeSpaceBatched,
           .batch_policy = policy,
           .term_batch_axes = batch_axes});

bool batched_kappa = false;
for (auto const& [node, annotations] : *batch_axes)
  for (auto const& annotation : annotations)
    for (auto const& [axis, mode_type] : annotation.axes)
      batched_kappa |= axis.space().base_key() == L"Κ" &&
                       mode_type == BatchModeType::Contracted;

The finite peak_threshold is a ceiling on the modeled peak. Among the schedules that fit it, the optimizer picks the lowest performance cost, then the fewest sliced modes, then the lowest peak, so nothing is sliced unless slicing is needed to fit the budget. If no schedule fits, it picks the lowest performance cost, then the lowest peak.

Exceeding the budget is not an error. The modeled peak is estimated for each sequant::Product separately, from dense extents; backend workspace, tile sizes, and intermediates shared between expressions can change the realized peak, so a finite budget does not guarantee that the process fits in that many bytes.

Executing a batched plan

Once optimize() has annotated an expression with a batching decision, actually running it numerically needs a real tensor backend (TiledArray, BTAS, or TAPP) and is reached through the same sequant::evaluate() entry points used for any other evaluation, plus sequant::make_evaluator() to adapt a BatchPolicy into the evaluator SeQuant’s cache consults. Two execution strategies are available, selected by BatchPolicy::scheduler: the default BatchScheduler::forest_descent evaluates one expression at a time, while BatchScheduler::ordered fuses all the expressions evaluated together (e.g. all residual equations of a coupled-cluster method) into a single schedule, so an intermediate several of them need can be computed once inside shared batch loops. Because these entry points are backend-dependent, they are outside the scope of a backend-agnostic example here; the internal architecture behind both strategies is documented for contributors in Batched Evaluation Architecture, and the cost model driving the decisions above in Cost Model and Single-Term Optimization.

To estimate the peak-memory impact of a BatchPolicy before running a real (and potentially expensive) calculation, SeQuant also provides a dry-run meter (sequant::eval::dryrun::meter()) that replays the policy-selected scheduler against a zero-data sizing backend, reporting the realized peak without allocating real tensors. Numerical accuracy still requires a real evaluation. See Dry runs and diagnostics.