pounce_cli/
nl_hessian_program.rs

1//! Precompiled symbolic-Hessian program for one `Tape`.
2//!
3//! `Tape::hessian_accumulate` runs forward-over-reverse AD at every
4//! call: for each tape variable `j` it (a) match-dispatches every op
5//! in a forward-tangent sweep, (b) zeros adj/adj_dot, (c)
6//! match-dispatches every op again in the reverse-over-tangent
7//! sweep, and (d) HashMap-looks-up every `Var(k)` slot to find its
8//! Hessian output position. On evaluator-bound problems (dirichlet,
9//! lane_emden, henon) that match-dispatch + symbolic-AD overhead is
10//! ~80% of total CPU.
11//!
12//! This module compiles all of that ONCE at tape-build time into a
13//! flat `Vec<HOp>` of pre-resolved primitive ops:
14//!
15//!   * Forward pass — one `Fwd*` op per `TapeOp`. Mirrors
16//!     `Tape::forward`.
17//!   * Per-`j` forward tangent — only the ops touching slots that
18//!     statically depend on `j` are emitted (the rest stay zero
19//!     from the per-`j` `ZeroRange` reset).
20//!   * Per-`j` reverse-over-tangent — only ops on slots reachable
21//!     backward from output, with all slot indices and Hessian
22//!     output pointers pre-resolved.
23//!
24//! ## Scratch layout
25//!
26//! The program reads/writes a single `&mut [f64]` arena of
27//! `n_slots` cells. We allocate four contiguous regions of length
28//! `n` (`n` = `tape.ops.len()`):
29//!
30//!   * `v[i]`        in slot `i`
31//!   * `dot[i]`      in slot `n + i`
32//!   * `adj[i]`      in slot `2n + i`
33//!   * `adj_dot[i]`  in slot `3n + i`
34//!
35//! Per-`j` setup zeros the `[n, 4n)` range and seeds `adj[n-1]`.
36//! Allocation pattern is intentionally trivial — finer-grained
37//! slot recycling buys little once the dispatch loop is the
38//! bottleneck, and a contiguous layout makes the per-`j`
39//! `ZeroRange` reset a single `memset`-friendly loop.
40
41use std::collections::HashMap;
42
43use super::nl_tape::{Tape, TapeOp, div_u_ln10, log_d2, log10_d2, sqrt_d2};
44
45/// One primitive operation in the compiled Hessian program.
46/// `dst`/`a`/`b`/etc. are `u32` offsets into the caller's scratch
47/// slice; see the module docs for the slot layout.
48#[derive(Debug, Clone, Copy)]
49pub enum HOp {
50    // ===== Forward pass =====
51    FwdLoadVar {
52        dst: u32,
53        x_idx: u32,
54    },
55    FwdLoadConst {
56        dst: u32,
57        c_idx: u32,
58    },
59    FwdAdd {
60        dst: u32,
61        a: u32,
62        b: u32,
63    },
64    FwdSub {
65        dst: u32,
66        a: u32,
67        b: u32,
68    },
69    FwdMul {
70        dst: u32,
71        a: u32,
72        b: u32,
73    },
74    FwdDiv {
75        dst: u32,
76        a: u32,
77        b: u32,
78    },
79    FwdPow {
80        dst: u32,
81        a: u32,
82        b: u32,
83    },
84    FwdNeg {
85        dst: u32,
86        a: u32,
87    },
88    FwdAbs {
89        dst: u32,
90        a: u32,
91    },
92    FwdSqrt {
93        dst: u32,
94        a: u32,
95    },
96    FwdExp {
97        dst: u32,
98        a: u32,
99    },
100    FwdLog {
101        dst: u32,
102        a: u32,
103    },
104    FwdLog10 {
105        dst: u32,
106        a: u32,
107    },
108    FwdSin {
109        dst: u32,
110        a: u32,
111    },
112    FwdCos {
113        dst: u32,
114        a: u32,
115    },
116
117    // ===== Scalar slot init =====
118    SetZero {
119        dst: u32,
120    },
121    SetOne {
122        dst: u32,
123    },
124
125    // ===== Bulk reset (start of each j) =====
126    ZeroRange {
127        start: u32,
128        len: u32,
129    },
130
131    // ===== Forward tangent (per j) =====
132    DotAdd {
133        dst: u32,
134        a: u32,
135        b: u32,
136    },
137    DotSub {
138        dst: u32,
139        a: u32,
140        b: u32,
141    },
142    /// dot[d] = dot[a]*v[b] + v[a]*dot[b]
143    DotMul {
144        dst: u32,
145        dot_a: u32,
146        vb: u32,
147        va: u32,
148        dot_b: u32,
149    },
150    /// dot[d] = (dot[a]*v[b] - v[a]*dot[b]) / (v[b]*v[b])
151    DotDiv {
152        dst: u32,
153        dot_a: u32,
154        vb: u32,
155        va: u32,
156        dot_b: u32,
157    },
158    /// dot[d] = 0.5 / v[d] * dot[a]  (v[d] = sqrt(v[a]))
159    DotSqrt {
160        dst: u32,
161        dot_a: u32,
162        vd: u32,
163    },
164    /// dot[d] = v[d] * dot[a]  (v[d] = exp(v[a]))
165    DotExp {
166        dst: u32,
167        dot_a: u32,
168        vd: u32,
169    },
170    DotLog {
171        dst: u32,
172        dot_a: u32,
173        va: u32,
174    },
175    DotLog10 {
176        dst: u32,
177        dot_a: u32,
178        va: u32,
179    },
180    DotSin {
181        dst: u32,
182        dot_a: u32,
183        va: u32,
184    },
185    DotCos {
186        dst: u32,
187        dot_a: u32,
188        va: u32,
189    },
190    DotNeg {
191        dst: u32,
192        dot_a: u32,
193    },
194    DotAbs {
195        dst: u32,
196        dot_a: u32,
197        va: u32,
198    },
199    /// Compound: dot[d] for Pow(a, b). Carries the runtime
200    /// `u != 0` / `u > 0` branches.
201    DotPow {
202        dst: u32,
203        va: u32,
204        vb: u32,
205        vd: u32,
206        dot_a: u32,
207        dot_b: u32,
208    },
209
210    // ===== Reverse + adj_dot update (per j) =====
211    // Each op consumes adj[i] (= `w`) and adj_dot[i] (= `wd`) of
212    // the consumer slot, then `+=`-accumulates into the adj /
213    // adj_dot of the operand slots.
214    RevAdd {
215        adj_a: u32,
216        adj_b: u32,
217        adj_dot_a: u32,
218        adj_dot_b: u32,
219        w: u32,
220        wd: u32,
221    },
222    RevSub {
223        adj_a: u32,
224        adj_b: u32,
225        adj_dot_a: u32,
226        adj_dot_b: u32,
227        w: u32,
228        wd: u32,
229    },
230    RevMul {
231        adj_a: u32,
232        adj_b: u32,
233        adj_dot_a: u32,
234        adj_dot_b: u32,
235        w: u32,
236        wd: u32,
237        va: u32,
238        vb: u32,
239        dot_a: u32,
240        dot_b: u32,
241    },
242    RevDiv {
243        adj_a: u32,
244        adj_b: u32,
245        adj_dot_a: u32,
246        adj_dot_b: u32,
247        w: u32,
248        wd: u32,
249        va: u32,
250        vb: u32,
251        dot_a: u32,
252        dot_b: u32,
253    },
254    RevPow {
255        adj_a: u32,
256        adj_b: u32,
257        adj_dot_a: u32,
258        adj_dot_b: u32,
259        w: u32,
260        wd: u32,
261        va: u32,
262        vb: u32,
263        vd: u32,
264        dot_a: u32,
265        dot_b: u32,
266    },
267    RevNeg {
268        adj_a: u32,
269        adj_dot_a: u32,
270        w: u32,
271        wd: u32,
272    },
273    RevAbs {
274        adj_a: u32,
275        adj_dot_a: u32,
276        w: u32,
277        wd: u32,
278        va: u32,
279    },
280    RevSqrt {
281        adj_a: u32,
282        adj_dot_a: u32,
283        w: u32,
284        wd: u32,
285        va: u32,
286        vd: u32,
287        dot_a: u32,
288    },
289    RevExp {
290        adj_a: u32,
291        adj_dot_a: u32,
292        w: u32,
293        wd: u32,
294        vd: u32,
295        dot_a: u32,
296    },
297    RevLog {
298        adj_a: u32,
299        adj_dot_a: u32,
300        w: u32,
301        wd: u32,
302        va: u32,
303        dot_a: u32,
304    },
305    RevLog10 {
306        adj_a: u32,
307        adj_dot_a: u32,
308        w: u32,
309        wd: u32,
310        va: u32,
311        dot_a: u32,
312    },
313    RevSin {
314        adj_a: u32,
315        adj_dot_a: u32,
316        w: u32,
317        wd: u32,
318        va: u32,
319        dot_a: u32,
320    },
321    RevCos {
322        adj_a: u32,
323        adj_dot_a: u32,
324        w: u32,
325        wd: u32,
326        va: u32,
327        dot_a: u32,
328    },
329
330    // ===== Output =====
331    /// values[hess_ptr] += weight * scratch[adj_dot_slot].
332    HessEmit {
333        hess_ptr: u32,
334        adj_dot_slot: u32,
335    },
336}
337
338/// Precompiled Hessian-of-one-tape program. Built once via
339/// [`HessianProgram::compile`]; executed many times.
340#[derive(Debug, Clone)]
341pub struct HessianProgram {
342    ops: Vec<HOp>,
343    consts: Vec<f64>,
344    n_slots: u32,
345}
346
347impl HessianProgram {
348    /// Build the program. The `hess_map` is the same `(row, col)
349    /// -> values-index` map that [`Tape::hessian_accumulate`] uses;
350    /// the compiler inlines each lookup into a `HessEmit` op.
351    ///
352    /// Returns `None` when `tape` contains an opcode the program path
353    /// cannot lower (see [`program_supports_op`]). A caller must fall
354    /// back to the `Tape` (`build_with_externals`) interpreter path for
355    /// those tapes — this is a *graceful* signal, not a panic, so a
356    /// problem built from arbitrary user `.nl` input can never crash the
357    /// process here (code review L28).
358    pub fn compile(tape: &Tape, hess_map: &HashMap<(usize, usize), usize>) -> Option<Self> {
359        // Gate up front: every downstream sweep (forward / tangent /
360        // reverse) and the dependence/reachability analyses only handle the
361        // supported opcode set. Reject unsupported tapes here so none of
362        // those match arms is ever reached with an opcode it can't lower.
363        if !tape.ops.iter().all(program_supports_op) {
364            return None;
365        }
366
367        let n = tape.ops.len() as u32;
368        let v_base = 0u32;
369        let dot_base = n;
370        let adj_base = 2 * n;
371        let adj_dot_base = 3 * n;
372        let n_slots = 4 * n;
373
374        let v_slot = |i: u32| v_base + i;
375        let dot_slot = |i: u32| dot_base + i;
376        let adj_slot = |i: u32| adj_base + i;
377        let adj_dot_slot = |i: u32| adj_dot_base + i;
378
379        let reachable = reachable_to_output(tape);
380        let var_indices = tape.variables();
381        // depends_on[k_idx][i] — does slot i depend on var_indices[k_idx]?
382        let depends_on: Vec<Vec<bool>> = (0..var_indices.len())
383            .map(|k_idx| depends_on_var(tape, var_indices[k_idx]))
384            .collect();
385
386        let mut consts: Vec<f64> = Vec::new();
387        let mut const_intern: HashMap<u64, u32> = HashMap::new();
388        let mut intern_const = |c: f64, consts: &mut Vec<f64>| -> u32 {
389            let bits = c.to_bits();
390            if let Some(&idx) = const_intern.get(&bits) {
391                return idx;
392            }
393            let idx = consts.len() as u32;
394            consts.push(c);
395            const_intern.insert(bits, idx);
396            idx
397        };
398
399        let mut ops: Vec<HOp> = Vec::new();
400
401        // ---- Forward pass ----
402        for (i, tape_op) in tape.ops.iter().enumerate() {
403            let i = i as u32;
404            let dst = v_slot(i);
405            let op = match *tape_op {
406                TapeOp::Const(c) => HOp::FwdLoadConst {
407                    dst,
408                    c_idx: intern_const(c, &mut consts),
409                },
410                TapeOp::Var(x_idx) => HOp::FwdLoadVar {
411                    dst,
412                    x_idx: x_idx as u32,
413                },
414                TapeOp::Add(a, b) => HOp::FwdAdd {
415                    dst,
416                    a: v_slot(a as u32),
417                    b: v_slot(b as u32),
418                },
419                TapeOp::Sub(a, b) => HOp::FwdSub {
420                    dst,
421                    a: v_slot(a as u32),
422                    b: v_slot(b as u32),
423                },
424                TapeOp::Mul(a, b) => HOp::FwdMul {
425                    dst,
426                    a: v_slot(a as u32),
427                    b: v_slot(b as u32),
428                },
429                TapeOp::Div(a, b) => HOp::FwdDiv {
430                    dst,
431                    a: v_slot(a as u32),
432                    b: v_slot(b as u32),
433                },
434                TapeOp::Pow(a, b) => HOp::FwdPow {
435                    dst,
436                    a: v_slot(a as u32),
437                    b: v_slot(b as u32),
438                },
439                TapeOp::Neg(a) => HOp::FwdNeg {
440                    dst,
441                    a: v_slot(a as u32),
442                },
443                TapeOp::Abs(a) => HOp::FwdAbs {
444                    dst,
445                    a: v_slot(a as u32),
446                },
447                TapeOp::Sqrt(a) => HOp::FwdSqrt {
448                    dst,
449                    a: v_slot(a as u32),
450                },
451                TapeOp::Exp(a) => HOp::FwdExp {
452                    dst,
453                    a: v_slot(a as u32),
454                },
455                TapeOp::Log(a) => HOp::FwdLog {
456                    dst,
457                    a: v_slot(a as u32),
458                },
459                TapeOp::Log10(a) => HOp::FwdLog10 {
460                    dst,
461                    a: v_slot(a as u32),
462                },
463                TapeOp::Sin(a) => HOp::FwdSin {
464                    dst,
465                    a: v_slot(a as u32),
466                },
467                TapeOp::Cos(a) => HOp::FwdCos {
468                    dst,
469                    a: v_slot(a as u32),
470                },
471                TapeOp::Funcall(_) => unreachable!(
472                    "HessianProgram path does not support AMPL external functions; \
473                     use the Tape (build_with_externals) path instead."
474                ),
475                TapeOp::Tan(_)
476                | TapeOp::Atan(_)
477                | TapeOp::Acos(_)
478                | TapeOp::Sinh(_)
479                | TapeOp::Cosh(_)
480                | TapeOp::Tanh(_)
481                | TapeOp::Asin(_)
482                | TapeOp::Acosh(_)
483                | TapeOp::Asinh(_)
484                | TapeOp::Erf(_)
485                | TapeOp::XLogX(_)
486                | TapeOp::CEntropy(_, _)
487                | TapeOp::Atanh(_)
488                | TapeOp::Atan2(_, _)
489                | TapeOp::Cmp(_, _, _)
490                | TapeOp::And(_, _)
491                | TapeOp::Or(_, _)
492                | TapeOp::Not(_)
493                | TapeOp::Select(_, _, _)
494                | TapeOp::Min(_, _)
495                | TapeOp::Max(_, _) => unreachable!(
496                    "HessianProgram path does not yet support tan/atan/acos, the \
497                     other transcendental opcodes including erf, atan2, min/max, or \
498                     conditional / logical opcodes; use the Tape \
499                     (build_with_externals) interpreter path instead."
500                ),
501            };
502            ops.push(op);
503        }
504
505        if n == 0 || var_indices.is_empty() {
506            return Some(HessianProgram {
507                ops,
508                consts,
509                n_slots,
510            });
511        }
512
513        // ---- Per-j forward-tangent + reverse-over-tangent ----
514        for (k_idx, &j) in var_indices.iter().enumerate() {
515            // Reset dot, adj, adj_dot for this j. Seed adj[n-1] = 1.
516            ops.push(HOp::ZeroRange {
517                start: dot_base,
518                len: 3 * n,
519            });
520            ops.push(HOp::SetOne {
521                dst: adj_slot(n - 1),
522            });
523
524            // Forward tangent: only emit ops for slots that
525            // statically depend on j (the rest stay zero from the
526            // ZeroRange above).
527            for (i, tape_op) in tape.ops.iter().enumerate() {
528                let i_u = i as u32;
529                if !depends_on[k_idx][i] {
530                    continue;
531                }
532                let dst = dot_slot(i_u);
533                let dot_op = match *tape_op {
534                    // Const: dot stays 0 (filtered above by
535                    // depends_on, since Const has no var-deps).
536                    TapeOp::Const(_) => continue,
537                    // Var(k): dot = 1 iff k == j, else 0. We only
538                    // get here if depends_on[k_idx][i] is true,
539                    // which for Var(k) means k == j.
540                    TapeOp::Var(_) => HOp::SetOne { dst },
541                    TapeOp::Add(a, b) => HOp::DotAdd {
542                        dst,
543                        a: dot_slot(a as u32),
544                        b: dot_slot(b as u32),
545                    },
546                    TapeOp::Sub(a, b) => HOp::DotSub {
547                        dst,
548                        a: dot_slot(a as u32),
549                        b: dot_slot(b as u32),
550                    },
551                    TapeOp::Mul(a, b) => HOp::DotMul {
552                        dst,
553                        dot_a: dot_slot(a as u32),
554                        vb: v_slot(b as u32),
555                        va: v_slot(a as u32),
556                        dot_b: dot_slot(b as u32),
557                    },
558                    TapeOp::Div(a, b) => HOp::DotDiv {
559                        dst,
560                        dot_a: dot_slot(a as u32),
561                        vb: v_slot(b as u32),
562                        va: v_slot(a as u32),
563                        dot_b: dot_slot(b as u32),
564                    },
565                    TapeOp::Pow(a, b) => HOp::DotPow {
566                        dst,
567                        va: v_slot(a as u32),
568                        vb: v_slot(b as u32),
569                        vd: v_slot(i_u),
570                        dot_a: dot_slot(a as u32),
571                        dot_b: dot_slot(b as u32),
572                    },
573                    TapeOp::Neg(a) => HOp::DotNeg {
574                        dst,
575                        dot_a: dot_slot(a as u32),
576                    },
577                    TapeOp::Abs(a) => HOp::DotAbs {
578                        dst,
579                        dot_a: dot_slot(a as u32),
580                        va: v_slot(a as u32),
581                    },
582                    TapeOp::Sqrt(a) => HOp::DotSqrt {
583                        dst,
584                        dot_a: dot_slot(a as u32),
585                        vd: v_slot(i_u),
586                    },
587                    TapeOp::Exp(a) => HOp::DotExp {
588                        dst,
589                        dot_a: dot_slot(a as u32),
590                        vd: v_slot(i_u),
591                    },
592                    TapeOp::Log(a) => HOp::DotLog {
593                        dst,
594                        dot_a: dot_slot(a as u32),
595                        va: v_slot(a as u32),
596                    },
597                    TapeOp::Log10(a) => HOp::DotLog10 {
598                        dst,
599                        dot_a: dot_slot(a as u32),
600                        va: v_slot(a as u32),
601                    },
602                    TapeOp::Sin(a) => HOp::DotSin {
603                        dst,
604                        dot_a: dot_slot(a as u32),
605                        va: v_slot(a as u32),
606                    },
607                    TapeOp::Cos(a) => HOp::DotCos {
608                        dst,
609                        dot_a: dot_slot(a as u32),
610                        va: v_slot(a as u32),
611                    },
612                    TapeOp::Funcall(_) => unreachable!(
613                        "HessianProgram path does not support AMPL external functions; \
614                         use the Tape (build_with_externals) path instead."
615                    ),
616                    TapeOp::Tan(_)
617                    | TapeOp::Atan(_)
618                    | TapeOp::Acos(_)
619                    | TapeOp::Sinh(_)
620                    | TapeOp::Cosh(_)
621                    | TapeOp::Tanh(_)
622                    | TapeOp::Asin(_)
623                    | TapeOp::Acosh(_)
624                    | TapeOp::Asinh(_)
625                    | TapeOp::Erf(_)
626                    | TapeOp::XLogX(_)
627                    | TapeOp::CEntropy(_, _)
628                    | TapeOp::Atanh(_)
629                    | TapeOp::Atan2(_, _)
630                    | TapeOp::Cmp(_, _, _)
631                    | TapeOp::And(_, _)
632                    | TapeOp::Or(_, _)
633                    | TapeOp::Not(_)
634                    | TapeOp::Select(_, _, _)
635                    | TapeOp::Min(_, _)
636                    | TapeOp::Max(_, _) => unreachable!(
637                        "HessianProgram path does not yet support tan/atan/acos, the \
638                         other transcendental opcodes including erf, atan2, min/max, or \
639                         conditional / logical opcodes; use the Tape \
640                         (build_with_externals) interpreter path instead."
641                    ),
642                };
643                ops.push(dot_op);
644            }
645
646            // Reverse-over-tangent: walk slots backward, emit only
647            // for reachable slots.
648            for i in (0..n as usize).rev() {
649                if !reachable[i] {
650                    continue;
651                }
652                let i_u = i as u32;
653                let w = adj_slot(i_u);
654                let wd = adj_dot_slot(i_u);
655                let tape_op = &tape.ops[i];
656                let rev_op = match *tape_op {
657                    TapeOp::Const(_) => continue,
658                    TapeOp::Var(k) => {
659                        // At a Var slot: if k >= j and hess_map has
660                        // an entry for (k, j), emit a HessEmit op.
661                        // No adj/adj_dot propagation (no operands).
662                        if k >= j {
663                            if let Some(&ptr) = hess_map.get(&(k, j)) {
664                                ops.push(HOp::HessEmit {
665                                    hess_ptr: ptr as u32,
666                                    adj_dot_slot: wd,
667                                });
668                            }
669                        }
670                        continue;
671                    }
672                    TapeOp::Add(a, b) => HOp::RevAdd {
673                        adj_a: adj_slot(a as u32),
674                        adj_b: adj_slot(b as u32),
675                        adj_dot_a: adj_dot_slot(a as u32),
676                        adj_dot_b: adj_dot_slot(b as u32),
677                        w,
678                        wd,
679                    },
680                    TapeOp::Sub(a, b) => HOp::RevSub {
681                        adj_a: adj_slot(a as u32),
682                        adj_b: adj_slot(b as u32),
683                        adj_dot_a: adj_dot_slot(a as u32),
684                        adj_dot_b: adj_dot_slot(b as u32),
685                        w,
686                        wd,
687                    },
688                    TapeOp::Mul(a, b) => HOp::RevMul {
689                        adj_a: adj_slot(a as u32),
690                        adj_b: adj_slot(b as u32),
691                        adj_dot_a: adj_dot_slot(a as u32),
692                        adj_dot_b: adj_dot_slot(b as u32),
693                        w,
694                        wd,
695                        va: v_slot(a as u32),
696                        vb: v_slot(b as u32),
697                        dot_a: dot_slot(a as u32),
698                        dot_b: dot_slot(b as u32),
699                    },
700                    TapeOp::Div(a, b) => HOp::RevDiv {
701                        adj_a: adj_slot(a as u32),
702                        adj_b: adj_slot(b as u32),
703                        adj_dot_a: adj_dot_slot(a as u32),
704                        adj_dot_b: adj_dot_slot(b as u32),
705                        w,
706                        wd,
707                        va: v_slot(a as u32),
708                        vb: v_slot(b as u32),
709                        dot_a: dot_slot(a as u32),
710                        dot_b: dot_slot(b as u32),
711                    },
712                    TapeOp::Pow(a, b) => HOp::RevPow {
713                        adj_a: adj_slot(a as u32),
714                        adj_b: adj_slot(b as u32),
715                        adj_dot_a: adj_dot_slot(a as u32),
716                        adj_dot_b: adj_dot_slot(b as u32),
717                        w,
718                        wd,
719                        va: v_slot(a as u32),
720                        vb: v_slot(b as u32),
721                        vd: v_slot(i_u),
722                        dot_a: dot_slot(a as u32),
723                        dot_b: dot_slot(b as u32),
724                    },
725                    TapeOp::Neg(a) => HOp::RevNeg {
726                        adj_a: adj_slot(a as u32),
727                        adj_dot_a: adj_dot_slot(a as u32),
728                        w,
729                        wd,
730                    },
731                    TapeOp::Abs(a) => HOp::RevAbs {
732                        adj_a: adj_slot(a as u32),
733                        adj_dot_a: adj_dot_slot(a as u32),
734                        w,
735                        wd,
736                        va: v_slot(a as u32),
737                    },
738                    TapeOp::Sqrt(a) => HOp::RevSqrt {
739                        adj_a: adj_slot(a as u32),
740                        adj_dot_a: adj_dot_slot(a as u32),
741                        w,
742                        wd,
743                        va: v_slot(a as u32),
744                        vd: v_slot(i_u),
745                        dot_a: dot_slot(a as u32),
746                    },
747                    TapeOp::Exp(a) => HOp::RevExp {
748                        adj_a: adj_slot(a as u32),
749                        adj_dot_a: adj_dot_slot(a as u32),
750                        w,
751                        wd,
752                        vd: v_slot(i_u),
753                        dot_a: dot_slot(a as u32),
754                    },
755                    TapeOp::Log(a) => HOp::RevLog {
756                        adj_a: adj_slot(a as u32),
757                        adj_dot_a: adj_dot_slot(a as u32),
758                        w,
759                        wd,
760                        va: v_slot(a as u32),
761                        dot_a: dot_slot(a as u32),
762                    },
763                    TapeOp::Log10(a) => HOp::RevLog10 {
764                        adj_a: adj_slot(a as u32),
765                        adj_dot_a: adj_dot_slot(a as u32),
766                        w,
767                        wd,
768                        va: v_slot(a as u32),
769                        dot_a: dot_slot(a as u32),
770                    },
771                    TapeOp::Sin(a) => HOp::RevSin {
772                        adj_a: adj_slot(a as u32),
773                        adj_dot_a: adj_dot_slot(a as u32),
774                        w,
775                        wd,
776                        va: v_slot(a as u32),
777                        dot_a: dot_slot(a as u32),
778                    },
779                    TapeOp::Cos(a) => HOp::RevCos {
780                        adj_a: adj_slot(a as u32),
781                        adj_dot_a: adj_dot_slot(a as u32),
782                        w,
783                        wd,
784                        va: v_slot(a as u32),
785                        dot_a: dot_slot(a as u32),
786                    },
787                    TapeOp::Funcall(_) => unreachable!(
788                        "HessianProgram path does not support AMPL external functions; \
789                         use the Tape (build_with_externals) path instead."
790                    ),
791                    TapeOp::Tan(_)
792                    | TapeOp::Atan(_)
793                    | TapeOp::Acos(_)
794                    | TapeOp::Sinh(_)
795                    | TapeOp::Cosh(_)
796                    | TapeOp::Tanh(_)
797                    | TapeOp::Asin(_)
798                    | TapeOp::Acosh(_)
799                    | TapeOp::Asinh(_)
800                    | TapeOp::Erf(_)
801                    | TapeOp::XLogX(_)
802                    | TapeOp::CEntropy(_, _)
803                    | TapeOp::Atanh(_)
804                    | TapeOp::Atan2(_, _)
805                    | TapeOp::Cmp(_, _, _)
806                    | TapeOp::And(_, _)
807                    | TapeOp::Or(_, _)
808                    | TapeOp::Not(_)
809                    | TapeOp::Select(_, _, _)
810                    | TapeOp::Min(_, _)
811                    | TapeOp::Max(_, _) => unreachable!(
812                        "HessianProgram path does not yet support tan/atan/acos, the \
813                         other transcendental opcodes including erf, atan2, min/max, or \
814                         conditional / logical opcodes; use the Tape \
815                         (build_with_externals) interpreter path instead."
816                    ),
817                };
818                ops.push(rev_op);
819            }
820        }
821
822        Some(HessianProgram {
823            ops,
824            consts,
825            n_slots,
826        })
827    }
828
829    pub fn n_slots(&self) -> usize {
830        self.n_slots as usize
831    }
832
833    pub fn n_ops(&self) -> usize {
834        self.ops.len()
835    }
836
837    /// Execute the program. `scratch` is overwritten throughout;
838    /// it must be at least [`n_slots`] long. `values` is the
839    /// shared Hessian-values buffer the caller is accumulating
840    /// into (same semantics as
841    /// [`Tape::hessian_accumulate`]'s `values`).
842    pub fn execute(&self, x: &[f64], weight: f64, scratch: &mut [f64], values: &mut [f64]) {
843        debug_assert!(scratch.len() >= self.n_slots as usize);
844        if self.ops.is_empty() || weight == 0.0 {
845            return;
846        }
847        let consts = &self.consts[..];
848        for &op in &self.ops {
849            match op {
850                HOp::FwdLoadVar { dst, x_idx } => {
851                    scratch[dst as usize] = x[x_idx as usize];
852                }
853                HOp::FwdLoadConst { dst, c_idx } => {
854                    scratch[dst as usize] = consts[c_idx as usize];
855                }
856                HOp::FwdAdd { dst, a, b } => {
857                    scratch[dst as usize] = scratch[a as usize] + scratch[b as usize];
858                }
859                HOp::FwdSub { dst, a, b } => {
860                    scratch[dst as usize] = scratch[a as usize] - scratch[b as usize];
861                }
862                HOp::FwdMul { dst, a, b } => {
863                    scratch[dst as usize] = scratch[a as usize] * scratch[b as usize];
864                }
865                HOp::FwdDiv { dst, a, b } => {
866                    scratch[dst as usize] = scratch[a as usize] / scratch[b as usize];
867                }
868                HOp::FwdPow { dst, a, b } => {
869                    scratch[dst as usize] = scratch[a as usize].powf(scratch[b as usize]);
870                }
871                HOp::FwdNeg { dst, a } => {
872                    scratch[dst as usize] = -scratch[a as usize];
873                }
874                HOp::FwdAbs { dst, a } => {
875                    scratch[dst as usize] = scratch[a as usize].abs();
876                }
877                HOp::FwdSqrt { dst, a } => {
878                    scratch[dst as usize] = scratch[a as usize].sqrt();
879                }
880                HOp::FwdExp { dst, a } => {
881                    scratch[dst as usize] = scratch[a as usize].exp();
882                }
883                HOp::FwdLog { dst, a } => {
884                    scratch[dst as usize] = scratch[a as usize].ln();
885                }
886                HOp::FwdLog10 { dst, a } => {
887                    scratch[dst as usize] = scratch[a as usize].log10();
888                }
889                HOp::FwdSin { dst, a } => {
890                    scratch[dst as usize] = scratch[a as usize].sin();
891                }
892                HOp::FwdCos { dst, a } => {
893                    scratch[dst as usize] = scratch[a as usize].cos();
894                }
895
896                HOp::SetZero { dst } => {
897                    scratch[dst as usize] = 0.0;
898                }
899                HOp::SetOne { dst } => {
900                    scratch[dst as usize] = 1.0;
901                }
902                HOp::ZeroRange { start, len } => {
903                    let s = start as usize;
904                    let e = s + len as usize;
905                    scratch[s..e].fill(0.0);
906                }
907
908                HOp::DotAdd { dst, a, b } => {
909                    scratch[dst as usize] = scratch[a as usize] + scratch[b as usize];
910                }
911                HOp::DotSub { dst, a, b } => {
912                    scratch[dst as usize] = scratch[a as usize] - scratch[b as usize];
913                }
914                HOp::DotMul {
915                    dst,
916                    dot_a,
917                    vb,
918                    va,
919                    dot_b,
920                } => {
921                    scratch[dst as usize] = scratch[dot_a as usize] * scratch[vb as usize]
922                        + scratch[va as usize] * scratch[dot_b as usize];
923                }
924                HOp::DotDiv {
925                    dst,
926                    dot_a,
927                    vb,
928                    va,
929                    dot_b,
930                } => {
931                    // Kahan form, as `Tape`'s Div tangent: `(ȧ - q·ḃ)/b` with
932                    // `q = a/b`, never `(ȧ·b - a·ḃ)/b²`, whose `b*b` overflows
933                    // for |b| > 1.3e154 and underflows below 1.5e-154.
934                    let v_b = scratch[vb as usize];
935                    let q = scratch[va as usize] / v_b;
936                    scratch[dst as usize] =
937                        (scratch[dot_a as usize] - q * scratch[dot_b as usize]) / v_b;
938                }
939                HOp::DotSqrt { dst, dot_a, vd } => {
940                    let svd = scratch[vd as usize];
941                    scratch[dst as usize] = if svd > 0.0 {
942                        scratch[dot_a as usize] * 0.5 / svd
943                    } else {
944                        0.0
945                    };
946                }
947                HOp::DotExp { dst, dot_a, vd } => {
948                    scratch[dst as usize] = scratch[dot_a as usize] * scratch[vd as usize];
949                }
950                HOp::DotLog { dst, dot_a, va } => {
951                    scratch[dst as usize] = scratch[dot_a as usize] / scratch[va as usize];
952                }
953                HOp::DotLog10 { dst, dot_a, va } => {
954                    scratch[dst as usize] =
955                        div_u_ln10(scratch[dot_a as usize], scratch[va as usize]);
956                }
957                HOp::DotSin { dst, dot_a, va } => {
958                    scratch[dst as usize] = scratch[dot_a as usize] * scratch[va as usize].cos();
959                }
960                HOp::DotCos { dst, dot_a, va } => {
961                    scratch[dst as usize] = -scratch[dot_a as usize] * scratch[va as usize].sin();
962                }
963                HOp::DotNeg { dst, dot_a } => {
964                    scratch[dst as usize] = -scratch[dot_a as usize];
965                }
966                HOp::DotAbs { dst, dot_a, va } => {
967                    scratch[dst as usize] = if scratch[va as usize] >= 0.0 {
968                        scratch[dot_a as usize]
969                    } else {
970                        -scratch[dot_a as usize]
971                    };
972                }
973                HOp::DotPow {
974                    dst,
975                    va,
976                    vb,
977                    vd,
978                    dot_a,
979                    dot_b,
980                } => {
981                    let u = scratch[va as usize];
982                    let r = scratch[vb as usize];
983                    let du = scratch[dot_a as usize];
984                    let dr = scratch[dot_b as usize];
985                    let mut result = 0.0;
986                    if r != 0.0 && u != 0.0 {
987                        result += r * u.powf(r - 1.0) * du;
988                    }
989                    if u > 0.0 {
990                        result += scratch[vd as usize] * u.ln() * dr;
991                    }
992                    scratch[dst as usize] = result;
993                }
994
995                HOp::RevAdd {
996                    adj_a,
997                    adj_b,
998                    adj_dot_a,
999                    adj_dot_b,
1000                    w,
1001                    wd,
1002                } => {
1003                    let w_v = scratch[w as usize];
1004                    let wd_v = scratch[wd as usize];
1005                    scratch[adj_a as usize] += w_v;
1006                    scratch[adj_b as usize] += w_v;
1007                    scratch[adj_dot_a as usize] += wd_v;
1008                    scratch[adj_dot_b as usize] += wd_v;
1009                }
1010                HOp::RevSub {
1011                    adj_a,
1012                    adj_b,
1013                    adj_dot_a,
1014                    adj_dot_b,
1015                    w,
1016                    wd,
1017                } => {
1018                    let w_v = scratch[w as usize];
1019                    let wd_v = scratch[wd as usize];
1020                    scratch[adj_a as usize] += w_v;
1021                    scratch[adj_b as usize] -= w_v;
1022                    scratch[adj_dot_a as usize] += wd_v;
1023                    scratch[adj_dot_b as usize] -= wd_v;
1024                }
1025                HOp::RevMul {
1026                    adj_a,
1027                    adj_b,
1028                    adj_dot_a,
1029                    adj_dot_b,
1030                    w,
1031                    wd,
1032                    va,
1033                    vb,
1034                    dot_a,
1035                    dot_b,
1036                } => {
1037                    let w_v = scratch[w as usize];
1038                    let wd_v = scratch[wd as usize];
1039                    let va_v = scratch[va as usize];
1040                    let vb_v = scratch[vb as usize];
1041                    let da_v = scratch[dot_a as usize];
1042                    let db_v = scratch[dot_b as usize];
1043                    scratch[adj_a as usize] += w_v * vb_v;
1044                    scratch[adj_b as usize] += w_v * va_v;
1045                    scratch[adj_dot_a as usize] += wd_v * vb_v + w_v * db_v;
1046                    scratch[adj_dot_b as usize] += wd_v * va_v + w_v * da_v;
1047                }
1048                HOp::RevDiv {
1049                    adj_a,
1050                    adj_b,
1051                    adj_dot_a,
1052                    adj_dot_b,
1053                    w,
1054                    wd,
1055                    va,
1056                    vb,
1057                    dot_a,
1058                    dot_b,
1059                } => {
1060                    let w_v = scratch[w as usize];
1061                    let wd_v = scratch[wd as usize];
1062                    let va_v = scratch[va as usize];
1063                    let vb_v = scratch[vb as usize];
1064                    // Kahan form, as `Tape::hessian_accumulate`'s Div arm: every
1065                    // 1/b² regrouped as (·/b)/b, so nothing squares or cubes
1066                    // `b` (vb2/vb3 overflowed for |b| > 1.3e154 and zeroed
1067                    // representable second-order terms). `q`/`qd` are the
1068                    // node's value and tangent, recomputed exactly as the
1069                    // forward and tangent sweeps produce them.
1070                    let da_v = scratch[dot_a as usize];
1071                    let db_v = scratch[dot_b as usize];
1072                    let q = va_v / vb_v;
1073                    let qd = (da_v - q * db_v) / vb_v;
1074                    scratch[adj_a as usize] += w_v / vb_v;
1075                    scratch[adj_dot_a as usize] += wd_v / vb_v - w_v * (db_v / vb_v) / vb_v;
1076                    scratch[adj_b as usize] -= w_v * q / vb_v;
1077                    scratch[adj_dot_b as usize] +=
1078                        -(wd_v * q) / vb_v + (w_v / vb_v) * (-qd + q * (db_v / vb_v));
1079                }
1080                HOp::RevPow {
1081                    adj_a,
1082                    adj_b,
1083                    adj_dot_a,
1084                    adj_dot_b,
1085                    w,
1086                    wd,
1087                    va,
1088                    vb,
1089                    vd,
1090                    dot_a,
1091                    dot_b,
1092                } => {
1093                    let w_v = scratch[w as usize];
1094                    let wd_v = scratch[wd as usize];
1095                    let u = scratch[va as usize];
1096                    let r = scratch[vb as usize];
1097                    let du = scratch[dot_a as usize];
1098                    let dr = scratch[dot_b as usize];
1099                    if r != 0.0 {
1100                        if u != 0.0 {
1101                            let p_a = r * u.powf(r - 1.0);
1102                            scratch[adj_a as usize] += w_v * p_a;
1103                            let mut dp_a = dr * u.powf(r - 1.0);
1104                            if u > 0.0 {
1105                                dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1106                            } else {
1107                                dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1108                            }
1109                            scratch[adj_dot_a as usize] += wd_v * p_a + w_v * dp_a;
1110                        } else if r >= 2.0 {
1111                            let p_a = 0.0;
1112                            scratch[adj_a as usize] += w_v * p_a;
1113                            let dp_a = if r == 2.0 {
1114                                2.0 * du
1115                            } else {
1116                                r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1117                            };
1118                            scratch[adj_dot_a as usize] += wd_v * p_a + w_v * dp_a;
1119                        }
1120                    }
1121                    if u > 0.0 {
1122                        let ln_u = u.ln();
1123                        let p_b = scratch[vd as usize] * ln_u;
1124                        scratch[adj_b as usize] += w_v * p_b;
1125                        let dur = scratch[vd as usize] * (r * du / u + dr * ln_u);
1126                        let dp_b = dur * ln_u + scratch[vd as usize] * du / u;
1127                        scratch[adj_dot_b as usize] += wd_v * p_b + w_v * dp_b;
1128                    }
1129                }
1130                HOp::RevNeg {
1131                    adj_a,
1132                    adj_dot_a,
1133                    w,
1134                    wd,
1135                } => {
1136                    scratch[adj_a as usize] -= scratch[w as usize];
1137                    scratch[adj_dot_a as usize] -= scratch[wd as usize];
1138                }
1139                HOp::RevAbs {
1140                    adj_a,
1141                    adj_dot_a,
1142                    w,
1143                    wd,
1144                    va,
1145                } => {
1146                    let s = if scratch[va as usize] >= 0.0 {
1147                        1.0
1148                    } else {
1149                        -1.0
1150                    };
1151                    scratch[adj_a as usize] += scratch[w as usize] * s;
1152                    scratch[adj_dot_a as usize] += scratch[wd as usize] * s;
1153                }
1154                HOp::RevSqrt {
1155                    adj_a,
1156                    adj_dot_a,
1157                    w,
1158                    wd,
1159                    va,
1160                    vd,
1161                    dot_a,
1162                } => {
1163                    let sv = scratch[vd as usize];
1164                    if sv > 0.0 {
1165                        let fp = 0.5 / sv;
1166                        // Shared with `Tape`: `sv*sv*sv` overflowed at
1167                        // u ≈ 2.6e205 where the true value is still ~1e-308.
1168                        let fpp = sqrt_d2(scratch[va as usize], sv);
1169                        let w_v = scratch[w as usize];
1170                        let wd_v = scratch[wd as usize];
1171                        scratch[adj_a as usize] += w_v * fp;
1172                        scratch[adj_dot_a as usize] +=
1173                            wd_v * fp + w_v * fpp * scratch[dot_a as usize];
1174                    }
1175                }
1176                HOp::RevExp {
1177                    adj_a,
1178                    adj_dot_a,
1179                    w,
1180                    wd,
1181                    vd,
1182                    dot_a,
1183                } => {
1184                    let ev = scratch[vd as usize];
1185                    let w_v = scratch[w as usize];
1186                    let wd_v = scratch[wd as usize];
1187                    scratch[adj_a as usize] += w_v * ev;
1188                    scratch[adj_dot_a as usize] += wd_v * ev + w_v * ev * scratch[dot_a as usize];
1189                }
1190                HOp::RevLog {
1191                    adj_a,
1192                    adj_dot_a,
1193                    w,
1194                    wd,
1195                    va,
1196                    dot_a,
1197                } => {
1198                    let u = scratch[va as usize];
1199                    let w_v = scratch[w as usize];
1200                    let wd_v = scratch[wd as usize];
1201                    scratch[adj_a as usize] += w_v / u;
1202                    scratch[adj_dot_a as usize] +=
1203                        wd_v / u + w_v * log_d2(u) * scratch[dot_a as usize];
1204                }
1205                HOp::RevLog10 {
1206                    adj_a,
1207                    adj_dot_a,
1208                    w,
1209                    wd,
1210                    va,
1211                    dot_a,
1212                } => {
1213                    let u = scratch[va as usize];
1214                    let w_v = scratch[w as usize];
1215                    let wd_v = scratch[wd as usize];
1216                    scratch[adj_a as usize] += div_u_ln10(w_v, u);
1217                    scratch[adj_dot_a as usize] +=
1218                        div_u_ln10(wd_v, u) + w_v * log10_d2(u) * scratch[dot_a as usize];
1219                }
1220                HOp::RevSin {
1221                    adj_a,
1222                    adj_dot_a,
1223                    w,
1224                    wd,
1225                    va,
1226                    dot_a,
1227                } => {
1228                    let u = scratch[va as usize];
1229                    let cu = u.cos();
1230                    let w_v = scratch[w as usize];
1231                    let wd_v = scratch[wd as usize];
1232                    scratch[adj_a as usize] += w_v * cu;
1233                    scratch[adj_dot_a as usize] +=
1234                        wd_v * cu + w_v * (-u.sin()) * scratch[dot_a as usize];
1235                }
1236                HOp::RevCos {
1237                    adj_a,
1238                    adj_dot_a,
1239                    w,
1240                    wd,
1241                    va,
1242                    dot_a,
1243                } => {
1244                    let u = scratch[va as usize];
1245                    let su = u.sin();
1246                    let w_v = scratch[w as usize];
1247                    let wd_v = scratch[wd as usize];
1248                    scratch[adj_a as usize] -= w_v * su;
1249                    scratch[adj_dot_a as usize] +=
1250                        wd_v * (-su) + w_v * (-u.cos()) * scratch[dot_a as usize];
1251                }
1252
1253                HOp::HessEmit {
1254                    hess_ptr,
1255                    adj_dot_slot,
1256                } => {
1257                    values[hess_ptr as usize] += weight * scratch[adj_dot_slot as usize];
1258                }
1259            }
1260        }
1261    }
1262}
1263
1264/// Whether [`HessianProgram::compile`] can lower a single opcode. The
1265/// program path covers smooth arithmetic plus `sin`/`cos`; every other
1266/// opcode — AMPL external `Funcall`, the remaining transcendentals
1267/// (`tan`/`atan`/`acos`/the hyperbolics/`asin`/`erf`…) and `atan2`, and the
1268/// `min`/`max`/conditional/logical family — is unsupported, so `compile`
1269/// returns `None` and the caller falls back to the `Tape`
1270/// (`build_with_externals`) interpreter path rather than panicking on user
1271/// input (code review L28). This is the single source of truth for the
1272/// supported set: the per-sweep match arms and the dependence/reachability
1273/// analyses are only ever reached with opcodes this predicate accepts, so
1274/// their unsupported branches are `unreachable!`.
1275fn program_supports_op(op: &TapeOp) -> bool {
1276    matches!(
1277        op,
1278        TapeOp::Const(_)
1279            | TapeOp::Var(_)
1280            | TapeOp::Add(_, _)
1281            | TapeOp::Sub(_, _)
1282            | TapeOp::Mul(_, _)
1283            | TapeOp::Div(_, _)
1284            | TapeOp::Pow(_, _)
1285            | TapeOp::Neg(_)
1286            | TapeOp::Abs(_)
1287            | TapeOp::Sqrt(_)
1288            | TapeOp::Exp(_)
1289            | TapeOp::Log(_)
1290            | TapeOp::Log10(_)
1291            | TapeOp::Sin(_)
1292            | TapeOp::Cos(_)
1293    )
1294}
1295
1296/// `out[i]` = does tape slot `i` contribute (transitively) to the
1297/// output slot `n-1`. Used to skip emitting reverse-pass ops for
1298/// dead slots.
1299fn reachable_to_output(tape: &Tape) -> Vec<bool> {
1300    let n = tape.ops.len();
1301    let mut r = vec![false; n];
1302    if n == 0 {
1303        return r;
1304    }
1305    r[n - 1] = true;
1306    for i in (0..n).rev() {
1307        if !r[i] {
1308            continue;
1309        }
1310        match tape.ops[i] {
1311            TapeOp::Const(_) | TapeOp::Var(_) => {}
1312            TapeOp::Add(a, b)
1313            | TapeOp::Sub(a, b)
1314            | TapeOp::Mul(a, b)
1315            | TapeOp::Div(a, b)
1316            | TapeOp::Pow(a, b)
1317            | TapeOp::Atan2(a, b)
1318            | TapeOp::CEntropy(a, b) => {
1319                r[a] = true;
1320                r[b] = true;
1321            }
1322            TapeOp::Neg(a)
1323            | TapeOp::Abs(a)
1324            | TapeOp::Sqrt(a)
1325            | TapeOp::Exp(a)
1326            | TapeOp::Log(a)
1327            | TapeOp::Log10(a)
1328            | TapeOp::Sin(a)
1329            | TapeOp::Cos(a)
1330            | TapeOp::Tan(a)
1331            | TapeOp::Atan(a)
1332            | TapeOp::Acos(a)
1333            | TapeOp::Sinh(a)
1334            | TapeOp::Cosh(a)
1335            | TapeOp::Tanh(a)
1336            | TapeOp::Asin(a)
1337            | TapeOp::Acosh(a)
1338            | TapeOp::Asinh(a)
1339            | TapeOp::Erf(a)
1340            | TapeOp::XLogX(a)
1341            | TapeOp::Atanh(a) => {
1342                r[a] = true;
1343            }
1344            TapeOp::Funcall(_) => unreachable!(
1345                "HessianProgram path does not support AMPL external functions; \
1346                 use the Tape (build_with_externals) path instead."
1347            ),
1348            TapeOp::Cmp(_, _, _)
1349            | TapeOp::And(_, _)
1350            | TapeOp::Or(_, _)
1351            | TapeOp::Not(_)
1352            | TapeOp::Select(_, _, _)
1353            | TapeOp::Min(_, _)
1354            | TapeOp::Max(_, _) => unreachable!(
1355                "HessianProgram path does not support conditional / logical / min-max \
1356                 opcodes; use the Tape (build_with_externals) path instead."
1357            ),
1358        }
1359    }
1360    r
1361}
1362
1363/// `out[i]` = does tape slot `i` transitively read from `Var(j)`.
1364/// Used to prune forward-tangent ops (slots with `out[i] = false`
1365/// have `dot[i] = 0` and the rest of the per-`j` pass can skip
1366/// them).
1367fn depends_on_var(tape: &Tape, j: usize) -> Vec<bool> {
1368    let n = tape.ops.len();
1369    let mut d = vec![false; n];
1370    for (i, op) in tape.ops.iter().enumerate() {
1371        d[i] = match *op {
1372            TapeOp::Const(_) => false,
1373            TapeOp::Var(k) => k == j,
1374            TapeOp::Add(a, b)
1375            | TapeOp::Sub(a, b)
1376            | TapeOp::Mul(a, b)
1377            | TapeOp::Div(a, b)
1378            | TapeOp::Pow(a, b)
1379            | TapeOp::CEntropy(a, b)
1380            | TapeOp::Atan2(a, b) => d[a] || d[b],
1381            TapeOp::Neg(a)
1382            | TapeOp::Abs(a)
1383            | TapeOp::Sqrt(a)
1384            | TapeOp::Exp(a)
1385            | TapeOp::Log(a)
1386            | TapeOp::Log10(a)
1387            | TapeOp::Sin(a)
1388            | TapeOp::Cos(a)
1389            | TapeOp::Tan(a)
1390            | TapeOp::Atan(a)
1391            | TapeOp::Acos(a)
1392            | TapeOp::Sinh(a)
1393            | TapeOp::Cosh(a)
1394            | TapeOp::Tanh(a)
1395            | TapeOp::Asin(a)
1396            | TapeOp::Acosh(a)
1397            | TapeOp::Asinh(a)
1398            | TapeOp::XLogX(a)
1399            | TapeOp::Erf(a)
1400            | TapeOp::Atanh(a) => d[a],
1401            TapeOp::Funcall(_) => unreachable!(
1402                "HessianProgram path does not support AMPL external functions; \
1403                 use the Tape (build_with_externals) path instead."
1404            ),
1405            TapeOp::Cmp(_, _, _)
1406            | TapeOp::And(_, _)
1407            | TapeOp::Or(_, _)
1408            | TapeOp::Not(_)
1409            | TapeOp::Select(_, _, _)
1410            | TapeOp::Min(_, _)
1411            | TapeOp::Max(_, _) => unreachable!(
1412                "HessianProgram path does not support conditional / logical / min-max \
1413                 opcodes; use the Tape (build_with_externals) path instead."
1414            ),
1415        };
1416    }
1417    d
1418}
1419
1420#[cfg(test)]
1421mod tests {
1422    use super::*;
1423    use crate::nl_reader::{BinOp, Expr, UnaryOp};
1424    use std::collections::BTreeSet;
1425    use std::sync::Arc;
1426
1427    fn cnst(c: f64) -> Expr {
1428        Expr::Const(c)
1429    }
1430    fn var(i: usize) -> Expr {
1431        Expr::Var(i)
1432    }
1433    fn add(a: Expr, b: Expr) -> Expr {
1434        Expr::Binary(BinOp::Add, Box::new(a), Box::new(b))
1435    }
1436    fn mul(a: Expr, b: Expr) -> Expr {
1437        Expr::Binary(BinOp::Mul, Box::new(a), Box::new(b))
1438    }
1439    fn pow(a: Expr, b: Expr) -> Expr {
1440        Expr::Binary(BinOp::Pow, Box::new(a), Box::new(b))
1441    }
1442    fn div(a: Expr, b: Expr) -> Expr {
1443        Expr::Binary(BinOp::Div, Box::new(a), Box::new(b))
1444    }
1445    fn sub(a: Expr, b: Expr) -> Expr {
1446        Expr::Binary(BinOp::Sub, Box::new(a), Box::new(b))
1447    }
1448    fn unary(op: UnaryOp, a: Expr) -> Expr {
1449        Expr::Unary(op, Box::new(a))
1450    }
1451
1452    /// Build the same shared (row, col) -> pos map both AD paths
1453    /// scatter into. Lower-triangle pairs, in tape.variables() order.
1454    fn build_hess_map(tape: &Tape) -> (HashMap<(usize, usize), usize>, Vec<(usize, usize)>) {
1455        let vars = tape.variables();
1456        let mut pairs: Vec<(usize, usize)> = Vec::new();
1457        let mut map: HashMap<(usize, usize), usize> = HashMap::new();
1458        for (ai, &vi) in vars.iter().enumerate() {
1459            for &vj in &vars[..=ai] {
1460                let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
1461                map.entry((r, c)).or_insert_with(|| {
1462                    let p = pairs.len();
1463                    pairs.push((r, c));
1464                    p
1465                });
1466            }
1467        }
1468        (map, pairs)
1469    }
1470
1471    /// Run both implementations against the same input and assert
1472    /// values match to a tight ULP-aligned tolerance.
1473    fn assert_program_matches_tape(tape: &Tape, x: &[f64], weight: f64) {
1474        let (hess_map, pairs) = build_hess_map(tape);
1475        let nnz = pairs.len();
1476
1477        let mut tape_vals = vec![0.0; nnz];
1478        tape.hessian_accumulate(x, weight, &hess_map, &mut tape_vals);
1479
1480        let program =
1481            HessianProgram::compile(tape, &hess_map).expect("tape uses only supported opcodes");
1482        let mut scratch = vec![0.0; program.n_slots()];
1483        let mut prog_vals = vec![0.0; nnz];
1484        program.execute(x, weight, &mut scratch, &mut prog_vals);
1485
1486        for (k, &(r, c)) in pairs.iter().enumerate() {
1487            let tol = tape_vals[k].abs().max(1.0) * 1e-12;
1488            assert!(
1489                (tape_vals[k] - prog_vals[k]).abs() < tol,
1490                "H[{},{}]: tape={:.6e} prog={:.6e}",
1491                r,
1492                c,
1493                tape_vals[k],
1494                prog_vals[k]
1495            );
1496        }
1497    }
1498
1499    #[test]
1500    fn matches_quadratic() {
1501        let e = add(
1502            add(
1503                mul(cnst(3.0), pow(var(0), cnst(2.0))),
1504                mul(cnst(2.0), mul(var(0), var(1))),
1505            ),
1506            pow(var(1), cnst(2.0)),
1507        );
1508        let tape = Tape::build(&e);
1509        assert_program_matches_tape(&tape, &[2.0, 3.0], 1.0);
1510        assert_program_matches_tape(&tape, &[-1.5, 0.7], 2.5);
1511    }
1512
1513    #[test]
1514    fn matches_transcendental() {
1515        let e = Expr::Sum(vec![
1516            unary(UnaryOp::Exp, var(0)),
1517            unary(UnaryOp::Sin, var(1)),
1518            unary(UnaryOp::Log, var(0)),
1519            unary(UnaryOp::Sqrt, var(1)),
1520            mul(var(0), var(1)),
1521            unary(UnaryOp::Cos, add(var(0), var(1))),
1522        ]);
1523        let tape = Tape::build(&e);
1524        assert_program_matches_tape(&tape, &[1.5, 2.0], 1.0);
1525        assert_program_matches_tape(&tape, &[0.3, 4.1], -0.4);
1526    }
1527
1528    #[test]
1529    fn matches_division() {
1530        let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
1531        let tape = Tape::build(&e);
1532        assert_program_matches_tape(&tape, &[0.5, 1.2], 1.0);
1533    }
1534
1535    #[test]
1536    fn division_log_log10_sqrt_accurate_at_extreme_magnitudes() {
1537        // This program carried the pre-fix forms after `Tape` was fixed:
1538        // division's `vb*vb`/`vb2*vb` (and `/(v_b*v_b)` in its tangent),
1539        // `log''`'s `-1/(u*u)`, `log10`'s `u*c`/`u*u*c`, and `sqrt''`'s
1540        // `sv*sv*sv`. Each overflowed and read `0` where the true entry is
1541        // representable. `assert_program_matches_tape`'s absolute 1e-12
1542        // floor cannot see that at these magnitudes, so this compares each
1543        // entry with a 100-digit reference, relatively.
1544        let e = Expr::Sum(vec![
1545            div(var(0), var(1)),
1546            div(var(2), var(3)),
1547            unary(UnaryOp::Log, var(4)),
1548            unary(UnaryOp::Log10, var(5)),
1549            unary(UnaryOp::Sqrt, var(6)),
1550        ]);
1551        let tape = Tape::build(&e);
1552        let x = [1e300, 1e200, 1e-10, 1e155, 1e155, 1e155, 1e206];
1553        let want: HashMap<(usize, usize), f64> = [
1554            ((1, 1), 2.0000000000000004e-300), // 2a/b³, b² = 1e400 overflows
1555            ((3, 2), -1e-310),                 // -1/b², b = 1e155
1556            ((4, 4), -1e-310),                 // log''
1557            ((5, 5), -4.3429448190326e-311),   // log10''
1558            ((6, 6), -2.5e-310),               // sqrt''
1559        ]
1560        .into_iter()
1561        .collect();
1562        let (hess_map, pairs) = build_hess_map(&tape);
1563        let program =
1564            HessianProgram::compile(&tape, &hess_map).expect("tape uses only supported opcodes");
1565        let mut scratch = vec![0.0; program.n_slots()];
1566        let mut prog_vals = vec![0.0; pairs.len()];
1567        program.execute(&x, 1.0, &mut scratch, &mut prog_vals);
1568        let mut tape_vals = vec![0.0; pairs.len()];
1569        tape.hessian_accumulate(&x, 1.0, &hess_map, &mut tape_vals);
1570        for (k, &rc) in pairs.iter().enumerate() {
1571            let w = want.get(&rc).copied().unwrap_or(0.0);
1572            for (who, got) in [("program", prog_vals[k]), ("tape", tape_vals[k])] {
1573                let ok = if w == 0.0 {
1574                    got == 0.0
1575                } else if w.abs() < f64::MIN_POSITIVE {
1576                    (got - w).abs() <= f64::from_bits(1)
1577                } else {
1578                    ((got - w) / w).abs() < 2e-15
1579                };
1580                assert!(ok, "{who} H{rc:?}: got {got:.17e}, want {w:.17e}");
1581            }
1582        }
1583    }
1584
1585    #[test]
1586    fn matches_through_cse() {
1587        let body = Arc::new(add(var(0), var(1)));
1588        let e = add(
1589            pow(Expr::Cse(body.clone()), cnst(2.0)),
1590            Expr::Cse(body.clone()),
1591        );
1592        let tape = Tape::build(&e);
1593        assert_program_matches_tape(&tape, &[1.0, 2.0], 1.0);
1594        assert_program_matches_tape(&tape, &[-0.5, 3.3], 0.7);
1595    }
1596
1597    #[test]
1598    fn matches_pow_chain() {
1599        // After Tier 1 this lowers to a Mul chain; verify both
1600        // paths agree on the lowered form too.
1601        let e = add(pow(var(0), cnst(3.0)), pow(var(1), cnst(-2.0)));
1602        let tape = Tape::build(&e);
1603        assert_program_matches_tape(&tape, &[1.7, 0.8], 1.0);
1604    }
1605
1606    #[test]
1607    fn matches_residual_pow_with_var_exponent() {
1608        // Pow where the exponent is variable (not constant), so
1609        // it survives Tier 1 and exercises the RevPow / DotPow
1610        // compound branches.
1611        let e = pow(var(0), var(1));
1612        let tape = Tape::build(&e);
1613        assert_program_matches_tape(&tape, &[2.5, 1.4], 1.0);
1614        assert_program_matches_tape(&tape, &[0.6, 2.1], -1.0);
1615    }
1616
1617    #[test]
1618    fn matches_sub_neg_abs() {
1619        let e = sub(
1620            unary(UnaryOp::Neg, var(0)),
1621            unary(UnaryOp::Abs, sub(var(1), var(0))),
1622        );
1623        let tape = Tape::build(&e);
1624        assert_program_matches_tape(&tape, &[1.0, -2.0], 1.0);
1625        assert_program_matches_tape(&tape, &[-3.5, 4.0], 0.9);
1626    }
1627
1628    #[test]
1629    fn slots_layout_matches_design() {
1630        let e = mul(var(0), var(1));
1631        let tape = Tape::build(&e);
1632        let (hess_map, _) = build_hess_map(&tape);
1633        let prog = HessianProgram::compile(&tape, &hess_map).expect("mul tape is supported");
1634        assert_eq!(prog.n_slots(), 4 * tape.ops.len());
1635    }
1636
1637    /// Sanity: the pruning analyses are consistent with the slot
1638    /// structure exposed via `hessian_sparsity()`.
1639    #[test]
1640    fn dependence_matches_hessian_sparsity_for_simple_case() {
1641        let e = add(unary(UnaryOp::Sin, var(0)), mul(var(1), var(2)));
1642        let tape = Tape::build(&e);
1643        let s: BTreeSet<(usize, usize)> = tape.hessian_sparsity();
1644        // (0,0) from sin, (2,1) from x1*x2, (1,1)/(2,2) NOT there
1645        // because Mul(x1, x2) emits cross only.
1646        assert!(s.contains(&(0, 0)));
1647        assert!(s.contains(&(2, 1)));
1648        assert_program_matches_tape(&tape, &[0.7, 1.1, 2.2], 1.0);
1649    }
1650
1651    /// Code review L28: a tape using an opcode the program path cannot
1652    /// lower (here `tan`) must make `compile` return `None` — a graceful
1653    /// fall-back-to-the-`Tape`-path signal — rather than panic. Previously
1654    /// the per-sweep match arms `panic!`'d on such ops, which would crash
1655    /// the process on arbitrary user `.nl` input if this path were ever
1656    /// wired into dispatch.
1657    #[test]
1658    fn unsupported_opcode_returns_none_instead_of_panicking() {
1659        // `tan(x0)` lowers to `TapeOp::Tan`, which the HessianProgram
1660        // compiler does not support.
1661        let e = unary(UnaryOp::Tan, var(0));
1662        let tape = Tape::build(&e);
1663        let (hess_map, _) = build_hess_map(&tape);
1664        assert!(
1665            HessianProgram::compile(&tape, &hess_map).is_none(),
1666            "tan() tape must fall back (None), not compile"
1667        );
1668
1669        // A fully-supported tape still compiles to `Some`, so the guard
1670        // rejects only genuinely-unsupported ops.
1671        let ok = mul(var(0), var(1));
1672        let ok_tape = Tape::build(&ok);
1673        let (ok_map, _) = build_hess_map(&ok_tape);
1674        assert!(
1675            HessianProgram::compile(&ok_tape, &ok_map).is_some(),
1676            "a supported (mul) tape must still compile"
1677        );
1678    }
1679}
pounce_cli/nl_hessian_program.rs

pounce_cli/
nl_hessian_program.rs
