[AutoDiff] determine_ad_stack_size: replace whole-CFG Bellman-Ford with SCC + DAG DP

[duburcqa]

`determine_ad_stack_size`: replace whole-CFG Bellman-Ford with SCC + DAG DP

Four commits, no behavior change. Reduces the AD-stack sizing pass from a per-stack global Bellman-Ford on the full CFG to a Tarjan SCC condensation plus per-stack DAG dynamic programming. On a large reverse-mode kernel that previously hung the compiler indefinitely (V=9531, ~9514 nodes inside cyclic SCCs, 4241 AD-stacks), the pass now returns in 163 ms total.

TL;DR

`ControlFlowGraph::determine_ad_stack_size()` runs once per kernel to size every adaptive `AdStackAllocaStmt`. The existing pass solves a max-prefix-sum longest-path problem with Bellman-Ford SPFA, run once per AD-stack against the full CFG. On large reverse-mode kernels with thousands of AD-stacks and multi-thousand-node loop bodies, the aggregate cost is K * V * E, and the autodiff push/pop pattern (paired in the same loop SCC, balanced cycle sums) puts each per-stack SPFA in its hard `|V|+1`-pushes-per-node ceiling.

This PR replaces the algorithm with:

• Tarjan SCC condensation of the CFG, computed once and shared across all AD-stacks
• Deduplication of AD-stacks with bit-identical `(increased_size, max_increased_size)` per-node row pairs
• Per representative: topological DP across the SCC DAG; inside each cyclic SCC, a single-pass DP on DFS-finish-time-ordered forward edges plus one back-edge relaxation check for positive-cycle detection
• Sign-based fast paths for the trivial cyclic-SCC cases

Output of the pass (every AD-stack's `max_size`, plus the positive-cycle flag that routes unresolvable shapes into the structural `SizeExpr` pre-pass) is byte-identical to the prior pass on the C++ unit test matrix and on a downstream end-to-end reverse-mode test that previously hung.

Why

Adaptive AD-stacks (`AdStackAllocaStmt` with `max_size = 0`) need their static upper bound resolved per kernel. The semantic problem is: for each AD-stack, what is the maximum running net push count along any walk from the kernel entry? The existing pass uses a queue-based Bellman-Ford SPFA with `|V|+1`-pushes-per-node positive-cycle detection, run once per AD-stack against the entire CFG.

On large autodiff kernels:

• K (number of AD-stacks) is in the thousands -- one alloca per autodiff variable
• V and E (CFG nodes/edges) are also in the thousands -- autodiff transforms inflate the IR several-fold
• The body of the autodiff range-for forms a single huge cyclic SCC, and stacks have paired push/pop in that SCC. Sum of `is` over the SCC = 0 (balanced cycle), so the SPFA never trips its positive-cycle short-circuit but iterates to its hard ceiling
• The aggregate cost reaches K * V * E, which on the largest kernel observed was effectively non-terminating

Mechanism

Four sequential commits, each a strict simplification on top of the prior version.

1. Drop hash overhead from the per-stack hot loop

`AdStackAllocaStmt*`-keyed `unordered_map<AdStackAllocaStmt*, vector>` for the per-stack `increased_size` and `max_increased_size` arrays is replaced with `vector<vector>` indexed by a contiguous int stack id (one hash insert at stack discovery, plain vector access in the hot loop). The per-stack vector references are hoisted out of the SPFA visit loop, and outgoing-edge node ids are precomputed once per source node so the inner loop walks an int vector instead of hashing `node_ids[next_node]` per edge traversal. AD-stacks with no push/pop in the CFG are skipped entirely (they were running a no-op SPFA before).

2. Deduplicate AD-stacks with identical push/pop rows

A sparse fingerprint of `(node_id, increased_size[node], max_increased_size[node])` triples per AD-stack groups stacks into equivalence classes. The longest-path computation runs once per representative; `(max_size, has_positive_loop)` is broadcast to every member of the class. Worst case (no duplicates) reduces to running SPFA once per stack -- strictly opportunistic.

3. Scope per-stack work to Tarjan SCCs

Tarjan SCC condensation of the CFG once, shared across all per-stack runs. Output: `scc_id[n]`, `scc_nodes[s]`, `dfs_finish[n]`. Per representative: walk the SCC DAG in topological order (sources first, since Tarjan emits SCCs in reverse-topological order); each non-cyclic SCC is processed in O(1) per node; bounded SPFA runs only inside cyclic SCCs, with cycle detection scoped to the SCC. Inter-SCC edges relax forward and are touched O(1) per stack since predecessor SCCs are already finalized.

4. Replace cyclic-SCC SPFA with DAG DP + back-edge check (the unlock)

The previous step still ran bounded SPFA inside each cyclic SCC. On autodiff kernels this is the band: paired push/pop in the same loop SCC means signs are mixed and SPFA iterates to its `|S|+1`-pushes-per-node ceiling on every stack.

Fix: during Tarjan, record DFS finish time per node. Split each cyclic SCC's intra-edges into a forward set (target finishes before source) and a back set (target finishes at or after source). Per representative:

```
for each cyclic SCC S in topological order:
if min_is(S) >= 0 and max_is(S) > 0:
positive cycle
// every node lies on some cycle (SCC property), and a cycle through a strictly-positive node
// with all non-negative is along it sums positive
elif min_is(S) == 0 == max_is(S):
spread max entry-side begin to every node in S in O(|S|)
else:
// single-pass DAG DP on forward edges
for each u in S in descending dfs_finish order:
for each forward edge u -> v:
relax max_size_at_node_begin[v]
// one relaxation check on back edges
for each u in S, each back edge u -> v:
if begin[u] + is[u] > begin[v]:
positive cycle
```

Correctness: any walk inside an SCC decomposes into a forward path along non-back edges plus zero or more closed cycles formed by a back-edge plus a forward path. For SCCs with no positive cycle, traversing a cycle adds a non-positive amount to the running size and so cannot improve `max_size` beyond what the forward DP already computed. The back-edge relaxation check after the DP detects the only failure mode (some back-edge would still improve a forward predecessor's value, which can only happen if a positive cycle exists for this stack).

Per-cyclic-SCC cost drops from O(|S| * |E_S|) to O(|S| + |E_S|).

Per-kernel measurements

Measured on a large reverse-mode autodiff test (full simulation step plus backward pass) on CPU. Same input on both columns:

Kernel	V	active stacks	reps	cyclic SCCs / nodes	dp ms (before)	dp ms (after)
`kernel_update_cartesian_space_c289_0`	567	456	226	1 / 564	178	0.5
`kernel_step_2_c501_0`	231	93	82	3 / 226	2.1	0.1
`kernel_compute_qacc_c303_0`	1520	584	584	2 / 1516	2700	4.2
`kernel_forward_dynamics_without_qacc_c305_0`	9531	4241	3959	15 / 9514	(never finished)	113

End-to-end test: previously hung indefinitely on the last kernel; now passes in 152 s (full simulation step plus backward pass plus correctness assertions).

Tests

• `tests/cpp/transforms/determine_ad_stack_size_test.cpp`: 8 existing scenarios pass unchanged (Basic, Loop, LoopPushOnlyBoundByTripCount, EmptyNodes, plus a 4-way `If` parametrization). These exercise the algorithm's full input shape: push-only loops, push-then-pop-paired bodies, branches, unreachable-from-entry sub-CFGs.
• `tests/python/test_ad_basics.py`: 147 tests pass unchanged.
• A downstream end-to-end reverse-mode rigid-body simulation test that previously hung now passes.

The cpp tests cover all the algorithm shapes that were formerly handled by the bounded SPFA, so the DAG-DP-plus-back-edge-check version is exercised for all of them. No new tests are necessary; the existing matrix already pins down the algorithm's input/output behavior.

Side-effect audit

Concern	Where checked	Verdict
Output for stacks with no positive cycle	DAG DP on forward edges processed in descending DFS finish-time order	Exact max-walk-sum on the back-edge-removed DAG, equals SCC's max-walk-sum when no positive cycle exists
Output for stacks with positive cycle	back-edge relaxation check after forward DP	Same `has_positive_loop` flag as before; routed to the structural `SizeExpr` pre-pass identically
Self-loops (singleton-SCC cyclic case)	classified as back-edges via `dfs_finish[u] == dfs_finish[v]`; the singleton SCC's cyclic flag walks `next_ids_intra_back`	Self-loop with positive `is` correctly flags positive cycle via the back-edge relaxation step
Inter-SCC propagation order	`next_ids_inter[u]` is the dedicated set; relaxed only after the SCC has converged	Predecessor-SCC values are finalized when an SCC is entered -- single-pass topological order on the SCC DAG matches Tarjan's emission order
Per-kernel determinism	fingerprint over `(node_id, is, mis)` triples; node ids are deterministic per CFG	Same input -> same representative -> same output
Public API	`ControlFlowGraph::determine_ad_stack_size()` signature unchanged; behavior contract unchanged (positive cycle -> `max_size = 0`; no compile-time fallback)	Caller in `irpass::determine_ad_stack_size` (`quadrants/transforms/determine_ad_stack_size.cpp`) is unmodified
Downstream IR passes / clients	`AdStackAllocaStmt::max_size` is the only output; no new fields	All consumers that read `max_size` continue to see identical values

[hughperkins]

Does this need to change user-visible doc in any way?

[duburcqa]

Does this need to change user-visible doc in any way?

I don't think it does. This is just a performance optimisation of the control graph walker used to generate the small kernel that compute the optimal adstack size at runtime before running the backward pass.

[hughperkins]

checklist:

• no doc changes needed

=> ok to merge

[github-actions]

Coverage Report (3d6bc5616)

Metric	Value
Diff coverage (changed lines only)	0%
Overall project coverage	73%

Total: 0 lines, 0 missing, 0% covered