Merge torchwright:main from /data/torchdoom2 (CP-SAT feasible warm-start b017bbd)

physicsrob · physicsrob · commit 8c5041102fdb · 2026-06-07T19:00:17.000Z
Brings the CP-SAT no-eager feasible warm-start hint (post-J d=4096) from the
parallel torchdoom2 checkout into main alongside the per-column gate-offset M
fix. Disjoint file sets (compiler/scheduler vs ops/), no conflicts.
diff --git a/docs/cpsat_scheduler.md b/docs/cpsat_scheduler.md
@@ -553,6 +553,36 @@ heuristic's layer count also tightens the search horizon
 (`max_layers = min(user_max, hint_n_layers + 1)`), which shrinks
 each `layer_var`'s domain.
 
+**The warm-start runs with eager freeing disabled (`eager_free=False`).**
+The default heuristic frees a node's columns *within* its consumer's
+layer and reuses them the same layer (within-layer reuse).  The
+layer-granular CP-SAT model — where a node occupies the residual
+stream for the half-open layer interval `[birth, cancel)` — cannot
+express that density, so the *eager* schedule is an **infeasible**
+hint: CP-SAT discards it, cold-searches, and times out into the
+silent heuristic fallback (the post-J DOOM symptom at d=4096/d_head=32:
+`optimize=2` returned 85 via fallback).  The no-eager schedule is one
+or two layers deeper (87 vs 85 at d=4096) but model-representable, so
+CP-SAT gets a *real feasible incumbent* to verify and improve — it
+walks 87 → 86 → … → 81–82 within the `optimize=2` budget (the exact
+floor varies run-to-run with CP-SAT's parallel workers).  The heuristic
+**fallback** (used only when CP-SAT finds nothing) keeps the default
+eager behavior, so a timeout never regresses below the eager depth.
+
+**Rollback-cancel correction.**  `LayerScheduler` speculatively
+allocates a node, then rolls the allocation back with `free(node)`
+when the op can't be committed under the dirty-cancel / head budget.
+That rollback is not a real death — the node is re-allocated (reborn)
+in a later layer, and may then die by `reassign` (free-add, never
+`free`).  Recording the rollback layer as the cancel produced hints
+with `cancel < birth` (e.g. freed at layer 6, actually born at 7),
+which is hard-infeasible in the model — exactly the 3 records that kept
+the otherwise-feasible no-eager hint from being accepted as a complete
+incumbent.  `_TrackingResidualStreamMap` therefore clears a node's
+stale cancel whenever it is re-allocated or reborn via `reassign`, so
+the emitted `hint_cancel` is internally consistent (`cancel ≥ birth+1`,
+or omitted when the node dies by reassign).
+
 ### Cancel-domain restriction
 
 The cancel decision space is the dominant LB-search cost when the
@@ -610,18 +640,25 @@ feasible but possibly non-optimal.  `forward_compile` accepts it —
 ### Geometry sensitivity
 
 The win-size from CP-SAT versus heuristic depends on residual-stream
-slack.  Measured on the headless DOOM graph (~4.4K nodes); each row
-records the ``optimize`` level (and corresponding budget) used:
-
-| geometry            | -O / budget | heuristic | CP-SAT | Δ    | first incumbent | OPTIMAL at |
-|---------------------|-------------|----------:|-------:|-----:|----------------:|-----------:|
-| d=2048, d_h=8192    | -O 1 (60s)  | 61        | (none) | n/a  | not in 60s      | not in 60s |
-| d=3072, d_h=8192    | -O 2 (180s) | 58        | 46     | -21% | ~27s            | ~80s       |
-| d=4096, d_h=4096    | -O 1 (60s)  | 59        | 46     | -22% | ~17s            | ~31s       |
+slack.  The first three rows were measured on an earlier headless DOOM
+graph (~4.4K nodes); the last row is the post-J graph (~13K nodes,
+the flat pass roughly doubled depth).  Each row records the
+``optimize`` level (and corresponding budget) used:
+
+| geometry                       | -O / budget | heuristic | CP-SAT | Δ    | first incumbent | OPTIMAL at |
+|--------------------------------|-------------|----------:|-------:|-----:|----------------:|-----------:|
+| d=2048, d_h=8192               | -O 1 (60s)  | 61        | (none) | n/a  | not in 60s      | not in 60s |
+| d=3072, d_h=8192               | -O 2 (180s) | 58        | 46     | -21% | ~27s            | ~80s       |
+| d=4096, d_h=4096               | -O 1 (60s)  | 59        | 46     | -22% | ~17s            | ~31s       |
+| post-J d=4096, d_h=4096, dh=32 | -O 2 (180s) | 85        | 81–82  | -4%  | ~70–80s         | never (LB 50) |
 
 At d=2048 the residual cumulative is the binding constraint and
 CP-SAT struggles to close the LB gap within budget.  At d=3072+
-CP-SAT converges optimally inside the budget.  The heuristic-
+CP-SAT converges optimally inside the budget.  On the much larger
+post-J graph CP-SAT no longer proves optimality (the lower bound
+sticks at 50) but still beats the heuristic — and only once the
+warm-start feeds a *feasible* (no-eager) hint; with the eager hint it
+finds no incumbent and falls back.  The heuristic-
 fallback behavior (when CP-SAT can't find an incumbent) is the
 right answer for d=2048 — users always get a schedule, just not
 the CP-SAT one.
@@ -643,7 +680,23 @@ the CP-SAT one.
   solution.  CP-SAT's first feasible reproduces the heuristic
   warm-start; stopping there returned no improvement over
   `optimize=0` while paying the model-build cost.  Removed.
-- **`repair_hint=True`** would have let CP-SAT actively complete
-  the partial hint into a feasible solution, but it conflicts
-  with `AddDecisionStrategy` (CP-SAT crashes with
-  "fixed_search != nullptr").  Not pursued.
+- **`repair_hint=True`** would have let CP-SAT actively complete the
+  partial hint into a feasible solution, but it is **unusable in this
+  OR-Tools version** (v9.15): it aborts the process with
+  `Check failed: heuristics.fixed_search != nullptr` — *with and
+  without* `AddDecisionStrategy`, and even with
+  `search_branching=FIXED_SEARCH` (where the crash just moves from
+  setup into a parallel worker mid-search).  (An earlier note here
+  claimed it merely "conflicts with `AddDecisionStrategy`"; that was
+  wrong — removing the strategy does not avoid the crash.)  And it
+  would not have helped regardless: measured on the post-J DOOM graph,
+  the ~70 s to first incumbent is **genuine model-size search**
+  (≈23 s presolve + ≈45 s LB/feasibility search on a ~80 k-variable
+  model with three coupled `AddCumulative` constraints), not a blocked
+  seed.  The raw hint (with the 3 infeasible `cancel < birth` records)
+  and the clean hint reach their first incumbent at the *same* time —
+  the soft `AddHint` silently drops the infeasible values either way —
+  so the cancel artifacts were a correctness curiosity, not the
+  performance lever.  The lever is feeding a *feasible* hint at all
+  (the no-eager schedule, see *Warm-start hints*) plus an `optimize≥2`
+  budget; CP-SAT then improves 87 → 81.
diff --git a/tests/compile/forward/test_residual_map.py b/tests/compile/forward/test_residual_map.py
@@ -140,3 +140,86 @@ def test_resolve_indices_concatenate():
     # Concatenate is NOT allocated — resolve_indices resolves through it
     result = rmap.resolve_indices(cat)
     assert result == idx_a + idx_b
+
+
+def test_tracking_map_clears_rollback_cancel_on_reallocate():
+    """A node freed by a *rolled-back* speculative allocation, then re-allocated
+    (reborn) at a later layer, must not keep the rollback layer as its cancel.
+
+    Regression for the CP-SAT warm-start hint producing ``cancel < birth``:
+    ``LayerScheduler`` speculatively allocates a node, fails its dirty-cancel/
+    head budget, and rolls the allocation back with ``free(node)`` at layer 6 —
+    but the node is really born at layer 7 and dies by ``reassign`` (free-add,
+    never ``free``).  The recorded ``cancel = 6 < birth = 7`` was hard-infeasible
+    in the CP-SAT model.  The tracking map clears the stale cancel on
+    (re)allocate / reassign so the emitted hint is internally consistent.
+    """
+    from torchwright.compiler.forward.compile import _TrackingResidualStreamMap
+
+    base = ResidualStreamMap(64)
+    tmap = _TrackingResidualStreamMap(base)
+    x = InputNode("x", 1, value_range=(-100.0, 100.0))
+
+    # Layer 6: speculative allocate then rollback free -> records cancel=6.
+    tmap.current_layer = 6
+    tmap.allocate(x)
+    tmap.free(x)
+    assert tmap.cancel_layer[x.node_id] == 6
+
+    # Layer 7: the node is really born here -> the stale cancel must be cleared.
+    tmap.current_layer = 7
+    tmap.allocate(x)
+    assert x.node_id not in tmap.cancel_layer
+
+    # A genuine later free re-records the real cancel.
+    tmap.current_layer = 9
+    tmap.free(x)
+    assert tmap.cancel_layer[x.node_id] == 9
+
+
+def test_node_deepcopy_preserves_identity():
+    """`copy.deepcopy` of a container that *references* nodes must keep node
+    identity (nodes are graph singletons keyed on node_id), while still
+    deep-copying the surrounding container.
+
+    Regression for the warm-start: `_run_heuristic_warm_start` does
+    `copy.deepcopy(residual_map)` to isolate its scheduler mutations.  Without
+    `Node.__deepcopy__` returning self, that cloned the `pos` node, breaking the
+    `n is self.pos_encoding` identity check so the warm-start freed `pos` early
+    and produced a hint the CP-SAT model (which reserves pos) rejected by exactly
+    len(pos) columns.
+    """
+    import copy
+
+    n = InputNode("x", 4, value_range=(-100.0, 100.0))
+    container = {"node": n, "list": [n]}
+    dup = copy.deepcopy(container)
+
+    assert dup["node"] is n  # node identity preserved (not cloned)
+    assert dup["list"][0] is n
+    assert dup is not container  # the container itself is still deep-copied
+    assert dup["list"] is not container["list"]
+
+
+def test_tracking_map_clears_stale_cancel_on_reassign():
+    """A node reborn via ``reassign`` (the free-add path) must not keep a stale
+    cancel from an earlier rolled-back allocation."""
+    from torchwright.compiler.forward.compile import _TrackingResidualStreamMap
+
+    base = ResidualStreamMap(64)
+    tmap = _TrackingResidualStreamMap(base)
+    y = InputNode("y", 1, value_range=(-100.0, 100.0))
+    w = InputNode("w", 1, value_range=(-100.0, 100.0))
+
+    # Earlier rolled-back allocation of y records a stale cancel.
+    tmap.current_layer = 6
+    tmap.allocate(y)
+    tmap.free(y)
+    assert tmap.cancel_layer[y.node_id] == 6
+
+    # Later, y is reborn by reassigning w's columns to it (free-add reuse).
+    tmap.current_layer = 8
+    tmap.allocate(w)
+    tmap.reassign(w, y)
+    assert y.node_id not in tmap.cancel_layer
+    assert tmap.is_allocated(y)
diff --git a/tests/compile/forward/test_scheduler.py b/tests/compile/forward/test_scheduler.py
@@ -970,3 +970,27 @@ def test_output_already_computed():
     scheduler = LayerScheduler(graph, D, D_HEAD, pos)
     # Should not raise — nothing to do
     attn_ops, mlp_ops, _biased = scheduler.schedule_layer(rmap, computed)
+
+
+def test_eager_free_flag_gates_within_layer_freeing():
+    """``eager_free=False`` makes ``_freshly_dead_inputs`` a no-op, so the
+    CP-SAT warm-start schedule never frees+reuses a column within a consumer's
+    layer — keeping it representable by the layer-granular CP-SAT model.  The
+    default ``eager_free=True`` (used by the heuristic fallback) surfaces the
+    freshly-dead intermediate so its column can be reclaimed in the same layer.
+    """
+    pos = _make_pos_encoding()
+    x = InputNode("x", D, value_range=(-100.0, 100.0))
+    a = _make_linear(x, D, "a")  # intermediate: a's only consumer is b
+    b = _make_linear(a, D, "b")
+    graph = GraphAnalyzer(b)
+
+    rmap = ResidualStreamMap(D * 4)
+    rmap.allocate(a)  # a is live; b has just been placed -> a is now dead
+
+    computed = {x, a, b}
+    eager = LayerScheduler(graph, D, D_HEAD, pos, eager_free=True)
+    noeager = LayerScheduler(graph, D, D_HEAD, pos, eager_free=False)
+
+    assert a in eager._freshly_dead_inputs(b, computed, rmap)
+    assert noeager._freshly_dead_inputs(b, computed, rmap) == []
diff --git a/torchwright/compiler/forward/compile.py b/torchwright/compiler/forward/compile.py
@@ -52,6 +52,21 @@ class _TrackingResidualStreamMap(ResidualStreamMap):
     records that value for every node that gets freed.  Nodes
     consumed via ``reassign`` (free-add path) don't go through
     ``free`` and are correctly omitted from the cancel hint.
+
+    **Rollback-free correction.**  ``LayerScheduler`` speculatively
+    allocates a node, then rolls the allocation back with ``free(node)``
+    when the op can't be committed under the dirty-cancel / head budget
+    (``scheduler.py`` ``residual_map.free(node); continue``).  That
+    rollback is not a real death — the node is re-allocated (reborn) in
+    a later layer.  Recording the rollback layer as the node's cancel
+    produced a hint with ``cancel < birth`` (e.g. freed at layer 6,
+    actually born at layer 7), which is hard-infeasible in the CP-SAT
+    model.  To keep the emitted ``hint_cancel`` internally consistent,
+    a node's stale cancel record is cleared whenever it is re-allocated
+    or reborn via ``reassign``: a node is only ever (re)allocated after
+    a premature rollback free, so any prior cancel for it must be
+    discarded and re-recorded by its genuine later free (or omitted
+    entirely if it dies by reassign, per the free-add design above).
     """
 
     def __init__(self, base: ResidualStreamMap) -> None:
@@ -64,10 +79,24 @@ def __init__(self, base: ResidualStreamMap) -> None:
         self.current_layer: int = 0
         self.cancel_layer: dict[int, int] = {}
 
+    def allocate(self, node: Node):  # type: ignore[override]
+        # A node reaching allocate after a recorded free was rolled back
+        # (its earlier free was premature); drop the stale cancel so the
+        # genuine death — a later free, or omission if it dies by
+        # reassign — is what the hint reflects.
+        self.cancel_layer.pop(node.node_id, None)
+        return super().allocate(node)
+
     def free(self, node: Node) -> None:  # type: ignore[override]
         self.cancel_layer[node.node_id] = self.current_layer
         super().free(node)
 
+    def reassign(self, old_node: Node, new_node: Node) -> None:  # type: ignore[override]
+        # The new node is born here (free-add reuse), so any stale cancel
+        # recorded for it by an earlier rolled-back allocation is wrong.
+        self.cancel_layer.pop(new_node.node_id, None)
+        super().reassign(old_node, new_node)
+
 
 def _run_heuristic_warm_start(
     graph: GraphAnalyzer,
@@ -105,6 +134,15 @@ def _run_heuristic_warm_start(
         admission_budget_fraction=admission_budget_fraction,
         policy=policy,
         pinned_nodes=overlay_pinned_inputs,
+        # The warm-start must hand CP-SAT a *model-representable* schedule.
+        # Eager (within-layer) freeing produces a shallower schedule that frees
+        # and reuses a column inside a consumer's layer — a density the
+        # layer-granular CP-SAT model cannot express, so the eager schedule is
+        # an infeasible hint and CP-SAT cold-searches and times out.  The
+        # no-eager schedule is deeper but feasible, giving the solver a real
+        # incumbent to improve from (d=4096/d_head=32: 87 -> ~81 vs the eager
+        # heuristic's 85).  The heuristic *fallback* below stays eager.
+        eager_free=False,
     )
     hint_layers: dict = {}
     hint_routing: dict = {}
@@ -703,19 +741,22 @@ def forward_compile(
             # error, so a layer-count-only measurement cannot tell it from
             # a real solve.  Surface it loudly (and optionally fatally via
             # ``require_solver``) so it can never masquerade as an
-            # optimizer result.  ``INFEASIBLE`` here is almost always a
-            # model bug (the heuristic found a feasible schedule yet
-            # CP-SAT's model proved none exists) rather than a hard
-            # instance — see ``docs/cpsat_scheduler.md``.
+            # optimizer result.  ``UNKNOWN`` means CP-SAT found no incumbent
+            # within the budget (expected at small budgets / hard geometries);
+            # ``INFEASIBLE`` on a graph the heuristic schedules would indicate
+            # a CP-SAT model bug — see ``docs/cpsat_scheduler.md``.  The
+            # fallback uses the EAGER heuristic (below), which is typically
+            # shallower than the no-eager warm-start hint (``hint_n_layers``).
             fallback_msg = (
                 f"CP-SAT returned no usable assignment "
                 f"(status={_stats.status_name}, "
                 f"{cpsat_time_budget_s:.0f}s budget); falling back to the "
-                f"heuristic schedule ({hint_n_layers} layers). The compile "
-                f"is valid but UNOPTIMIZED — optimize>0 did not take "
-                f"effect. status=INFEASIBLE on a graph the heuristic "
-                f"schedules indicates a CP-SAT model bug, not a hard "
-                f"instance."
+                f"eager heuristic schedule (the no-eager warm-start hint was "
+                f"{hint_n_layers} layers; the eager fallback is typically "
+                f"shallower). The compile is valid but UNOPTIMIZED — "
+                f"optimize>0 did not take effect. status=INFEASIBLE (vs "
+                f"UNKNOWN/timeout) on a graph the heuristic schedules would "
+                f"indicate a CP-SAT model bug, not a hard instance."
             )
             if require_solver:
                 raise RuntimeError(
@@ -728,6 +769,9 @@ def forward_compile(
                 for line in _stats.solver_log.splitlines()[-40:]:
                     print(f"  {line}")
                 print("--- end CP-SAT solver log ---")
+            # Eager heuristic fallback (default eager_free=True): a timeout
+            # never regresses below the eager heuristic's depth, even though
+            # the CP-SAT hint was the deeper no-eager schedule.
             scheduler = LayerScheduler(
                 graph,
                 d,
diff --git a/torchwright/compiler/forward/cpsat_scheduler.py b/torchwright/compiler/forward/cpsat_scheduler.py
@@ -984,14 +984,39 @@ def build_cpsat_model(
     resid_intervals: List = []
     resid_demands: List[int] = []
     for n in residual_nodes:
-        size = model.NewIntVar(1, max_layers + 1, f"rsz_n{n.node_id}")
-        model.Add(size == cancel_layer[n.node_id] - layer_var[n.node_id])
-        iv = model.NewIntervalVar(
-            layer_var[n.node_id],
-            size,
-            cancel_layer[n.node_id],
-            f"riv_n{n.node_id}",
-        )
+        if isinstance(n, Add) and n.node_id in is_free:
+            # Free-add reuses a dead addend's already-allocated residual
+            # columns (`reassign` in both `LayerScheduler._schedule_attn_
+            # sublayer` and the directed replay), so it adds NO fresh
+            # residual column at its birth layer.  The reused addend's own
+            # interval already covers that layer — `cancel_consumer_lb`
+            # forces the addend's cancel >= layer[A] + 1 because A is one of
+            # its consumers — so giving the Add a full interval starting at
+            # layer[A] double-counts the shared column for exactly one layer.
+            # Shift the Add's residual start by `is_free[A]`: a free-add
+            # (is_free=1) starts one layer later (the addend covers layer[A]),
+            # a compute-add (is_free=0) allocates fresh columns and starts at
+            # layer[A].  This is the residual-cumulative analogue of the
+            # BIRTH-dirty / attention-head free-add gating above, and removes
+            # the residual over-count that rejected schedules the heuristic
+            # compiles (the dead addend and its Add never occupy two distinct
+            # columns at the add layer).
+            start = model.NewIntVar(0, max_layers, f"rstart_n{n.node_id}")
+            model.Add(start == layer_var[n.node_id] + is_free[n.node_id])
+            size = model.NewIntVar(0, max_layers + 1, f"rsz_n{n.node_id}")
+            model.Add(size == cancel_layer[n.node_id] - start)
+            iv = model.NewIntervalVar(
+                start, size, cancel_layer[n.node_id], f"riv_n{n.node_id}"
+            )
+        else:
+            size = model.NewIntVar(1, max_layers + 1, f"rsz_n{n.node_id}")
+            model.Add(size == cancel_layer[n.node_id] - layer_var[n.node_id])
+            iv = model.NewIntervalVar(
+                layer_var[n.node_id],
+                size,
+                cancel_layer[n.node_id],
+                f"riv_n{n.node_id}",
+            )
         resid_intervals.append(iv)
         resid_demands.append(len(n))
     # Freeable inputs occupy residual columns from layer 0 until their cancel
diff --git a/torchwright/compiler/forward/scheduler.py b/torchwright/compiler/forward/scheduler.py
diff --git a/torchwright/graph/node.py b/torchwright/graph/node.py
