Program Listing for File eval.hpp¶
↰ Return to documentation for file (SeQuant/core/eval/eval.hpp)
#ifndef SEQUANT_EVAL_EVAL_HPP
#define SEQUANT_EVAL_EVAL_HPP
#include <SeQuant/core/eval/fwd.hpp>
#include <SeQuant/core/batch_policy.hpp>
#include <SeQuant/core/container.hpp>
#include <SeQuant/core/eval/cache_manager.hpp>
#include <SeQuant/core/eval/cell_registry.hpp>
#include <SeQuant/core/eval/eval_node.hpp>
#include <SeQuant/core/eval/lifetime_mask.hpp>
#include <SeQuant/core/eval/ordered_dump.hpp>
#include <SeQuant/core/eval/result.hpp>
#include <SeQuant/core/eval/schedule_dump.hpp>
#include <SeQuant/core/eval/slicing_signature.hpp>
#include <SeQuant/core/expr.hpp>
#include <SeQuant/core/io/serialization/serialization.hpp>
#include <SeQuant/core/logger.hpp>
#include <SeQuant/core/meta.hpp>
#include <SeQuant/core/optimize/optimize.hpp>
#include <SeQuant/core/utility/exception.hpp>
#include <SeQuant/core/utility/macros.hpp>
#include <SeQuant/core/utility/string.hpp>
#include <range/v3/range/operations.hpp>
#include <algorithm>
#include <any>
#include <atomic>
#include <chrono>
#include <cstdlib>
#include <deque>
#include <iostream>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <vector>
// Headers for process_rss_bytes() — see log::process_rss_bytes() below.
#if defined(__APPLE__)
#include <mach/mach.h>
#elif defined(__linux__)
#include <unistd.h>
#include <fstream>
#endif
namespace sequant {
namespace log {
using Duration = std::chrono::nanoseconds;
struct Bytes {
size_t value;
};
[[nodiscard]] inline bool printing() noexcept {
return Logger::instance().eval.level > 0;
}
template <typename T, typename... Ts>
[[nodiscard]] inline auto bytes(T const& arg, Ts const&... args) {
auto one = [](auto const& a) -> size_t {
if constexpr (requires {
static_cast<bool>(a);
a->size_in_bytes();
}) {
// Smart-pointer-like operand: tolerate null so callers (e.g. the
// EvalOp::Adjoint dispatcher, which leaves `right` unevaluated) can
// pass an empty ResultPtr without an external guard.
return a ? a->size_in_bytes() : size_t{0};
} else if constexpr (requires { a->size_in_bytes(); })
return a->size_in_bytes();
else
return a.size_in_bytes();
};
return Bytes{(one(arg) + ... + one(args))};
}
template <typename N, bool F, typename... Ts>
[[nodiscard]] inline Bytes bytes(CacheManager<N, F> const& cache,
Ts const&... args) {
if (!printing()) return Bytes{0};
return Bytes{cache.size_in_bytes() + (size_t{0} + ... + bytes(args).value)};
}
[[nodiscard]] inline auto to_string(Bytes bs) noexcept {
return std::format("{}B", bs.value);
}
[[nodiscard]] inline std::size_t process_rss_bytes() noexcept {
#if defined(__APPLE__)
::task_vm_info_data_t vm_info{};
::mach_msg_type_number_t vm_count = TASK_VM_INFO_COUNT;
if (::task_info(::mach_task_self(), TASK_VM_INFO,
reinterpret_cast<::task_info_t>(&vm_info),
&vm_count) == KERN_SUCCESS &&
vm_count >= TASK_VM_INFO_COUNT) {
return static_cast<std::size_t>(vm_info.phys_footprint);
}
// Fallback: raw resident-set size (larger; includes shared pages).
::mach_task_basic_info_data_t info{};
::mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;
if (::task_info(::mach_task_self(), MACH_TASK_BASIC_INFO,
reinterpret_cast<::task_info_t>(&info),
&count) != KERN_SUCCESS) {
return 0;
}
return static_cast<std::size_t>(info.resident_size);
#elif defined(__linux__)
// /proc/self/statm columns are page counts:
// total resident shared text lib data dt
std::ifstream f("/proc/self/statm");
std::size_t pages_total = 0, pages_resident = 0;
if (!(f >> pages_total >> pages_resident)) return 0;
static const long page_size = ::sysconf(_SC_PAGESIZE);
if (page_size <= 0) return 0;
return pages_resident * static_cast<std::size_t>(page_size);
#else
return 0;
#endif
}
[[nodiscard]] inline Bytes rss() noexcept { return Bytes{process_rss_bytes()}; }
enum struct EvalMode {
Constant,
Variable,
Power,
Tensor,
Permute,
Product,
MultByPhase,
Sum,
SumInplace,
Symmetrize,
Antisymmetrize,
Unknown
};
[[nodiscard]] EvalMode eval_mode(meta::eval_node auto const& node) {
if (node.leaf()) {
return node->is_constant() ? EvalMode::Constant
: node->is_variable() ? EvalMode::Variable
: node->is_power() ? EvalMode::Power
: node->is_tensor() ? EvalMode::Tensor
: EvalMode::Unknown;
} else {
return node->is_product() ? EvalMode::Product
: node->is_sum() ? EvalMode::Sum
: node->is_adjoint() ? EvalMode::Permute
: EvalMode::Unknown;
}
}
[[nodiscard]] constexpr auto to_string(EvalMode mode) noexcept {
return (mode == EvalMode::Constant) ? "Constant"
: (mode == EvalMode::Variable) ? "Variable"
: (mode == EvalMode::Power) ? "Power"
: (mode == EvalMode::Tensor) ? "Tensor"
: (mode == EvalMode::Permute) ? "Permute"
: (mode == EvalMode::Product) ? "Product"
: (mode == EvalMode::MultByPhase) ? "MultByPhase"
: (mode == EvalMode::Sum) ? "Sum"
: (mode == EvalMode::SumInplace) ? "SumInplace"
: (mode == EvalMode::Symmetrize) ? "Symmetrize"
: (mode == EvalMode::Antisymmetrize) ? "Antisymmetrize"
: "??";
}
enum struct CacheMode { Store, Access, Release };
[[nodiscard]] constexpr auto to_string(CacheMode mode) noexcept {
return (mode == CacheMode::Store) ? "Store"
: (mode == CacheMode::Access) ? "Access"
: "Release";
}
enum struct TermMode { Begin, End };
[[nodiscard]] constexpr auto to_string(TermMode mode) noexcept {
return (mode == TermMode::Begin) ? "Begin" : "End";
}
// clang-format off
// clang-format on
struct EvalStat {
EvalMode mode;
Duration time;
Bytes mem_result{};
Bytes mem_alloc{};
Bytes mem_hwmark{};
std::optional<Bytes> mem_left{};
std::optional<Bytes> mem_right{};
};
struct CacheStat {
CacheMode mode;
size_t key;
int curr_life, max_life;
size_t num_alive;
Bytes entry_memory;
Bytes total_memory;
};
template <typename Arg, typename... Args>
void log(Arg const& arg, Args const&... args) {
auto& l = Logger::instance();
if (l.eval.level > 0) write_log(l, arg, std::format(" | {}", args)..., '\n');
}
template <typename... Args>
auto eval(EvalStat const& stat, Args const&... args) {
if (!printing()) return; // nothing to format/emit; skip rss() and formatting
auto const result_s = std::format("result={}", to_string(stat.mem_result));
auto const alloc_s = std::format("alloc={}", to_string(stat.mem_alloc));
auto const hw_s = std::format("hw={}", to_string(stat.mem_hwmark));
// Reduce this rank's RSS to the value to report (e.g. sum over ranks = true
// total app memory). This runs on every rank (printing() is level>0,
// identical across ranks), so an injected collective reducer is matched.
auto const rss_local = process_rss_bytes();
auto const& rss_reduce = Logger::instance().eval.rss_reduce;
auto const rss_s = std::format(
"rss={}",
to_string(Bytes{rss_reduce ? rss_reduce(rss_local) : rss_local}));
// Optional backend-supplied suffix (already reduced across ranks); omitted
// entirely when the hook is unset or returns empty.
auto const& heap_stats = Logger::instance().eval.heap_stats;
auto const heap_s = heap_stats ? heap_stats() : std::string{};
auto emit = [&](auto const&... trailer) {
if (stat.mem_left) {
SEQUANT_ASSERT(stat.mem_right);
log("Eval", //
to_string(stat.mode), //
stat.time, //
std::format("left={}", to_string(*stat.mem_left)), //
std::format("right={}", to_string(*stat.mem_right)), //
result_s, alloc_s, hw_s, rss_s, //
trailer...);
} else {
log("Eval", //
to_string(stat.mode), //
stat.time, //
result_s, alloc_s, hw_s, rss_s, trailer...);
}
};
if (heap_s.empty())
emit(args...);
else
emit(heap_s, args...);
}
template <typename... Args>
auto cache(CacheStat const& stat, Args const&... args) {
log("Cache", //
to_string(stat.mode), //
std::format("key={}", stat.key), //
std::format("life={}/{}", stat.curr_life, stat.max_life), //
std::format("alive={}", stat.num_alive), //
std::format("entry={}", to_string(stat.entry_memory)), //
std::format("total={}", to_string(stat.total_memory)), //
args...);
}
template <typename N, bool F, typename... Args>
auto cache(N const& node, CacheManager<N, F>& cm, Args const&... args) {
// Structured runtime-schedule event, emitted to the cache's ScheduleSink when
// one is wired (independent of the trace level). Keyed by node->hash_value()
// -- the same identity the IR emitter (schedule_dump.hpp) writes -- so the
// schedule visualizer joins runtime lifetimes onto the IR DAG by hash. Mode:
// Store (first build), Access (reuse), Release (last use). Store-count > 1
// for a hash means the value was rebuilt (recompute).
if (auto* const sink = cm.schedule_sink(); sink && sink->os && !sink->fired) {
auto const cl = cm.life(node);
auto const ml = cm.max_life(node);
char const* const evm = (cl == 0) ? "Release"
: (cl + 1 == ml) ? "Store"
: "Access";
*sink->os << "SCHEDULE_RUN_EVENT {\"hash\":\"" << node->hash_value()
<< "\",\"mode\":\"" << evm << "\",\"life\":" << cl
<< ",\"max_life\":" << ml << "}\n";
}
if (!printing()) return; // skip the entry/total size walks and formatting
using CacheMode::Access;
using CacheMode::Release;
using CacheMode::Store;
auto const key = hash::value(*node);
auto const cur_l = cm.life(node);
auto const max_l = cm.max_life(node);
bool const release = cur_l == 0;
bool const store = cur_l + 1 == max_l;
cache(CacheStat{.mode = store ? Store
: release ? Release
: Access,
.key = key,
.curr_life = cur_l,
.max_life = max_l,
.num_alive = cm.alive_count(),
.entry_memory = {cm.entry_size_in_bytes(node)},
.total_memory = {bytes(cm)}},
args...);
}
inline auto term(TermMode mode, std::string_view term) {
log("Term", to_string(mode), term);
}
inline void release_after_op() {
if (auto const& rel = Logger::instance().eval.release_memory; rel) rel();
}
[[nodiscard]] auto label(meta::eval_node auto const& node) {
// Guard on the structural leaf() (no children), not the semantic
// node->is_primary(): a non-primary leaf (a constant/variable leaf, or any
// leaf whose expr is not a primary tensor) is is_primary()==false yet has no
// children, so the else branch's node.left()/node.right() would throw
// (checked_ptr_access) -- the trace-only crash that fails every eval test
// under SEQUANT_EVAL_TRACE.
return node.leaf() ? node->label()
: std::format("{} {} {} -> {}", node.left()->label(),
(node->is_product() ? "*"
: node->is_sum() ? "+"
: "??"), //
node.right()->label(), node->label());
}
template <typename BatchContext>
std::string scope_annot(BatchContext const& active) {
std::string scope;
for (auto const& e : active) {
if (!scope.empty()) scope += ",";
scope += toUtf8(e.axis.space().base_key());
}
return std::format("scope={{{}}}", scope);
}
template <meta::eval_node Node, typename BatchContext>
std::string slice_home_annot(Node const& node, BatchContext const& active) {
std::string canon, sliced;
auto const& idxs = node->canon_indices();
for (auto const& ix : idxs) {
canon += toUtf8(ix.full_label());
canon += " ";
}
container::svector<std::size_t> spos;
for (auto const& sm : node->sliced_modes())
if (auto const p = index_position(node, sm);
p && std::find(spos.begin(), spos.end(), *p) == spos.end())
spos.push_back(*p);
std::sort(spos.begin(), spos.end());
for (auto const p : spos)
sliced += std::format("{}:{} ", p,
toUtf8(std::wstring(idxs[p].space().base_key())));
auto annot = std::format("canon=[{}] sliced=[{}] {}", canon, sliced,
scope_annot(active));
// Append any schedule-derived per-node metadata (e.g. the value's home
// scope + use scopes -- properties the running annotation cannot see). Empty
// provider => nothing appended => the base annotation as it stands.
if (auto const& nm = Logger::instance().eval.node_meta; nm)
annot += " " + nm(node->hash_value());
return annot;
}
template <meta::eval_node Node, typename BatchContext>
[[nodiscard]] std::string label(Node const& node, BatchContext const& active) {
return std::format("{} {}", label(node), slice_home_annot(node, active));
}
} // namespace log
// implementation details of the eval engine; prefer sequant::detail over an
// unnamed namespace in a header (see CppCoreGuidelines SF.21)
namespace detail {
template <typename F, typename... Args>
[[nodiscard]] log::Duration timed_eval_inplace(F&& fun, Args&&... args)
requires(std::is_invocable_r_v<void, F, Args...>)
{
using Clock = std::chrono::high_resolution_clock;
auto tstart = Clock::now();
std::forward<F>(fun)(std::forward<Args>(args)...);
auto tend = Clock::now();
return {tend - tstart};
}
template <typename T>
constexpr bool is_cache_manager_v = false;
template <typename N, bool F>
constexpr bool is_cache_manager_v<CacheManager<N, F>> = true;
// True if any of Args is a CacheManager. The fold over `||` is well-defined
// (and false) for an empty pack, so there is no tuple_element underflow to
// guard against. Used to keep the variadic cache-appending
// evaluate()/evaluate_impl() forwarders below from matching a call that already
// carries a cache -- e.g. the scope-executor overload
// evaluate(forest, policy, layout, leaf, cache, mode_order, guard), whose
// CacheManager is not the last argument. A last-argument-only check would fire
// on that overload, append a second cache, and fail to resolve.
template <typename... Args>
concept any_type_is_cache_manager =
(... || is_cache_manager_v<std::remove_cvref_t<Args>>);
template <typename... Args>
auto&& arg0(Args&&... args) {
return std::get<0>(std::forward_as_tuple(std::forward<Args>(args)...));
}
auto&& node0(auto&& val) { return std::forward<decltype(val)>(val); }
auto&& node0(std::ranges::range auto&& rng) {
return ranges::front(std::forward<decltype(rng)>(rng));
}
enum struct CacheCheck { Checked, Unchecked };
} // namespace detail
enum struct Trace {
On,
Off,
Default =
#ifdef SEQUANT_EVAL_TRACE
On
#else
Off
#endif
};
static_assert(Trace::Default == Trace::On || Trace::Default == Trace::Off);
// implementation details of the eval engine; prefer sequant::detail over an
// unnamed namespace in a header (see CppCoreGuidelines SF.21)
namespace detail {
[[nodiscard]] consteval bool trace(Trace t) noexcept { return t == Trace::On; }
} // namespace detail
[[nodiscard]] inline Index::index_vector contracted_indices(
meta::eval_node auto const& node) {
Index::index_vector result;
if (node.leaf() || !node->is_product()) return result;
auto const& l = node.left()->canon_indices();
auto const& r = node.right()->canon_indices();
auto const& c = node->canon_indices();
auto contains = [](auto const& vec, Index const& ix) {
return std::find(vec.begin(), vec.end(), ix) != vec.end();
};
for (Index const& ix : l)
if (contains(r, ix) && !contains(c, ix)) result.push_back(ix);
return result;
}
template <typename IndexPredicate>
[[nodiscard]] inline std::optional<Index> batch_axis(
meta::eval_node auto const& node, IndexPredicate const& accept) {
std::optional<Index> best;
for (Index const& ix : contracted_indices(node)) {
if (!accept(ix)) continue;
if (!best ||
best->space().approximate_size() < ix.space().approximate_size())
best = ix;
}
return best;
}
[[nodiscard]] inline std::optional<Index> batch_axis(
meta::eval_node auto const& node) {
return batch_axis(node, [](Index const&) { return true; });
}
// index_position() lives in SeQuant/core/eval/slicing_signature.hpp (shared
// with the hoist path) and is available here through that include.
template <typename Node>
[[nodiscard]] std::optional<std::pair<Node, std::size_t>> find_leaf_carrying(
Node const& node, Index const& ix) {
if (node.leaf()) {
if (auto const p = index_position(node, ix)) return std::pair{node, *p};
return std::nullopt;
}
if (auto found = find_leaf_carrying(node.left(), ix)) return found;
return find_leaf_carrying(node.right(), ix);
}
// The tree-walking evaluation engine, and the forest-descent path's own: the
// batched forest evaluator re-enters this (evaluate_impl), never `evaluate`,
// so `evaluate` reads as the single outermost call and the whole batched
// recursion is contained in evaluate_impl. External callers use the thin
// `evaluate` overloads (below); internal re-entries (the scatter/contraction
// per-block re-evaluations and the hoisted-invariant builds) call
// evaluate_impl directly. The ordered (table-driven) executor does not come
// through here at all -- it computes each cell from its own table reads (\c
// eval::detail::compute_cell, ordered_executor.hpp), sharing this file's
// per-op kernels (apply_one_op_traced, sum_in_place_traced, fetch_leaf_traced,
// apply_canon_phase, try_custom_eval, note_fresh_build) rather than this
// stack machine.
template <meta::can_evaluate Node>
[[nodiscard]] ResultPtr apply_one_op(Node const& node, ResultPtr const& left,
ResultPtr const& right) {
if (node->op_type() == EvalOp::Adjoint) {
// Unary: only the left operand; the right child is the Constant(1)
// sentinel.
std::array<std::any, 2> const adj_ann{node.left()->annot(), node->annot()};
return left->adjoint(adj_ann);
}
std::array<std::any, 3> const ann{node.left()->annot(), node.right()->annot(),
node->annot()};
if (node->op_type() == EvalOp::Sum) return left->sum(*right, ann);
SEQUANT_ASSERT(node->op_type() == EvalOp::Product);
bool const de_nest =
node.left()->tot() && node.right()->tot() && !node->tot();
return left->prod(*right, ann, de_nest ? DeNest::True : DeNest::False);
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate Node, typename N,
bool FHC>
[[nodiscard]] ResultPtr apply_canon_phase(Node const& nd, ResultPtr res,
CacheManager<N, FHC>& cache) {
auto phase = nd->canon_phase();
if (phase == 1) return res;
ResultPtr post;
auto const _ph0 = std::chrono::steady_clock::now();
auto time =
detail::timed_eval_inplace([&]() { post = res->mult_by_phase(phase); });
eval::EvalImplTimeline::note_phase(_ph0);
if constexpr (detail::trace(EvalTrace)) {
size_t hwmark = log::bytes(cache, post).value;
if (!cache.alive(nd)) hwmark += log::bytes(res).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
auto stat = log::EvalStat{
.mode = log::EvalMode::MultByPhase,
.time = time,
.mem_result = log::bytes(post),
.mem_alloc = log::bytes(post),
.mem_hwmark = {cache.note_working_set(hwmark, nd->hash_value())}};
log::eval(stat, std::format("{} * {}", phase, nd->label()),
log::slice_home_annot(nd, cache.batch_context()));
}
return post;
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate Node, typename N,
bool FHC>
[[nodiscard]] ResultPtr apply_one_op_traced(Node const& node,
ResultPtr const& left,
ResultPtr const& right,
CacheManager<N, FHC>& cache) {
// The contraction annotation triple the shaped-product hook receives. A
// unary (Adjoint) node never reaches the hook, and its right child is the
// Constant(1) sentinel, so the triple is only built for a binary op.
std::array<std::any, 3> const ann =
node->op_type() == EvalOp::Adjoint
? std::array<std::any, 3>{}
: std::array<std::any, 3>{node.left()->annot(), node.right()->annot(),
node->annot()};
ResultPtr result;
log::Duration time{};
if (node->op_type() != EvalOp::Product) {
time = detail::timed_eval_inplace(
[&]() { result = apply_one_op(node, left, right); });
} else {
// Consult the shaped-product hook (if set) before evaluating the
// product. The hook receives the node (wrapped in a std::any as a
// std::reference_wrapper so the full IR node is inspectable) plus the
// evaluated operands and annotations; a non-null return *replaces*
// the normal product (e.g. a shape-constrained emission of it), a
// null return declines and the standard prod() below runs. An empty
// hook is never consulted; default-empty => the standard product runs.
auto const _tp0 = std::chrono::steady_clock::now(); // node-eval start
if (auto const& hook = cache.shaped_product_hook(); hook) {
time = detail::timed_eval_inplace([&]() {
result = hook(std::any{std::cref(node)}, *left, *right, ann);
});
}
if (!result) {
// Sentinel so the recompute tally below fires only when a DryRun
// prod actually computed fresh flops for this node. DryRunOps::prod
// early-returns without setting last_op_flops for a scalar*tensor
// product (and it is never set at all by the wet TA backend); a
// stale value from a previous op must not be attributed here (it
// would fold garbage into the identity-keyed tally). prod sets
// last_op_flops >= 0 only on the real contraction path.
eval::detail::last_op_flops() = -1.0;
time = detail::timed_eval_inplace(
[&]() { result = apply_one_op(node, left, right); });
// Record this product build against node's identity (keyed by the
// exact cache identity: hash-bin + Bliss, so 64-bit hash collisions
// are not folded) in the (root) cache's recompute tally. Deduped at
// the (value, slice) granularity using actual replay FLOPs (DryRun
// prod just stashed this build's realized cost in last_op_flops): a
// value built over distinct slices is tiling (not recompute); the
// same value rebuilt at the same slice -- including a node
// invariant to an enclosing loop, rebuilt every block -- is
// recompute. The slice signature is the live batch context
// projected onto the modes node carries (find_leaf_carrying), so
// an invariant's projection is identical (empty for that loop)
// every block and its rebuilds fold. A no-op unless the dry-run
// replay enabled the tally on the root cache (the wet TA path
// leaves it disabled), so nothing is tallied off the costing path.
if (eval::detail::last_op_flops() >= 0.0) {
std::string slice_sig;
for (auto const& entry : cache.batch_context()) {
Index const& ix = entry.axis;
auto const& blk = entry.range;
if (find_leaf_carrying(node, ix).has_value()) {
slice_sig += toUtf8(ix.full_label());
slice_sig += ':';
slice_sig += std::to_string(blk.first);
slice_sig += ';';
}
}
cache.tally_build(node, slice_sig, eval::detail::last_op_flops(),
eval::detail::last_op_exec());
}
}
// node-eval done: fence + accumulate the contraction's full cost.
eval::EvalImplTimeline::note_prod(_tp0);
}
SEQUANT_ASSERT(result);
if constexpr (detail::trace(EvalTrace)) {
// Skip an operand's bytes only when it aliases a cache buffer that is
// already counted (locally in bytes(cache,...) or up-chain in
// chain_residency()). chain_holds() tests pointer identity against
// every alive entry on the scope chain: a cached child fetched full
// aliases and is skipped; an apply_phase / sliced / permuted child is
// a distinct buffer and is added.
size_t hwmark = log::bytes(cache, result).value;
if (!cache.chain_holds(left)) hwmark += log::bytes(left).value;
if (right && !cache.chain_holds(right)) hwmark += log::bytes(right).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
log::eval(log::EvalStat{.mode = log::eval_mode(node),
.time = time,
.mem_result = log::bytes(result),
.mem_alloc = log::bytes(result),
.mem_hwmark = {cache.note_working_set(
hwmark, node->hash_value())},
.mem_left = log::bytes(left),
.mem_right = log::bytes(right)},
log::label(node, cache.batch_context()) + " | L." +
left->trange_annot() +
(right ? " R." + right->trange_annot() : std::string{}) +
" O." + result->trange_annot());
}
log::release_after_op();
return result;
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate Node, typename N,
bool FHC>
[[nodiscard]] ResultPtr sum_in_place_traced(Node const& node, ResultPtr left,
ResultPtr const& right,
CacheManager<N, FHC>& cache) {
std::array<std::any, 3> const ann{node.left()->annot(), node.right()->annot(),
node->annot()};
ResultPtr result;
log::Duration time{};
// The accumulator (left) is used as is, in its own layout --
// never repermuted -- since binarize() pins a marked Sum's
// canon_indices_ to its left operand's (see EvalExpr::binarize
// (Sum)'s make_sum lambda), so this node's own target layout is
// left's layout. The right addend generally is not already in
// that layout (each summand of the original N-ary Sum
// canonicalizes independently), so -- exactly as the allocating
// sum() path below permutes both operands into the result's
// layout (see e.g. ResultTensorBTAS::sum) -- it must be permuted
// into left's layout (ann[0]) before the raw elementwise
// add_inplace; skipping this for a tensor Sum silently adds
// mismatched layouts. Scalars carry no layout: ResultScalar::
// permute() is unimplemented (throws) and none is needed, since
// ann[0]/ann[1]/ann[2] are all trivially empty for a scalar Sum.
bool const needs_permute = node->result_type() == ResultType::Tensor;
ResultPtr right_aligned = right;
log::Duration perm_time{};
if (needs_permute) {
perm_time = detail::timed_eval_inplace([&]() {
right_aligned = right->permute(std::array<std::any, 2>{ann[1], ann[0]});
});
}
if constexpr (detail::trace(EvalTrace)) {
if (needs_permute) {
// Mirrors the top-level evaluate(node, layout, ...) Permute
// event above: a genuine fresh allocation for the (typically
// much smaller, single-term) right addend, logged separately
// from the SumInplace event below.
size_t hwmark = log::bytes(cache, right_aligned).value;
if (!cache.chain_holds(right)) hwmark += log::bytes(right).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
log::eval(log::EvalStat{.mode = log::EvalMode::Permute,
.time = perm_time,
.mem_result = log::bytes(right_aligned),
.mem_alloc = log::bytes(right_aligned),
.mem_hwmark = {cache.note_working_set(
hwmark, node->hash_value())}},
log::label(node, cache.batch_context()));
}
}
// Move the accumulator out of left (leaving it null; it is not
// read again below) and add the (now aligned) right addend into
// it: zero allocation here, unlike the fresh buffer sum() below
// returns.
time = detail::timed_eval_inplace([&]() {
result = std::move(left);
result->add_inplace(*right_aligned);
});
if constexpr (detail::trace(EvalTrace)) {
// Mirrors the forest-level SumInplace accounting (the
// Nodes-range evaluate() overload, below): bytes(cache, result)
// already counts the (mutated) former-accumulator buffer, since
// result now is that buffer, so only the (aligned) right addend
// is added, and only if it is not already resident on the scope
// chain. mem_alloc is zero -- SumInplace itself allocates
// nothing (the right-alignment allocation, if any, was already
// logged above as its own Permute event) -- and mem_left/
// mem_right are left unset, matching the "SumInplace | --- |
// 0B" row of the EvalStat doc table above.
size_t hwmark = log::bytes(cache, result).value;
if (!cache.chain_holds(right_aligned))
hwmark += log::bytes(right_aligned).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
log::eval(log::EvalStat{.mode = log::EvalMode::SumInplace,
.time = time,
.mem_result = log::bytes(result),
.mem_alloc = {0},
.mem_hwmark = {cache.note_working_set(
hwmark, node->hash_value())}},
log::label(node, cache.batch_context()));
}
log::release_after_op();
return result;
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate Node, typename N,
bool FHC>
[[nodiscard]] ResultPtr try_custom_eval(Node const& node,
CacheManager<N, FHC>& cache) {
SEQUANT_ASSERT(!node.leaf());
auto const& custom_eval = cache.custom_evaluator();
if (!custom_eval) return nullptr;
ResultPtr intercepted;
auto time = detail::timed_eval_inplace(
[&]() { intercepted = custom_eval(node, cache); });
if (!intercepted) return nullptr;
if constexpr (detail::trace(EvalTrace)) {
size_t hwmark = log::bytes(cache, intercepted).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
log::eval(log::EvalStat{.mode = log::eval_mode(node),
.time = time,
.mem_result = log::bytes(intercepted),
.mem_alloc = log::bytes(intercepted),
.mem_hwmark = {cache.note_working_set(
hwmark, node->hash_value())}},
log::label(node, cache.batch_context()));
}
log::release_after_op();
return intercepted;
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate Node, typename F,
typename N, bool FHC>
requires meta::leaf_node_evaluator<Node, F>
[[nodiscard]] ResultPtr fetch_leaf_traced(Node const& node,
F const& leaf_evaluator,
CacheManager<N, FHC>& cache) {
ResultPtr result;
auto time =
detail::timed_eval_inplace([&]() { result = leaf_evaluator(node); });
if constexpr (detail::trace(EvalTrace)) {
size_t hwmark = log::bytes(cache, result).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
log::eval(log::EvalStat{.mode = log::eval_mode(node),
.time = time,
.mem_result = log::bytes(result),
.mem_alloc = log::bytes(result),
.mem_hwmark = {cache.note_working_set(
hwmark, node->hash_value())}},
log::label(node, cache.batch_context()));
}
log::release_after_op();
return result;
}
template <meta::can_evaluate Node, typename N, bool FHC>
void note_fresh_build(Node const& node, CacheManager<N, FHC>& cache) {
// Diagnostic (SEQUANT_UT_BUILD_METER): count actual builds at this single
// chokepoint (non-leaf only = real contraction executions). No-op when off.
if (!node.leaf())
eval::BuildMeter::on_build(node->hash_value(),
eval::BuildMeter::enabled()
? log::label(node, cache.batch_context())
: std::string{});
// Per-op build event: finish_phase_b is the single choke point every
// freshly computed node passes through -- leaves, custom-eval subtrees, and
// standard contractions -- so this counts every build (cached or not),
// giving the schedule visualizer full per-node recompute coverage (the
// cache Store/ Access/Release events above cover only cached nodes). Keyed
// by hash_value() to join onto the IR DAG. Emitted to the cache's
// ScheduleSink when one is wired (set_schedule_sink); no sink => no dump.
if (auto* const sink = cache.schedule_sink();
sink && sink->os && !sink->fired) {
std::ostream& os = *sink->os;
// ctx = the active batch loops (mode -> block offset) this build ran
// under. The visualizer counts distinct ctx projections onto the modes a
// node depends on: builds at a repeated projected slice are avoidable
// recompute (a value rebuilt where an enclosing loop it is invariant to
// advanced) vs the inherent per-block work batching requires.
// Leaves are accessed, not computed (leaf_evaluator just hands back a
// ref to a precomputed input), so tag them "Fetch" -- cheap, not
// recompute. Only internal nodes (contractions) are real "Build" work.
char const* const bmode = node.leaf() ? "Fetch" : "Build";
os << "SCHEDULE_RUN_EVENT {\"hash\":\"" << node->hash_value()
<< "\",\"mode\":\"" << bmode << "\"";
// sig = the avoidable-recompute join key, matching the IR node record.
// Internal nodes only; leaves (Fetch) are not tallied. Lets the
// visualizer join this hash to the replay's per-node avoidable without
// reconstructing the signature in the renderer.
// Per-build flops (cm->flops) keyed to this node by hash above; the
// recompute rollup and the visualizer join on the topological hash, not
// a dummy-/slice-dependent signature.
if (!node.leaf()) os << ",\"flops\":" << eval::detail::last_op_flops();
os << ",\"ctx\":[";
bool first = true;
for (auto const& entry : cache.batch_context()) {
Index const& ix = entry.axis;
auto const& blk = entry.range;
// dep = does the node's subtree carry this loop mode (free or
// contracted below)? If not, the node is invariant to it and rebuilding
// per block of it is avoidable recompute. find_leaf_carrying works in
// the node's own label space, so no alpha-renaming reconciliation.
bool const dep = find_leaf_carrying(node, ix).has_value();
os << (first ? "" : ",") << "[\"" << toUtf8(ix.full_label()) << "\","
<< blk.first << "," << (dep ? 1 : 0) << "]";
first = false;
}
os << "]}\n";
}
}
template <Trace EvalTrace = Trace::Default,
detail::CacheCheck Cache = detail::CacheCheck::Checked,
meta::can_evaluate Node, typename F, typename N, bool FHC>
requires meta::leaf_node_evaluator<Node, F>
ResultPtr evaluate_impl(Node const& node, //
F const& leaf_evaluator, //
CacheManager<N, FHC>& cache) {
// Diagnostic (SEQUANT_UT_EVALIMPL): split CC-eval into time inside
// top-level evaluate_impl (body) vs the gap between successive top-level
// calls (executor/call-site machinery). See EvalImplTimeline.
eval::EvalImplTimeline::Scope _tl_evalimpl;
// B-full accounting: an entry into this engine. The ordered executor
// computes every cell itself (detail::compute_cell, ordered_executor.hpp),
// so an ordered run must leave this counter at zero -- see OrderedOpCounts.
++eval::detail::ordered_op_counts_slot().probes;
// Multiply a (possibly cached) result by its node's canonicalization phase
// (apply_canon_phase, above -- shared with the ordered executor's own
// compute_cell).
auto apply_phase = [&cache](auto const& nd, ResultPtr res) -> ResultPtr {
return apply_canon_phase<EvalTrace>(nd, std::move(res), cache);
};
// Slice-on-use: slice a value fetched at the Enter stage to the current batch
// block for the `hops` innermost enclosing batch loops that `nd` carries and
// that the value does not yet have baked in. `hops` == number of enclosing
// loops crossed to reach the value's lifetime scope (0 for a local hit / a
// freshly built value in this scope; d == batch_context().size() for a fresh
// leaf whose lifetime is top). The slice set is exactly
// (use scope minus lifetime scope) intersect carried(nd): the `hops`
// innermost batch_context entries, filtered by index_position(nd, axis).
// slice_mode is non-mutating, so a cached full value is left undisturbed.
// Empty batch_context (the off path) => d == hops == 0 => the loop is empty
// and the value is returned unchanged.
auto slice_to_use = [&cache](ResultPtr value, auto const& nd,
std::size_t hops) -> ResultPtr {
auto const& ctx = cache.batch_context();
std::size_t const d = ctx.size();
// hops (parent links access_at crossed) must not exceed d (batch_context
// entries): each realized loop pushes exactly one entry and wires at most
// one parent link, so hops <= d always. A violation would underflow
// `d - hops` and silently under-slice (oversized result); assert loudly.
SEQUANT_ASSERT(hops <= d);
// The batched forest evaluator pushes each member's own physical axis as
// `exact_axis`; a slice fires only where that is set (an intra-tree
// exact match on `nd`), so this is a forest-descent primitive only. The
// ordered executor does not come through here at all: it reads every
// operand from its already-sliced cell (CellReadResolver::fetch).
for (std::size_t i = d - hops; i < d; ++i) {
if (!ctx[i].exact_axis) continue;
auto const p_new = index_position(nd, *ctx[i].exact_axis);
if (p_new) {
auto const& blk = ctx[i].range;
value = value->slice_mode(*p_new, blk.first, blk.second);
}
}
return value;
};
// One entry of the explicit evaluation stack. `stage` records how far a node
// has progressed; `left`/`right` hold its evaluated operands; `store_after`
// marks a Checked node that exists in the cache map but has not been stored
// yet, so its computed result must be cached (this replaces the recursive
// wrapper's `evaluate<..., Unchecked>` re-entry).
enum class Stage { Enter, NeedLeft, NeedRight, NeedLeftAdj };
struct Frame {
// Non-owning pointer into the tree being evaluated (which outlives this
// call): a Node member would deep-copy the whole subtree into every frame
// (binary_node.hpp), so evaluating the left-leaning Sum-tree binarize
// builds for a whole equation -- one spine node per summand, each subtree
// larger than the last -- would cost O(terms^2) nodes. Fatal on a UCC BCH
// energy (thousands of terms).
Node const* node_p;
Node const& nd() const { return *node_p; }
bool checked;
Stage stage = Stage::Enter;
bool store_after = false;
ResultPtr left = {};
ResultPtr right = {};
};
// Finalize a freshly computed Phase-B result: if this Checked node needs
// storing, cache it (phase-applied) and hand back the phase-applied cached
// pointer -- exactly the recursive Checked wrapper's store path. Otherwise
// pass the raw result through unchanged.
auto finish_phase_b = [&cache, &apply_phase](Frame const& f,
ResultPtr rb) -> ResultPtr {
note_fresh_build(f.nd(), cache);
if (!f.store_after) return rb;
auto ptr =
cache.store_and_access(f.nd(), apply_phase(f.nd(), std::move(rb)));
if constexpr (detail::trace(EvalTrace))
log::cache(f.nd(), cache, log::label(f.nd(), cache.batch_context()));
return apply_phase(f.nd(), ptr);
};
// A `std::deque` is used so that a reference to the top frame stays valid
// across push_back (which reallocates a `std::vector`).
std::deque<Frame> stk;
stk.push_back(Frame{.node_p = &node,
.checked = (Cache == detail::CacheCheck::Checked)});
ResultPtr ret; // result handed up by the frame that most recently finalized
// Deliver `r` to the parent frame and pop the just-completed frame.
auto finalize = [&stk, &ret](ResultPtr r) {
ret = std::move(r);
stk.pop_back();
};
while (!stk.empty()) {
Frame& f = stk.back();
switch (f.stage) {
case Stage::Enter: {
// --- Checked cache wrapper: a hit returns directly; a miss on a node
// that exists in the map schedules a store once computed. ---
if (f.checked) {
if (auto m = cache.access_at(f.nd()); m.ptr) {
if constexpr (detail::trace(EvalTrace))
log::cache(f.nd(), cache,
log::label(f.nd(), cache.batch_context()));
// Slice-on-use: a value fetched `m.hops` scopes up does not have
// this scope's (and any intervening) batch slices baked in, so
// slice it to the current block for the loops the fetch crossed.
// A local hit (hops == 0) or the off path (empty batch_context)
// is a no-op, leaving apply_phase()'s own result.
finalize(slice_to_use(apply_phase(f.nd(), m.ptr), f.nd(), m.hops));
break;
}
f.store_after = cache.exists(f.nd());
}
// --- Custom-evaluator interception (non-leaf only): a non-null result
// short-circuits the subtree -- children are never pushed. This is
// the subtree pruning batched eval relies on; see the class note. A
// null return declines to the standard scheme below. ---
if (!f.nd().leaf()) {
if (ResultPtr intercepted =
try_custom_eval<EvalTrace>(f.nd(), cache)) {
finalize(finish_phase_b(f, std::move(intercepted)));
break;
}
}
// --- Leaf. ---
if (f.nd().leaf()) {
// The full leaf (traced and cached full), via the shared leaf
// fetch (fetch_leaf_traced, above).
ResultPtr result =
fetch_leaf_traced<EvalTrace>(f.nd(), leaf_evaluator, cache);
// Store the full leaf under its canonical key (a block slice would
// corrupt the cache), then return it sliced to the current block: a
// freshly built leaf's lifetime is top, so every enclosing carried
// batch loop is unbaked and is sliced here (hops == batch_context
// size). On the off path (empty batch_context) this is a no-op.
ResultPtr stored = finish_phase_b(f, std::move(result));
finalize(slice_to_use(stored, f.nd(), cache.batch_context().size()));
break;
}
// --- Internal node: request the left operand (always Checked). The
// stage must advance before the push (push may grow the deque). ---
f.stage = (f.nd()->op_type() == EvalOp::Adjoint) ? Stage::NeedLeftAdj
: Stage::NeedLeft;
stk.push_back(Frame{.node_p = &f.nd().left(), .checked = true});
break;
}
case Stage::NeedLeftAdj: {
// Unary IR op (Adjoint): only the left operand is evaluated; the right
// child is the Constant(1) sentinel kept to preserve FullBinaryNode's
// invariant, and is intentionally never pushed.
f.left = std::move(ret);
SEQUANT_ASSERT(f.left);
ResultPtr result;
auto time = detail::timed_eval_inplace(
[&]() { result = apply_one_op(f.nd(), f.left, f.right); });
if constexpr (detail::trace(EvalTrace)) {
// `right` is null here (see log::bytes() null tolerance).
size_t hwmark = log::bytes(cache, result).value;
if (!cache.chain_holds(f.left)) hwmark += log::bytes(f.left).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
log::eval(log::EvalStat{.mode = log::eval_mode(f.nd()),
.time = time,
.mem_result = log::bytes(result),
.mem_alloc = log::bytes(result),
.mem_hwmark = {cache.note_working_set(
hwmark, f.nd()->hash_value())},
.mem_left = log::bytes(f.left),
.mem_right = log::bytes(f.right)},
log::label(f.nd(), cache.batch_context()));
}
log::release_after_op();
finalize(finish_phase_b(f, std::move(result)));
break;
}
case Stage::NeedLeft: {
f.left = std::move(ret);
SEQUANT_ASSERT(f.left);
f.stage = Stage::NeedRight;
stk.push_back(Frame{.node_p = &f.nd().right(), .checked = true});
break;
}
case Stage::NeedRight: {
f.right = std::move(ret);
SEQUANT_ASSERT(f.left);
SEQUANT_ASSERT(f.right);
ResultPtr result;
// In-place accumulation is eligible only when f.left's provenance is
// known to be an evaluation-local, exclusively-owned buffer:
// - !f.nd().left().leaf(): a leaf's ResultPtr comes straight out of
// the caller-supplied leaf_evaluator, whose provenance this engine
// cannot see -- a memoizing evaluator (the norm for AO integrals /
// amplitudes reused across calls, e.g. rand_tensor_yield in the
// unit tests) can hand out the same buffer to unrelated callers,
// so mutating it here would silently corrupt those other reads.
// Only an internal node's result is guaranteed freshly built by
// this evaluation (prod()/sum()/permute() always allocate), so
// only internal nodes are safe candidates.
// - !cache.chain_holds_shared(f.left): even an internal node's result
// must not be a live, shared entry on the CacheManager's scope chain
// -- i.e. not referenced again elsewhere in this evaluation's
// tree/forest (multi-use, so some other read is pending). This is a
// runtime condition, not merely an assert: under this Release build
// SEQUANT_ASSERT expands to a no-op (SEQUANT_ASSERT_ENABLED
// undefined), so a guard that only asserted would be elided and the
// mutation would silently corrupt a value shared across roots (see
// the ordered-executor multi-root path, where a subexpression CSE'd
// across two independent roots is homed resident and read by both).
// chain_holds_shared() tests, by pointer identity, whether a live
// entry with more than one consumer (max_life > 1 -- a resident home
// or a not-yet-drained multi-use CSE entry) still holds f.left; a
// private single-use accumulator is not reported (a transient
// running total is held by no entry, and a single-use CSE entry
// moves its buffer out on its sole read, so it no longer holds it),
// so in-place still fires for the private common case -- see
// CacheManager::chain_holds_shared's own doc comment.
// - Explicit value cells: when a CellReadResolver is
// wired (the ordered executor's table-driven path), the same
// provenance question is answered from the table's own life
// instead -- eval::CellReadResolver::operand_drained(hash) is
// true iff the left operand's most recent table read spent its
// source's last life (a private/transient value, never a table
// read at all, is reported drained too: nothing else could be
// sharing it, exactly as chain_holds_shared() reports "not held"
// for the same case) -- see that method's own doc comment. The
// forest path (no resolver wired) keeps chain_holds_shared.
// A marked Sum whose left child is a leaf (the common case for the
// innermost Sum of a chain, whose left is the chain seed), or whose
// left operand is a shared cache-resident value, therefore falls back
// to the allocating sum() below despite being marked -- one bounded
// extra allocation, not per term -- and every other (non-leaf-seeded,
// private) Sum in the chain still accumulates in place from there on,
// since each already-computed running total is a fresh,
// evaluation-local buffer.
bool const inplace_eligible =
f.nd()->op_type() == EvalOp::Sum && f.nd()->accumulate_in_place() &&
!f.nd().left().leaf() && !cache.chain_holds_shared(f.left);
if (inplace_eligible) {
// The accumulate-in-place Sum (sum_in_place_traced, above),
// shared with the ordered executor's own compute_cell.
result = sum_in_place_traced<EvalTrace>(f.nd(), std::move(f.left),
f.right, cache);
finalize(finish_phase_b(f, std::move(result)));
break;
}
// The op itself, with its hook / tally / timing / trace bookkeeping
// (apply_one_op_traced, above) -- shared verbatim with the ordered
// executor's own compute_cell.
result = apply_one_op_traced<EvalTrace>(f.nd(), f.left, f.right, cache);
finalize(finish_phase_b(f, std::move(result)));
break;
}
}
}
return ret;
}
template <Trace EvalTrace = Trace::Default,
detail::CacheCheck Cache = detail::CacheCheck::Checked,
meta::can_evaluate Node, typename F, typename N, bool FHC>
requires meta::leaf_node_evaluator<Node, F>
ResultPtr evaluate(Node const& node, F const& leaf_evaluator,
CacheManager<N, FHC>& cache) {
return evaluate_impl<EvalTrace, Cache>(node, leaf_evaluator, cache);
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate Node, typename F,
typename N, bool FHC>
requires meta::leaf_node_evaluator<Node, F> //
ResultPtr evaluate(Node const& node, //
auto const& layout, //
F const& leaf_evaluator, //
CacheManager<N, FHC>& cache) {
// if the layout is not the default constructed value need to permute
bool const perm = layout != decltype(layout){};
std::string xpr;
if constexpr (detail::trace(EvalTrace)) {
xpr = toUtf8(io::serialization::to_string(to_expr(node)));
log::term(log::TermMode::Begin, xpr);
}
struct {
ResultPtr pre, post;
} result;
result.pre = evaluate_impl<EvalTrace>(node, leaf_evaluator, cache);
auto time = detail::timed_eval_inplace([&]() {
result.post = perm ? result.pre->permute(
std::array<std::any, 2>{node->annot(), layout})
: result.pre;
});
SEQUANT_ASSERT(result.post);
// logging
if constexpr (detail::trace(EvalTrace)) {
if (perm) {
// result.pre aliases a cache buffer only when the inner evaluate returned
// it unchanged (node cached at some scope, no mult_by_phase fresh alloc);
// chain_holds() tests that by pointer identity across the scope chain. A
// permuted/phase-shifted pre is a distinct buffer and is added.
size_t hwmark = log::bytes(cache, result.post).value;
if (!cache.chain_holds(result.pre))
hwmark += log::bytes(result.pre).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
auto stat = log::EvalStat{
.mode = log::EvalMode::Permute,
.time = time,
.mem_result = log::bytes(result.post),
.mem_alloc = log::bytes(result.post),
.mem_hwmark = {cache.note_working_set(hwmark, node->hash_value())}};
log::eval(stat, node->label(),
log::slice_home_annot(node, cache.batch_context()));
}
log::term(log::TermMode::End, xpr);
}
return result.post;
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate_range Nodes,
typename F, typename N, bool FHC>
requires meta::leaf_node_evaluator<std::ranges::range_value_t<Nodes>, F>
ResultPtr evaluate(Nodes const& nodes, //
auto const& layout, //
F const& leaf_evaluator, CacheManager<N, FHC>& cache) {
ResultPtr result;
for (auto&& n : nodes) {
if (!result) {
result = evaluate<EvalTrace>(n, layout, leaf_evaluator, cache);
continue;
}
ResultPtr pre = evaluate<EvalTrace>(n, layout, leaf_evaluator, cache);
auto time =
detail::timed_eval_inplace([&]() { result->add_inplace(*pre); });
// logging
if constexpr (detail::trace(EvalTrace)) {
// SumInplace allocates nothing: it writes into the accumulator. hwmark
// counts the cache plus both operands live at this moment; skip pre's
// bytes only when pre aliases a chain-resident cache buffer (fetched
// full). pre comes back from the permute-wrapping evaluate, so a
// permuted/phase-shifted read is a distinct buffer with its own pointer
// and is added; chain_holds() decides by pointer identity.
size_t hwmark = log::bytes(cache, result).value;
if (!cache.chain_holds(pre)) hwmark += log::bytes(pre).value;
hwmark += cache.parent() ? cache.parent()->chain_residency() : 0;
auto stat = log::EvalStat{
.mode = log::EvalMode::SumInplace,
.time = time,
.mem_result = log::bytes(result),
.mem_alloc = {0},
.mem_hwmark = {cache.note_working_set(hwmark, n->hash_value())}};
log::eval(stat, n->label(),
log::slice_home_annot(n, cache.batch_context()));
}
}
return result;
}
template <Trace EvalTrace = Trace::Default, meta::can_evaluate_range Nodes,
typename F, typename N, bool FHC>
requires meta::leaf_node_evaluator<std::ranges::range_value_t<Nodes>, F>
ResultPtr evaluate(Nodes const& nodes, //
F const& leaf_evaluator, CacheManager<N, FHC>& cache) {
using annot_type = decltype([](std::ranges::range_value_t<Nodes> const& n) {
return n->annot();
});
static_assert(std::is_default_constructible_v<annot_type>);
return evaluate(nodes, annot_type{}, leaf_evaluator, cache);
}
template <Trace EvalTrace = Trace::Default, typename node_t, typename F,
typename N, bool FHC>
requires meta::leaf_node_evaluator<node_t, F>
container::svector<ResultPtr> evaluate_multiroot(
container::svector<node_t> const& roots,
container::svector<std::string> const& layouts,
[[maybe_unused]] F const& leaf_evaluator, CacheManager<N, FHC>& cache) {
static_assert(
std::is_same_v<node_t, N>,
"evaluate_multiroot: the roots' node type must match the cache's key "
"type");
if (layouts.size() != roots.size())
throw Exception(
"evaluate_multiroot: layouts.size() must equal roots.size() -- one "
"layout per root is required");
auto const& drv = cache.multiroot_driver();
if (!drv)
throw Exception(
"evaluate_multiroot: no multiroot driver installed on the cache -- "
"there is no per-root fallback; install one via "
"cache.set_multiroot_driver(...) (e.g. ordered_executor.hpp's "
"evaluate_ordered_multiroot, closed over an OrderedSchedule built "
"from the SAME roots)");
return drv(roots, layouts, cache);
}
template <Trace EvalTrace = Trace::Default, typename... Args>
requires(!detail::any_type_is_cache_manager<Args...>)
ResultPtr evaluate(Args&&... args) {
using Node = std::remove_cvref_t<decltype(detail::node0(
detail::arg0(std::forward<Args>(args)...)))>;
auto cache = CacheManager<Node>::empty();
return evaluate<EvalTrace>(std::forward<Args>(args)..., cache);
}
template <Trace EvalTrace = Trace::Default, typename... Args>
requires(!detail::any_type_is_cache_manager<Args...>)
ResultPtr evaluate_impl(Args&&... args) {
using Node = std::remove_cvref_t<decltype(detail::node0(
detail::arg0(std::forward<Args>(args)...)))>;
auto cache = CacheManager<Node>::empty();
return evaluate_impl<EvalTrace>(std::forward<Args>(args)..., cache);
}
template <Trace EvalTrace = Trace::Default, typename... Args>
ResultPtr evaluate_symm(Args&&... args) {
ResultPtr pre = evaluate<EvalTrace>(std::forward<Args>(args)...);
SEQUANT_ASSERT(pre);
ResultPtr result;
auto time = detail::timed_eval_inplace([&]() { result = pre->symmetrize(); });
// logging
if constexpr (detail::trace(EvalTrace)) {
// cache is owned by the inner evaluate call and out of scope here;
// hwmark reflects only the local working set (pre + freshly allocated
// result both live during the symmetrize op).
auto stat = log::EvalStat{.mode = log::EvalMode::Symmetrize,
.time = time,
.mem_result = log::bytes(result),
.mem_alloc = log::bytes(result),
.mem_hwmark = log::bytes(pre, result)};
log::eval(
stat,
detail::node0(detail::arg0(std::forward<Args>(args)...))->label());
}
return result;
}
template <Trace EvalTrace = Trace::Default, typename... Args>
ResultPtr evaluate_antisymm(Args&&... args) {
ResultPtr pre = evaluate<EvalTrace>(std::forward<Args>(args)...);
SEQUANT_ASSERT(pre);
auto const& n0 = detail::node0(detail::arg0(std::forward<Args>(args)...));
ResultPtr result;
auto time = detail::timed_eval_inplace(
[&]() { result = pre->antisymmetrize(n0->as_tensor().bra_rank()); });
// logging
if constexpr (detail::trace(EvalTrace)) {
// See Symmetrize for the rationale on hwmark.
auto stat = log::EvalStat{.mode = log::EvalMode::Antisymmetrize,
.time = time,
.mem_result = log::bytes(result),
.mem_alloc = log::bytes(result),
.mem_hwmark = log::bytes(pre, result)};
log::eval(stat, n0->label());
}
return result;
}
struct accept_any_index {
bool operator()(Index const&) const noexcept { return true; }
};
struct no_scope_guard {};
struct make_no_scope_guard {
no_scope_guard operator()(std::size_t /*n_batches*/) const noexcept {
return {};
}
};
struct never_volatile {
template <typename Node>
bool operator()(Node const&) const noexcept {
return false;
}
};
template <typename Node, typename Pred>
[[nodiscard]] bool subtree_any(Node const& n, Pred const& pred) {
// Small-buffer stack: the hot callers are per-node / per-value
// (place_at_this_level's `collect` asks this of every non-leaf node of every
// member root, and the ordered path's volatile_of asks it once per value),
// and the recursion this replaced allocated nothing -- so the pending
// frontier, which for anything but a deep spine is a handful of pointers,
// stays on the stack. It still grows onto the heap for the spine case, which
// is the point.
container::svector<Node const*, 32> stack{&n};
while (!stack.empty()) {
Node const& c = *stack.back();
stack.pop_back();
if (pred(c)) return true;
if (!c.leaf()) {
stack.push_back(&c.right());
stack.push_back(&c.left());
}
}
return false;
}
namespace detail {
template <typename TreeNode, bool FHC>
struct BatchedScratch {
CacheManager<TreeNode, FHC> cache;
std::vector<TreeNode const*> seeds;
};
template <typename TreeNode, bool FHC, typename Members>
[[nodiscard]] BatchedScratch<TreeNode, FHC> make_batched_scratch(
Members const& members, CacheManager<TreeNode, FHC> const& real) {
using Hasher = TreeNodeHasher<TreeNode, FHC>;
using Comp = TreeNodeEqualityComparator<TreeNode>;
// The batch's external modes: obtained exactly as the evaluator obtains them
// (partition each member root's node_slice_mask() by BatchModeType). An
// External mode is an external that survives free onto a node's result, so a
// node carrying one is not batch-invariant under that mode -- its value
// depends on the external slice. When the caller nests an External mode
// outside a Contracted one (the External mode is sliced by an outer re-entry,
// then this scratch batches an inner Contracted mode), a persistent
// intermediate that carries the External mode but not the Contracted `mode`
// would look seedable under `mode` alone -- yet seeding its full
// (unsliced-external) value would be wrong under the outer slice. Tracking
// the External modes in the signature (below) forbids seeding/sharing such
// nodes. When there is no External mode this list is empty and every
// External-derived test is a no-op, leaving the Contracted-only behavior
// untouched.
container::svector<Index> ext_axes;
for (auto const& [root, mode] : members) {
if (root->leaf()) continue;
for (auto const& [ix, knd] : (*root)->node_slice_mask())
if (knd == BatchModeType::External &&
std::find(ext_axes.begin(), ext_axes.end(), ix) == ext_axes.end())
ext_axes.push_back(ix);
}
struct Meta {
std::size_t count = 0;
std::optional<std::size_t> sig;
// Positions of each External mode (in ext_axes order) in this node's
// canonical result indices, or nullopt if the node does not carry it. This
// is a function of the (canonical) node alone, so it is identical across
// all occurrences of a canonically-equal node.
container::svector<std::optional<std::size_t>> ext_sig;
bool consistent = true;
};
std::unordered_map<TreeNode const*, Meta, Hasher, Comp> meta;
// The external-mode signature (positions of each ext axis on n's result, or
// absent) is exactly slicing_signature(n, ext_axes); the batch-mode `sig`
// below is the single-mode case, index_position(n, mode).
auto ext_sig_of = [&ext_axes](TreeNode const& n) {
return slicing_signature(n, ext_axes);
};
// Iterative pre-order (an explicit stack, not recursion): a member root can
// be the node `evaluate` was called on -- a forest root, i.e. the residual's
// single in-place Sum tree, whose left spine is as deep as the term count --
// so a recursive descent here is O(spine) in call frames. Pushing the right
// child before the left one keeps the pop order the recursion's pre-order, so
// which occurrence is the `first` one (and so every recorded signature and
// count) is unchanged.
auto visit = [&meta, &ext_sig_of](TreeNode const& root,
Index const& mode) -> void {
std::vector<TreeNode const*> stack{&root};
while (!stack.empty()) {
TreeNode const& n = *stack.back();
stack.pop_back();
if (n.leaf()) continue;
auto const sig = index_position(n, mode);
auto const esig = ext_sig_of(n);
auto const [it, first] = meta.try_emplace(&n);
auto& e = it->second;
if (first) {
e.sig = sig;
e.ext_sig = esig;
} else if (e.sig != sig || e.ext_sig != esig) {
e.consistent = false;
}
++e.count;
// Prune a re-encounter only when its signature matches the first one:
// canonical equality maps canonical position p to position p, so an
// equal signature here implies the descendants' signatures equal those
// already recorded on the first walk (deeper accesses shared and
// counted). A differing signature gives no such guarantee -- descend so
// descendants' signatures under this occurrence are recorded too;
// otherwise a descendant sliced differently only under this (unshared,
// pruned) occurrence could pass the guard and serve wrong slices. The
// extra descendant counts are real accesses: an inconsistently-sliced
// occurrence is evaluated per occurrence, not served from the scratch at
// n. The External signature is invariant across occurrences (a function
// of the canonical node), so folding it into the match only tightens the
// guard.
if (!first && e.sig == sig && e.ext_sig == esig) continue;
stack.push_back(&n.right());
stack.push_back(&n.left());
}
};
for (auto const& [root, mode] : members) {
if (root->leaf()) continue;
// member roots are accumulated by the caller, not cached here.
visit(root->left(), mode);
visit(root->right(), mode);
}
std::unordered_map<TreeNode, std::size_t, Hasher, Comp> reg;
std::unordered_set<TreeNode, Hasher, Comp> seed_keys;
std::vector<TreeNode const*> seeds;
for (auto const& [ptr, e] : meta) {
if (!e.consistent) continue; // ambiguous slicing: never share
// A node carrying any batched External mode has an external slice a
// seeded/home-read full value would ignore -- so it is never
// shareable-full.
bool const carries_ext =
std::any_of(e.ext_sig.begin(), e.ext_sig.end(),
[](auto const& p) { return p.has_value(); });
// Seed an alive persistent batch-invariant real entry into the scratch
// (persistent so it survives reset()), else register a repeated subnode.
bool const seedable =
!e.sig && !carries_ext && real.persistent(*ptr) && real.alive(*ptr);
if (seedable) {
seeds.push_back(ptr);
seed_keys.insert(*ptr);
reg.emplace(*ptr, e.count); // count ignored for persistent entries
} else if (e.count >= 2) {
reg.emplace(*ptr, e.count);
}
}
auto is_persistent = [seed_keys = std::move(seed_keys)](TreeNode const& n) {
return seed_keys.contains(n);
};
CacheManager<TreeNode, FHC> scratch{std::move(reg), std::move(is_persistent)};
return {std::move(scratch), std::move(seeds)};
}
} // namespace detail
using PeakSink = std::atomic<double>*;
template <Trace EvalTrace = Trace::Default, typename F,
typename IndexPredicate = accept_any_index,
typename ScopeGuardFactory = make_no_scope_guard,
typename IsVolatile = never_volatile>
[[nodiscard]] auto make_batched_custom_evaluator(
F leaf_evaluator,
std::function<std::size_t(Index const&)> target_batch_size,
IndexPredicate accept = {}, ScopeGuardFactory make_scope_guard = {},
IsVolatile is_volatile = {}, bool persistent_only = false,
std::size_t depth = 0, PeakSink peak = nullptr) {
return [leaf_evaluator = std::move(leaf_evaluator),
target_batch_size = std::move(target_batch_size), accept, is_volatile,
persistent_only, depth, peak,
make_scope_guard](auto const& node, auto& cache) -> ResultPtr {
// Runaway backstop: nesting re-enters this evaluator on the per-batch
// scratch (see the reinstall below), incrementing depth once per nested
// mode. Real trees nest a handful of modes deep; a large depth signals a
// non-terminating re-entry (e.g. a mis-annotated mode that never shrinks).
SEQUANT_ASSERT(depth < 8);
// Slice a value to the current batch block for the `hops` innermost
// enclosing batch loops that `nd` carries (see the Enter-stage
// slice_to_use in evaluate() -- this is the same primitive, reachable from
// the closure-internal probes below that bypass the Enter stage). `cache`
// is the cache the closure fired on, so cache.batch_context() is this
// node's enclosing context.
auto slice_to_use = [&cache](ResultPtr value, auto const& nd,
std::size_t hops) -> ResultPtr {
auto const& ctx = cache.batch_context();
std::size_t const d = ctx.size();
// See the assert on the evaluate() copy of this lambda: hops <= d always
// (one batch_context push + <=1 parent link per level); a violation would
// underflow `d - hops` and silently under-slice.
SEQUANT_ASSERT(hops <= d);
for (std::size_t i = d - hops; i < d; ++i) {
auto const& axis = ctx[i].axis;
auto const& blk = ctx[i].range;
if (auto const p = index_position(nd, axis))
value = value->slice_mode(*p, blk.first, blk.second);
}
return value;
};
// A leaf evaluator that slices each fetched leaf to the enclosing batch
// blocks, over the whole enclosing nest. Used at the three closure-internal
// sites (pick_sliceable probe, carrier_full pre-size, hoist build) that
// consume leaf slicing but bypass the Enter stage; feeding the raw
// leaf_evaluator there would un-slice the enclosing loops and break the
// "K is not re-picked" invariant (a re-entered probe would see K's full
// extent and re-batch it). The main value path does not use this: it re-
// enters evaluate() with the raw leaf_evaluator and lets the Enter-stage
// slice_to_use do the slicing.
auto sliced_leaf = [&](auto const& ln) -> ResultPtr {
return slice_to_use(leaf_evaluator(ln), ln, cache.batch_context().size());
};
// Synthesized DAG-scope level for a forest-evaluator push: this firing
// realizes exactly one loop over `cache`'s own enclosing context, so every
// push site below shares the same depth (`cache.batch_context().size() +
// 1`, matching build_ordered_schedule's `d + 1` convention -- see
// DagScopeLevel's doc comment) and differs only in the pushed axis's
// space. Plumbing only: nothing resolves by `level` on this path -- the
// forest evaluator resolves by exact axis (`exact_axis`, filled at each
// push site below).
auto const synth_level = [&cache](Index const& ax) -> DagScopeLevel {
return DagScopeLevel{.depth = cache.batch_context().size() + 1,
.space = std::wstring(ax.space().base_key())};
};
// Mode selection is sliceability-aware and realizes the optimizer's
// multi-mode nesting one mode per depth level. candidate_axes lists this
// node's batch modes in the optimizer's annotated order (see
// EvalExpr::node_slice_mask), keeping the accepted annotations.
//
// The optimizer's annotations are authoritative at every depth: a node
// carrying no accepted annotation means "do not batch this node", and is
// left unbatched. There is deliberately no heuristic fallback -- batching
// is only ever realized where the peak-constrained optimizer asked for it,
// so every realized batch loop is one the cost model priced. (Callers
// therefore cannot batch without a peak budget: no budget => the optimizer
// emits no annotations => nothing batches. See BatchPolicy::peak_threshold
// and, on the MPQC side, validate_batch_config.)
auto candidate_axes =
[&accept](auto const& n) -> container::svector<Index> {
container::svector<Index> out;
for (auto const& entry : n->node_slice_mask())
if (accept(entry.first)) out.push_back(entry.first);
return out;
};
// Pick the first candidate mode that is actually sliceable (partitions into
// > 1 batch) in this (possibly already-outer-sliced) context, returning it
// together with its realized partition. A mode already sliced by an outer
// re-entry yields a single batch on the sliced leaf and is skipped, so a
// nested re-entry on the same node advances to the node's next annotated
// mode -- realizing `for K-batch: for mu1-batch: replay` at one multi-mode
// node. The recursive reinstall below walks one mode per depth level (the
// depth < 8 backstop bounds the re-entry).
auto pick_sliceable = [&](auto const& n)
-> std::optional<std::pair<
Index, container::svector<std::pair<std::size_t, std::size_t>>>> {
BackendArrayOps const* const aops = cache.array_ops();
auto const& ectx = cache.batch_context();
// Is ix's mode already sliced by an enclosing block, so it must not be
// re-picked? Mirrors slice_to_use exactly: an exact context axis slices
// ix.
auto already_sliced = [&](Index const& ix) {
for (auto const& e : ectx)
if (e.axis == ix) return true;
return false;
};
for (Index const& ix : candidate_axes(n)) {
if (already_sliced(ix)) continue;
SEQUANT_ASSERT(aops &&
"batched forest eval requires backend array-ops "
"(CacheManager::set_array_ops)");
auto b = aops->axis_batches(ix, target_batch_size(ix));
if (b.size() > 1) return std::make_pair(ix, std::move(b));
}
return std::nullopt;
};
// Persistence gate (opt-in via persistent_only): when set, decline to batch
// any subtree containing a volatile leaf -- such a subtree is rebuilt every
// evaluation, so batching pays the partition + relaxed-screening cost each
// pass to amortize over nothing. By default (persistent_only == false) we
// batch across the board: slicing the batch mode reduces the footprint of
// any mode-carrying intermediate regardless of volatility, and the cost
// model credits it accordingly, so the runtime must realize it too. (When
// is_volatile is never_volatile the gate is moot either way.)
if (persistent_only && subtree_any(node, is_volatile)) {
return nullptr;
}
auto picked = pick_sliceable(node);
if (!picked) {
return nullptr; // no accepted, sliceable mode (nothing to gain)
}
Index const K = std::move(picked->first);
auto const batches = std::move(picked->second);
using node_t = std::remove_cvref_t<decltype(node)>;
using member_t = std::pair<node_t const*, Index>;
TreeNodeEqualityComparator<node_t> const eq;
// Classify the picked mode by BatchModeType. K is a batch mode of `node`;
// the optimizer stamps it contracted (summed away -> block partials
// accumulate) or external (an external index free on the node's result ->
// block partials are disjoint slices, scattered into a pre-sized result).
// The depth-0 heuristic fallback only ever yields a contracted index, so an
// mode absent from node_slice_mask() is Contracted, so an unannotated
// node takes the Contracted-only path with no External entry.
BatchModeType picked_kind = BatchModeType::Contracted;
for (auto const& [ix, knd] : node->node_slice_mask())
if (ix == K) {
picked_kind = knd;
break;
}
// Per-level placement (order-aware only). This
// firing realizes a batch loop over `K` at runtime `depth`. A
// member-subtree node invariant to this loop (it does not carry `K` on its
// result) is built once at its home level and served to every batch body
// through the scope chain, rather than rebuilt per batch. A node's
// residency (the batch modes it is variant to) is its \c sliced_modes():
// the cross-occurrence lifetime-mask meet of all batched modes -- External
// (occ) and Contracted (aux) alike -- that live on the node's own result
// slots (consistent placement across occurrences -> CSE). A node variant to
// an outer aux loop carries that aux free on a result slot, so the aux mode
// survives the meet into sliced_modes.
// A node is invariant to this loop iff `K` is not in its residency. Its
// home level is the deepest enclosing batch_context entry whose mode is in
// the residency (the innermost enclosing loop it is variant to); -1 (the
// chain root / run-term cache) if it is invariant to the whole nest. The
// node is built once at that level (walk-up), sliced to its home blocks,
// and reused across this loop's batches. A node carrying `K` (K in its
// residency) is loop-local: it is left to inline evaluation (descend),
// which finds any deeper-hoisted invariants through the chain -- so
// descending never rebuilds them.
//
// The walk-up decides placement from the residency-derived home level, and
// hoists a node iff `K` is not in its sliced_modes (an External carrier has
// K in sliced_modes -> loop-local -> descended, never hoisted).
// The order-aware gate is the emitted `batch_order_aware()` bit (true for
// every node the order-aware cost model emitted, including a whole-nest
// invariant whose residency is empty): on the off path every node is
// order-blind, so `targets` is empty, set_parent is not wired, and the
// per-batch replay is left untouched. The bit is a positive signal an
// empty residency cannot provide -- it is what distinguishes an off-path
// all-full node (do not hoist) from an order-aware whole-nest invariant
// (hoist to the root).
auto place_at_this_level =
[&](auto& scratch_cache, auto& parent_cache,
std::vector<node_t const*> const& member_roots) {
// The enclosing batch loops (strictly outer to this firing); this
// level's mode K and any inner loop are not in it. A node is
// hoistable here iff every residency (sliced_modes) mode is one of
// these outer loops: then it is invariant to this loop and to every
// inner loop, so its home is its deepest enclosing residency level
// and it is built once there. If a residency mode is K (loop-local)
// or an inner mode (its home is a deeper loop), it is not all-outer
// -> descend, so the deeper level handles it sliced. The deepest
// residency being outer <=> all residency outer (deepest is the max);
// a node carrying an inner mode is deliberately not hoisted at this
// outer level, since hoisting it would hold an aux carrier full over
// aux at the outer occ level.
auto const& ectx = parent_cache.batch_context(); // enclosing loops
auto in_ectx = [&ectx](Index const& m) -> bool {
for (auto const& e : ectx)
if (e.axis == m) return true;
return false;
};
auto residency_all_outer = [&in_ectx](node_t const& n) -> bool {
for (auto const& ix : n->sliced_modes())
if (!in_ectx(ix)) return false;
return true;
};
// Placement here is purely the seed, so an order-aware,
// residency-all-outer node is hoisted to its seed home (full on any
// demoted mode), including a node carrying an external
// node_slice_mask() stamp absent from its sliced_modes. A value
// cached at its seed home is the same value the descended path
// produces -- the Enter-stage slice-on-use slices it to the block
// when a nested external loop consumes it -- so this is a placement
// choice only, never a change of result. The table-driven ordered
// executor makes its own placement decisions on cells.
std::vector<node_t const*> targets;
// Iterative pre-order (explicit stack): a member root can be the node
// `evaluate` was called on -- a forest root, i.e. the residual's
// single in-place Sum tree, whose left spine is as deep as the term
// count -- so a recursive descent here is O(spine) in call frames.
// Right pushed before left, so the pop order is the recursion's
// pre-order and `targets` comes out in the same order.
auto collect = [&](node_t const& root) -> void {
std::vector<node_t const*> stack{&root};
while (!stack.empty()) {
node_t const& n = *stack.back();
stack.pop_back();
if (n.leaf()) continue;
if (n->batch_order_aware() && residency_all_outer(n) &&
!subtree_any(n, is_volatile)) {
auto shares = [&](node_t const* p) { return eq(*p, n); };
if (std::none_of(targets.begin(), targets.end(), shares))
targets.push_back(&n);
continue; // built as a unit -- do not descend into it
}
stack.push_back(&n.right());
stack.push_back(&n.left());
}
};
for (node_t const* m : member_roots) {
if (m->leaf()) continue;
collect(m->left());
collect(m->right());
}
if (targets.empty()) return;
// Wire the scope chain only when there is something to hoist, which
// keeps the off path unwired.
scratch_cache.set_parent(&parent_cache);
auto in_residency = [](node_t const& n, Index const& m) -> bool {
auto const& sm = n->sliced_modes();
return std::find(sm.begin(), sm.end(), m) != sm.end();
};
for (node_t const* dptr : targets) {
node_t const& d = *dptr;
// Home level = deepest enclosing-context entry whose mode is in d's
// residency (sliced_modes); -1 => invariant to the whole nest
// (chain root).
int rl = -1;
for (int i = static_cast<int>(ectx.size()) - 1; i >= 0; --i)
if (in_residency(d, ectx[i].axis)) {
rl = i;
break;
}
// Locate the level-rl cache by walking up from parent_cache (the
// level depth-1 cache): rl == -1 => the chain root (the real/term
// cache); rl >= 0 => the scratch (depth-1 - rl) hops up. Runtime
// nest depth aligns with the scope-chain position (each realized
// loop = one context entry = one scratch level).
auto* target = &parent_cache;
if (rl == -1) {
while (target->parent()) target = target->parent();
} else {
// Release-safe guard (SEQUANT_ASSERT elides in release): never
// walk the chain off its end and dereference a null parent.
if (rl > static_cast<int>(depth) - 1)
throw Exception(
"hoist home level not strictly outer to this loop");
for (int lvl = static_cast<int>(depth) - 1; lvl > rl; --lvl) {
auto* const p = target->parent();
if (!p)
throw Exception(
"hoist walk-up exceeded the scope chain (a single-batch "
"sliced mode may have shifted the runtime nest depth)");
target = p;
}
}
target->ensure_hoist_slot(d);
if (target->alive(d)) continue; // built already in a broader scope
// Build the whole invariant once on a fresh cache via the variadic
// evaluate(n, sliced_leaf) (empty cache, no custom evaluator, so no
// re-entry into this batched evaluator). `sliced_leaf` slices the
// enclosing loops d carries (up to its home level); the loops it
// does not carry pass through unsliced (built full over its deeper
// / invariant modes). Store under the same canonical-phase
// convention the batched member store uses.
ResultPtr built = evaluate_impl<EvalTrace>(d, sliced_leaf);
if (auto const ph = d->canon_phase(); ph != 1)
built = built->mult_by_phase(ph);
(void)target->store_and_access(d, std::move(built));
}
};
// DEBUG (behavior-neutral): log the trigger's depth, picked mode + kind,
// and its full node_slice_mask() annotation + result indices, to diagnose
// nested re-batching of a single aux mode. Emitted only when tracing is on.
if (log::printing()) {
std::string annot;
for (auto const& [ix, knd] : node->node_slice_mask()) {
annot += toUtf8(ix.full_label());
annot += (knd == BatchModeType::External ? ":ext " : ":con ");
}
std::string res;
for (auto const& ix : node->canon_indices()) {
res += toUtf8(ix.full_label());
res += " ";
}
auto scope_ctx = cache.batch_context();
scope_ctx.push_back({K, synth_level(K), {0, 0}, K});
log::log(
"BatchAxes",
std::format("depth={} picked={}:{} nbatches={} annot=[{}] "
"result=[{}] {}",
depth, toUtf8(K.full_label()),
picked_kind == BatchModeType::External ? "ext" : "con",
batches.size(), annot, res, log::scope_annot(scope_ctx)));
}
if (picked_kind == BatchModeType::External) {
// Scatter branch. K survives to node's result as a free external mode,
// so the per-block partials are disjoint slices of one result (not
// summands of a contraction): they are write_into_slice()d into a
// pre-sized result, never add_inplace()d. Inner batch modes -- of either
// kind, at this node or a descendant -- still nest through the same
// per-block reinstall the contracted path uses, so External composes with
// Contracted: within each external block the reinstalled evaluator slices
// any remaining mode (a contracted mode accumulates within the block, an
// inner external mode scatters into its own block-local result). K itself
// is not re-picked on the re-entry: the pushed batch_context makes the
// re-entry's sliced_leaf slice K's carrier to this block, so the block
// yields a single batch on the sliced leaf and is skipped (same invariant
// as the contracted nesting).
auto const dest_mode = index_position(node, K);
SEQUANT_ASSERT(dest_mode &&
"external batch mode is not free on the node's result");
// Backend array-ops from the cache chain -- the same source the ordered
// (DAG) executor reads, so forest and DAG build identical scatter
// destinations from the node's own (unsliced) index list.
BackendArrayOps const* const aops = cache.array_ops();
SEQUANT_ASSERT(aops &&
"batched external-mode scatter requires backend array-ops "
"(CacheManager::set_array_ops)");
// A single-node scratch: an external mode is not a
// persistent-final sharing mode, so the group/replay machinery (which
// co-batches cross-term contracted finals) does not apply -- scatter just
// this node. The scratch still dedups repeats within the node's subtree.
std::vector<member_t> solo{{&node, K}};
auto bs = detail::make_batched_scratch(solo, cache);
bs.cache.set_array_ops(cache.array_ops()); // inherit backend ops
for (auto const* s : bs.seeds)
(void)bs.cache.store_and_access(*s, cache.access(*s));
place_at_this_level(bs.cache, cache, std::vector<node_t const*>{&node});
auto const scope_guard = make_scope_guard(batches.size());
(void)scope_guard;
if (log::printing())
log::log("BatchScatter", "Begin",
std::format("external mode over {} blocks", batches.size()));
ResultPtr dest;
for (auto const& [e_lo, e_hi] : batches) {
if (e_lo == e_hi) continue;
// Per-block slice marker for trace post-processing (avoidable-recompute
// accounting), mirroring the contracted BatchGroup path: tags every op
// replayed in this external block with the enclosing scatter loop's
// (mode, element-range-low). Without it, the same expression scattered
// into disjoint external blocks would carry an identical touched-slice
// signature and be miscounted as a duplicate rebuild, when each block
// is in fact legitimate (1/nblocks of the work). Gated on printing() so
// it is inert unless a trace is being emitted.
if (log::printing())
log::log("BatchIter", toUtf8(std::wstring(K.full_label())), e_lo);
bs.cache.reset();
// Extend the enclosing batch context by this block and set it on the
// scratch, so the re-entry's Enter-stage slice-on-use (and its own
// sliced_leaf) slices every leaf carrying K to this block and composes
// inner slices on top. Slice-on-use also covers a cached intermediate
// fetched from an ancestor scope. The raw leaf_evaluator
// is threaded down; the Enter stage does the slicing.
auto ctx = cache.batch_context();
ctx.push_back({K, synth_level(K), {e_lo, e_hi}, K});
bs.cache.set_batch_context(std::move(ctx));
bs.cache.set_custom_evaluator(make_batched_custom_evaluator<EvalTrace>(
std::function<ResultPtr(node_t const&)>{leaf_evaluator},
target_batch_size, accept, make_scope_guard, is_volatile,
persistent_only, depth + 1, peak));
ResultPtr part =
evaluate_impl<EvalTrace>(node, leaf_evaluator, bs.cache);
// Pre-size the full-extent zero destination from the node's own
// (unsliced) index list on the first block; the backend realizes it
// (flat or nested) with no array in the DAG consulted.
if (!dest) dest = aops->make_zeros(node->canon_indices());
dest->write_into_slice(*part, *dest_mode, e_lo, e_hi);
if (peak) {
const double cand =
static_cast<double>(bs.cache.working_set_hwmark());
double cur = peak->load(std::memory_order_relaxed);
while (cand > cur && !peak->compare_exchange_weak(
cur, cand, std::memory_order_relaxed)) {
}
}
}
if (log::printing()) log::log("BatchScatter", "End");
SEQUANT_ASSERT(dest);
return dest;
}
// The replay group: the trigger plus every registered persistent key that
// is not yet alive and batches over a mode with the identical realized
// partition. All compatible persistent finals stream over the batch mode
// in the same passes, so sub-intermediates shared between them (wherever
// the scratch's slicing-signature guard admits sharing -- equal canonical
// positions of the batch mode plus equal element ranges imply identical
// slices) are evaluated once per batch instead of once per consumer.
// The cost of considering a candidate is one leaf evaluation (the
// mode_batches probe). With an unregistered (empty) real cache the group
// is just the trigger.
std::vector<member_t> group{{&node, K}};
cache.for_each_key([&](node_t const& k) {
if (!cache.persistent(k) || cache.alive(k)) return;
if (eq(k, node)) return; // the trigger occupies its own slot
if (subtree_any(k, is_volatile)) return; // defensive: P implies NV
auto const pk = pick_sliceable(k);
if (!pk) return;
// Join iff this member's first sliceable mode realizes the identical
// partition as the trigger (so all members stream over the same batches).
if (pk->second != batches) return;
group.emplace_back(&k, pk->first);
});
// Layer by nesting: a member whose subtree contains another member
// evaluates in a later layer, with the inner result by then alive in the
// real cache -- seeded into the outer pass when slice-free w.r.t. the
// outer batch mode, re-derived sliced (correct, unshared) otherwise.
// Iterative (explicit stack), for the same reason as the walks above: a
// group member can be the node `evaluate` was called on, i.e. a forest root
// whose left spine is as deep as the term count. Right pushed before left
// keeps the pop order the recursion's pre-order, so the same node is found
// first (and the answer, a bool, is order-independent anyway).
auto contains = [&eq](node_t const& outer, node_t const& inner) -> bool {
if (outer.leaf()) return false;
std::vector<node_t const*> stack{&outer.right(), &outer.left()};
while (!stack.empty()) {
node_t const& n = *stack.back();
stack.pop_back();
if (eq(n, inner)) return true;
if (!n.leaf()) {
stack.push_back(&n.right());
stack.push_back(&n.left());
}
}
return false;
};
std::vector<std::vector<member_t>> layers;
{
std::vector<member_t> remaining = std::move(group);
while (!remaining.empty()) {
std::vector<member_t> layer, rest;
for (auto const& m : remaining) {
bool const outer = std::any_of(
remaining.begin(), remaining.end(), [&](member_t const& o) {
return m.first != o.first && contains(*m.first, *o.first);
});
(outer ? rest : layer).push_back(m);
}
SEQUANT_ASSERT(!layer.empty()); // containment is a strict order
layers.push_back(std::move(layer));
remaining = std::move(rest);
}
}
// Trace: the batched path co-evaluates a GROUP -- the trigger plus any
// cross-term persistent finals that slice over the same aux partition --
// streaming them together over the aux batches in one pass (so a
// sub-intermediate shared between members is computed once per batch, not
// once per consumer). The members are SIBLINGS computed alongside each
// other, NOT a term hierarchy; the per-op Eval lines below interleave
// across members and batches. Bracket the group and list its members so
// those ops can be attributed. Distinct "BatchGroup"/"BatchMember" labels
// (not Term|Begin/End) to avoid implying nesting; the top-level evaluate
// still emits the enclosing per-term Term markers.
if (log::printing()) {
std::size_t n_members = 0;
for (auto const& layer : layers) n_members += layer.size();
auto scope_ctx = cache.batch_context();
scope_ctx.push_back({K, synth_level(K), {0, 0}, K});
log::log(
"BatchGroup", "Begin",
std::format("{} members co-evaluated over {} aux batches {}",
n_members, batches.size(), log::scope_annot(scope_ctx)));
for (auto const& layer : layers)
for (auto const& mk : layer)
log::log("BatchMember",
toUtf8(io::serialization::to_string(to_expr(*mk.first))));
}
{
// Structured BatchGroup for the visualizer: the co-evaluation unit (its
// batch mode, block count, and member node hashes) so the DAG can draw a
// subgraph enclosing the siblings streamed together over K -- the runtime
// batching structure the IR forest cannot show. Gated by
// SEQUANT_SCHED_DUMP.
static bool const sched_dump =
eval::detail::dump_enabled("SEQUANT_SCHED_DUMP");
if (sched_dump) {
std::cerr << "SCHEDULE_RUN_GROUP {\"kind\":\""
<< (picked_kind == BatchModeType::External ? "external"
: "contracted")
<< "\",\"mode\":\"" << toUtf8(K.full_label())
<< "\",\"blocks\":" << batches.size() << ",\"members\":[";
bool gfirst = true;
for (auto const& layer : layers)
for (auto const& mk : layer) {
std::cerr << (gfirst ? "" : ",") << '"' << (*mk.first)->hash_value()
<< '"';
gfirst = false;
}
std::cerr << "]}\n";
}
}
// RAII scope for the batched partial contractions; a backend-supplied
// factory may relax block-sparse screening here (scaled by the batch count)
// so per-batch screening does not drop contributions that survive over the
// full batch mode. Held for the entire loop below, including the per-batch
// evaluate() calls that may re-enter this evaluator on an inner annotated
// node (see the reinstall's `make_scope_guard` argument): the inner
// level's own guard is then constructed and destroyed while this (outer)
// guard is still alive, so a backend that relaxes screening scaled by its
// own level's batch count composes multiplicatively across nesting depth
// (net relaxation = product of batch counts over all alive levels).
auto const scope_guard = make_scope_guard(batches.size());
(void)scope_guard;
ResultPtr trigger_result;
for (auto const& layer : layers) {
// The layer's scratch cache: registered from the member subtrees (same
// canonical-equality counting as the real cache), so repeated subtrees
// -- canonically-equal siblings within a member as well as
// sub-intermediates shared between members -- are evaluated once per
// batch. Carries no custom evaluator (no re-interception) and keeps the
// partial, sliced intermediates out of the real cache; reset() between
// batches drops the previous batch's partials, while pre-seeded alive
// persistent entries (registered persistent in the scratch) survive.
auto bs = detail::make_batched_scratch(layer, cache);
bs.cache.set_array_ops(cache.array_ops()); // inherit backend ops
for (auto const* s : bs.seeds)
(void)bs.cache.store_and_access(*s, cache.access(*s));
{
std::vector<node_t const*> roots;
roots.reserve(layer.size());
for (auto const& mk : layer) roots.push_back(mk.first);
place_at_this_level(bs.cache, cache, roots);
}
std::vector<ResultPtr> acc(layer.size());
for (auto const& [e_lo, e_hi] : batches) {
if (e_lo == e_hi) continue;
// Per-slice marker for trace post-processing (avoidable-recompute
// accounting): tags every op replayed in this iteration with the
// enclosing batch loop's (mode, element-range-low). The mode label K
// and the range low bound uniquely identify this slice within the
// (single-mode-per-level) nesting. Gated on printing() so it is inert
// unless a trace is being emitted.
if (log::printing())
log::log("BatchIter", toUtf8(std::wstring(K.full_label())), e_lo);
bs.cache.reset();
for (std::size_t m = 0; m != layer.size(); ++m) {
auto const& [mem, Km] = layer[m];
// Extend the enclosing batch context by this member's block and set
// it on the scratch, so the re-entry's Enter-stage slice-on-use (and
// its own sliced_leaf) slices every leaf carrying Km to this block
// and composes inner slices on top. Slice-on-use also covers a
// cached intermediate fetched from an ancestor scope.
// Rebuilt from `cache.batch_context()` (the enclosing context) each
// member so contexts do not accumulate across members. The raw
// leaf_evaluator is threaded down (type-erased into a std::function
// so this template's self-instantiation is finite: the leaf-evaluator
// type is std::function at every deeper level); the Enter stage does
// the slicing.
auto ctx = cache.batch_context();
ctx.push_back({Km, synth_level(Km), {e_lo, e_hi}, Km});
bs.cache.set_batch_context(std::move(ctx));
bs.cache.set_custom_evaluator(
make_batched_custom_evaluator<EvalTrace>(
std::function<ResultPtr(node_t const&)>{leaf_evaluator},
target_batch_size, accept, make_scope_guard, is_volatile,
persistent_only, depth + 1, peak));
ResultPtr part =
evaluate_impl<EvalTrace>(*mem, leaf_evaluator, bs.cache);
if (!acc[m])
acc[m] = std::move(part);
else
acc[m]->add_inplace(*part);
}
// Fold this batch's scratch high-watermark into the global sink. The
// next iteration calls bs.cache.reset(), which zeroes the scratch
// hwmark, so the fold must happen here (per batch), not after the
// batches loop. The loop is serial (the nested evaluate() re-entry is
// serial too), but the sink is atomic; a relaxed fetch-max CAS keeps it
// correct regardless. A null sink skips the fold entirely, leaving
// existing (sink-less) callers byte-unchanged.
if (peak) {
const double cand =
static_cast<double>(bs.cache.working_set_hwmark());
double cur = peak->load(std::memory_order_relaxed);
while (cand > cur && !peak->compare_exchange_weak(
cur, cand, std::memory_order_relaxed)) {
}
}
}
// Store the members into the real cache under the canonical-phase
// convention (mirroring evaluate()'s Checked store), eagerly per layer
// so later layers can seed them. The trigger is returned instead: its
// Checked wrapper stores it (a direct store here would double-decay a
// non-persistent trigger's life count).
for (std::size_t m = 0; m != layer.size(); ++m) {
auto const* mem = layer[m].first;
if (mem == &node) {
trigger_result = std::move(acc[m]);
continue;
}
ResultPtr v = std::move(acc[m]);
if (auto const ph = (*mem)->canon_phase(); ph != 1)
v = v->mult_by_phase(ph);
(void)cache.store_and_access(*mem, std::move(v));
}
}
if (log::printing()) {
auto scope_ctx = cache.batch_context();
scope_ctx.push_back({K, synth_level(K), {0, 0}, K});
log::log("BatchGroup", "End", log::scope_annot(scope_ctx));
}
SEQUANT_ASSERT(trigger_result);
return trigger_result;
};
}
template <Trace EvalTrace = Trace::Default, class F,
class ScopeGuardFactory = make_no_scope_guard>
[[nodiscard]] auto make_evaluator(BatchPolicy const& policy, F yielder,
ScopeGuardFactory make_scope_guard = {},
PeakSink peak = nullptr) {
auto is_volatile_node = [p = policy.is_volatile_leaf](auto const& n) -> bool {
if (!n.leaf() || !n->is_tensor()) return false;
return p && p(n->as_tensor());
};
// BatchPolicy docs: an empty is_batchable_index or batch_target_size means
// "no batching". Forwarding an empty std::function would instead throw
// std::bad_function_call from batch_axis()/target_batch_size() at evaluation
// time, so when either is unset, substitute predicates that decline batching
// (accept nothing => batch_axis returns nullopt => target_batch_size is never
// called) rather than partially-filled ones.
// Runtime accept = the derived union of both batchability roles: a mode is
// accepted at runtime if it is batchable in either the contracted or the
// external role (see BatchPolicy::is_batchable_index()).
std::function<bool(Index const&)> accept = policy.is_batchable_index();
std::function<std::size_t(Index const&)> target = policy.batch_target_size;
if (!accept || !target) {
accept = [](Index const&) { return false; };
target = [](Index const&) -> std::size_t { return 0; };
}
return make_batched_custom_evaluator<EvalTrace>(
std::move(yielder), std::move(target), std::move(accept),
std::move(make_scope_guard), std::move(is_volatile_node),
policy.persistent_only, /*depth=*/0, peak);
}
} // namespace sequant
#endif // SEQUANT_EVAL_EVAL_HPP