eval: suppress the depth-0 heuristic batch-axis fallback under a batched objective

The batched evaluator's candidate_axes() falls back, at depth 0 on a node
carrying no accepted optimizer annotation, to batching the largest contracted
batchable index (batch_axis()). That legacy single-axis convenience predates
the per-node batch-axis annotations the batched (peak) objectives now emit.
Under such an objective the annotations are authoritative: an un-annotated
node means "do not batch here", so the fallback wrongly batches nodes the
peak-threshold objective declined -- measured on aux-only water-8 as 30 of 37
aux triggers, +34% product ops, and a small accuracy perturbation from the
per-batch screening relaxation.

Add BatchPolicy::suppress_heuristic_fallback, read only by the runtime
evaluator (make_evaluator -> make_batched_custom_evaluator); default false
preserves the legacy behavior byte-for-byte. Also add a behavior-neutral
BatchAxes trace line (depth / picked axis+kind / annotation / result indices),
gated on log::printing(), to diagnose which batch axes are annotation- vs
heuristic-driven.

Author: evaleev

SeQuant/core/batch_policy.hpp

@@ -26,6 +26,10 @@ struct BatchPolicy {
   /// yet contracted at no node -- is eligible for batching; its per-slice size
   /// comes from \c batch_target_size(ix) like any batchable index. Default
   /// false = no spectator batching (byte-identical to non-spectator behavior).
+  /// Necessary but not sufficient: spectator axes are emitted only under a
+  /// TIME-FIRST objective (DenseTimeSpaceBatched) and only when the selected
+  /// root's modeled peak exceeds \c peak_threshold. Spectator batching is
+  /// therefore currently unavailable under the space-first objectives.
   bool batch_spectator_indices = false;
 
   /// If true, restrict batching to persistent (amplitude-independent) subtrees,
@@ -39,6 +43,18 @@ struct BatchPolicy {
   /// Read identically by the single-term optimizer and the runtime evaluator.
   bool persistent_only = false;
 
+  /// If true, the runtime batched evaluator does NOT apply its depth-0
+  /// heuristic fallback (batching the largest contracted batchable index when a
+  /// node carries no accepted optimizer annotation). Set this when a batched
+  /// objective is active: the optimizer then emits per-node batch-axis
+  /// annotations that are AUTHORITATIVE, so an un-annotated node means "do not
+  /// batch here" and the legacy single-axis fallback would wrongly override
+  /// that (batching nodes the peak-threshold objective declined). Default false
+  /// preserves the legacy single-axis convenience for callers/objectives that
+  /// do not annotate. Read only by the runtime evaluator (make_evaluator); the
+  /// optimizer ignores it.
+  bool suppress_heuristic_fallback = false;
+
   /// Footprint multiplier for the in-flight batch contribution that co-resides
   /// with a batch-accumulated intermediate (K += contribution). 0 = ignore
   /// (default); ~1 = full contribution materialized; backend-specific (TA's
@@ -48,11 +64,22 @@ struct BatchPolicy {
   /// batchable index.
   double accumulation_factor = 0.0;
 
-  /// Peak-memory budget in BYTES for the batched objective. The single-term
-  /// optimizer minimizes flops among schedules whose modeled peak is <=
-  /// peak_threshold, falling back to min-peak (best effort) when none fit.
-  /// Default +infinity => every schedule feasible => min flops => no batching.
-  /// This is the *enable* trigger for batching (a finite value turns it on).
+  /// Peak-memory budget in BYTES for the batched objectives. Its meaning
+  /// DIFFERS between them:
+  ///
+  /// - SPACE-FIRST (DenseSpaceTimeBatched): a hard feasibility gate. The
+  ///   single-term optimizer minimizes flops among schedules whose modeled peak
+  ///   is <= peak_threshold, falling back to min-peak (best effort) when none
+  ///   fit. Default +infinity => every schedule feasible => min flops => no
+  ///   batching, i.e. here a finite value is the *enable* trigger for batching.
+  ///
+  /// - TIME-FIRST (DenseTimeSpaceBatched): NOT a feasibility gate. Root
+  ///   selection ignores it entirely (peak breaks exact flop ties only), so it
+  ///   can neither constrain the schedule's peak nor enable CONTRACTED-axis
+  ///   batching (which is emitted regardless). Its ONLY effect is to trigger
+  ///   EXTERNAL (spectator) axis emission, together with
+  ///   \c batch_spectator_indices: axes are emitted iff the selected root's
+  ///   modeled peak exceeds this threshold.
   double peak_threshold = std::numeric_limits<double>::infinity();
 };
 

SeQuant/core/eval/eval.hpp

@@ -1361,9 +1361,11 @@ template <typename F, typename IndexPredicate = accept_any_index,
     F le, 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) {
+    std::size_t depth = 0, PeakSink peak = nullptr,
+    bool suppress_heuristic_fallback = false) {
   return [le = std::move(le), target_batch_size = std::move(target_batch_size),
           accept, is_volatile, persistent_only, depth, peak,
+          suppress_heuristic_fallback,
           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
@@ -1381,12 +1383,19 @@ template <typename F, typename IndexPredicate = accept_any_index,
     // scratch would re-batch inner nodes of an unannotated subtree (extra,
     // unintended work) rather than realizing the optimizer's chosen multi-axis
     // nesting.
-    auto candidate_axes = [&accept,
-                           depth](auto const& n) -> container::svector<Index> {
+    auto candidate_axes = [&accept, depth, suppress_heuristic_fallback](
+                              auto const& n) -> container::svector<Index> {
       container::svector<Index> out;
       for (auto const& entry : n->batch_axes())
         if (accept(entry.first)) out.push_back(entry.first);
-      if (out.empty() && depth == 0)
+      // Depth-0 heuristic fallback: when the node carries no accepted
+      // annotation, batch the largest contracted batchable index (legacy
+      // single-axis behavior). Under a batched objective the optimizer's
+      // annotations are authoritative -- an empty annotation means "do not
+      // batch this node" -- so suppress_heuristic_fallback (BatchPolicy) skips
+      // this fallback to honor the optimizer's peak-threshold-gated decision.
+      // Default false = byte-identical (fallback active).
+      if (out.empty() && depth == 0 && !suppress_heuristic_fallback)
         if (auto const h = batch_axis(n, accept)) out.push_back(*h);
       return out;
     };
@@ -1448,6 +1457,28 @@ template <typename F, typename IndexPredicate = accept_any_index,
         break;
       }
 
+    // DEBUG (behavior-neutral): log the trigger's depth, picked axis + kind,
+    // and its full batch_axes() annotation + result indices, to diagnose nested
+    // re-batching of a single aux axis. Emitted only when tracing is on.
+    if (log::printing()) {
+      std::string annot;
+      for (auto const& [ix, knd] : node->batch_axes()) {
+        annot += toUtf8(ix.full_label());
+        annot += (knd == AxisKind::External ? ":ext " : ":con ");
+      }
+      std::string res;
+      for (auto const& ix : node->canon_indices()) {
+        res += toUtf8(ix.full_label());
+        res += " ";
+      }
+      log::log(
+          "BatchAxes",
+          std::format("depth={} picked={}:{} nbatch={} annot=[{}] result=[{}]",
+                      depth, toUtf8(K.full_label()),
+                      picked_kind == AxisKind::External ? "ext" : "con",
+                      batches.size(), annot, res));
+    }
+
     if (picked_kind == AxisKind::External) {
       // SCATTER branch. K survives to node's result as a free spectator mode,
       // so the per-block partials are disjoint slices of one result (not
@@ -1502,7 +1533,7 @@ template <typename F, typename IndexPredicate = accept_any_index,
         bs.cache.set_custom_evaluator(make_batched_custom_evaluator(
             std::function<ResultPtr(node_t const&)>{le_g}, target_batch_size,
             accept, make_scope_guard, is_volatile, persistent_only, depth + 1,
-            peak));
+            peak, suppress_heuristic_fallback));
         ResultPtr part = evaluate(node, le_g, bs.cache);
         // Pre-size on the first block (learns the result's non-axis extents and
         // kind from the block partial; the external axis is widened to full).
@@ -1655,7 +1686,7 @@ template <typename F, typename IndexPredicate = accept_any_index,
           bs.cache.set_custom_evaluator(make_batched_custom_evaluator(
               std::function<ResultPtr(node_t const&)>{le_g}, target_batch_size,
               accept, make_scope_guard, is_volatile, persistent_only, depth + 1,
-              peak));
+              peak, suppress_heuristic_fallback));
           ResultPtr part = evaluate(*mem, le_g, bs.cache);
           if (!acc[m])
             acc[m] = std::move(part);
@@ -1751,7 +1782,8 @@ template <class F, class ScopeGuardFactory = make_no_scope_guard>
   return make_batched_custom_evaluator(
       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);
+      policy.persistent_only, /*depth=*/0, peak,
+      policy.suppress_heuristic_fallback);
 }
 
 }  // namespace sequant