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};
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                    let v_b = scratch[vb as usize];
932                    scratch[dst as usize] = (scratch[dot_a as usize] * v_b
933                        - scratch[va as usize] * scratch[dot_b as usize])
934                        / (v_b * v_b);
935                }
936                HOp::DotSqrt { dst, dot_a, vd } => {
937                    let svd = scratch[vd as usize];
938                    scratch[dst as usize] = if svd > 0.0 {
939                        scratch[dot_a as usize] * 0.5 / svd
940                    } else {
941                        0.0
942                    };
943                }
944                HOp::DotExp { dst, dot_a, vd } => {
945                    scratch[dst as usize] = scratch[dot_a as usize] * scratch[vd as usize];
946                }
947                HOp::DotLog { dst, dot_a, va } => {
948                    scratch[dst as usize] = scratch[dot_a as usize] / scratch[va as usize];
949                }
950                HOp::DotLog10 { dst, dot_a, va } => {
951                    scratch[dst as usize] =
952                        scratch[dot_a as usize] / (scratch[va as usize] * std::f64::consts::LN_10);
953                }
954                HOp::DotSin { dst, dot_a, va } => {
955                    scratch[dst as usize] = scratch[dot_a as usize] * scratch[va as usize].cos();
956                }
957                HOp::DotCos { dst, dot_a, va } => {
958                    scratch[dst as usize] = -scratch[dot_a as usize] * scratch[va as usize].sin();
959                }
960                HOp::DotNeg { dst, dot_a } => {
961                    scratch[dst as usize] = -scratch[dot_a as usize];
962                }
963                HOp::DotAbs { dst, dot_a, va } => {
964                    scratch[dst as usize] = if scratch[va as usize] >= 0.0 {
965                        scratch[dot_a as usize]
966                    } else {
967                        -scratch[dot_a as usize]
968                    };
969                }
970                HOp::DotPow {
971                    dst,
972                    va,
973                    vb,
974                    vd,
975                    dot_a,
976                    dot_b,
977                } => {
978                    let u = scratch[va as usize];
979                    let r = scratch[vb as usize];
980                    let du = scratch[dot_a as usize];
981                    let dr = scratch[dot_b as usize];
982                    let mut result = 0.0;
983                    if r != 0.0 && u != 0.0 {
984                        result += r * u.powf(r - 1.0) * du;
985                    }
986                    if u > 0.0 {
987                        result += scratch[vd as usize] * u.ln() * dr;
988                    }
989                    scratch[dst as usize] = result;
990                }
991
992                HOp::RevAdd {
993                    adj_a,
994                    adj_b,
995                    adj_dot_a,
996                    adj_dot_b,
997                    w,
998                    wd,
999                } => {
1000                    let w_v = scratch[w as usize];
1001                    let wd_v = scratch[wd as usize];
1002                    scratch[adj_a as usize] += w_v;
1003                    scratch[adj_b as usize] += w_v;
1004                    scratch[adj_dot_a as usize] += wd_v;
1005                    scratch[adj_dot_b as usize] += wd_v;
1006                }
1007                HOp::RevSub {
1008                    adj_a,
1009                    adj_b,
1010                    adj_dot_a,
1011                    adj_dot_b,
1012                    w,
1013                    wd,
1014                } => {
1015                    let w_v = scratch[w as usize];
1016                    let wd_v = scratch[wd as usize];
1017                    scratch[adj_a as usize] += w_v;
1018                    scratch[adj_b as usize] -= w_v;
1019                    scratch[adj_dot_a as usize] += wd_v;
1020                    scratch[adj_dot_b as usize] -= wd_v;
1021                }
1022                HOp::RevMul {
1023                    adj_a,
1024                    adj_b,
1025                    adj_dot_a,
1026                    adj_dot_b,
1027                    w,
1028                    wd,
1029                    va,
1030                    vb,
1031                    dot_a,
1032                    dot_b,
1033                } => {
1034                    let w_v = scratch[w as usize];
1035                    let wd_v = scratch[wd as usize];
1036                    let va_v = scratch[va as usize];
1037                    let vb_v = scratch[vb as usize];
1038                    let da_v = scratch[dot_a as usize];
1039                    let db_v = scratch[dot_b as usize];
1040                    scratch[adj_a as usize] += w_v * vb_v;
1041                    scratch[adj_b as usize] += w_v * va_v;
1042                    scratch[adj_dot_a as usize] += wd_v * vb_v + w_v * db_v;
1043                    scratch[adj_dot_b as usize] += wd_v * va_v + w_v * da_v;
1044                }
1045                HOp::RevDiv {
1046                    adj_a,
1047                    adj_b,
1048                    adj_dot_a,
1049                    adj_dot_b,
1050                    w,
1051                    wd,
1052                    va,
1053                    vb,
1054                    dot_a,
1055                    dot_b,
1056                } => {
1057                    let w_v = scratch[w as usize];
1058                    let wd_v = scratch[wd as usize];
1059                    let va_v = scratch[va as usize];
1060                    let vb_v = scratch[vb as usize];
1061                    let vb2 = vb_v * vb_v;
1062                    let vb3 = vb2 * vb_v;
1063                    let da_v = scratch[dot_a as usize];
1064                    let db_v = scratch[dot_b as usize];
1065                    scratch[adj_a as usize] += w_v / vb_v;
1066                    scratch[adj_dot_a as usize] += wd_v / vb_v + w_v * (-db_v / vb2);
1067                    scratch[adj_b as usize] += w_v * (-va_v / vb2);
1068                    scratch[adj_dot_b as usize] +=
1069                        wd_v * (-va_v / vb2) + w_v * (-da_v / vb2 + 2.0 * va_v * db_v / vb3);
1070                }
1071                HOp::RevPow {
1072                    adj_a,
1073                    adj_b,
1074                    adj_dot_a,
1075                    adj_dot_b,
1076                    w,
1077                    wd,
1078                    va,
1079                    vb,
1080                    vd,
1081                    dot_a,
1082                    dot_b,
1083                } => {
1084                    let w_v = scratch[w as usize];
1085                    let wd_v = scratch[wd as usize];
1086                    let u = scratch[va as usize];
1087                    let r = scratch[vb as usize];
1088                    let du = scratch[dot_a as usize];
1089                    let dr = scratch[dot_b as usize];
1090                    if r != 0.0 {
1091                        if u != 0.0 {
1092                            let p_a = r * u.powf(r - 1.0);
1093                            scratch[adj_a as usize] += w_v * p_a;
1094                            let mut dp_a = dr * u.powf(r - 1.0);
1095                            if u > 0.0 {
1096                                dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1097                            } else {
1098                                dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1099                            }
1100                            scratch[adj_dot_a as usize] += wd_v * p_a + w_v * dp_a;
1101                        } else if r >= 2.0 {
1102                            let p_a = 0.0;
1103                            scratch[adj_a as usize] += w_v * p_a;
1104                            let dp_a = if r == 2.0 {
1105                                2.0 * du
1106                            } else {
1107                                r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1108                            };
1109                            scratch[adj_dot_a as usize] += wd_v * p_a + w_v * dp_a;
1110                        }
1111                    }
1112                    if u > 0.0 {
1113                        let ln_u = u.ln();
1114                        let p_b = scratch[vd as usize] * ln_u;
1115                        scratch[adj_b as usize] += w_v * p_b;
1116                        let dur = scratch[vd as usize] * (r * du / u + dr * ln_u);
1117                        let dp_b = dur * ln_u + scratch[vd as usize] * du / u;
1118                        scratch[adj_dot_b as usize] += wd_v * p_b + w_v * dp_b;
1119                    }
1120                }
1121                HOp::RevNeg {
1122                    adj_a,
1123                    adj_dot_a,
1124                    w,
1125                    wd,
1126                } => {
1127                    scratch[adj_a as usize] -= scratch[w as usize];
1128                    scratch[adj_dot_a as usize] -= scratch[wd as usize];
1129                }
1130                HOp::RevAbs {
1131                    adj_a,
1132                    adj_dot_a,
1133                    w,
1134                    wd,
1135                    va,
1136                } => {
1137                    let s = if scratch[va as usize] >= 0.0 {
1138                        1.0
1139                    } else {
1140                        -1.0
1141                    };
1142                    scratch[adj_a as usize] += scratch[w as usize] * s;
1143                    scratch[adj_dot_a as usize] += scratch[wd as usize] * s;
1144                }
1145                HOp::RevSqrt {
1146                    adj_a,
1147                    adj_dot_a,
1148                    w,
1149                    wd,
1150                    va: _,
1151                    vd,
1152                    dot_a,
1153                } => {
1154                    let sv = scratch[vd as usize];
1155                    if sv > 0.0 {
1156                        let fp = 0.5 / sv;
1157                        let fpp = -0.25 / (sv * sv * sv);
1158                        let w_v = scratch[w as usize];
1159                        let wd_v = scratch[wd as usize];
1160                        scratch[adj_a as usize] += w_v * fp;
1161                        scratch[adj_dot_a as usize] +=
1162                            wd_v * fp + w_v * fpp * scratch[dot_a as usize];
1163                    }
1164                }
1165                HOp::RevExp {
1166                    adj_a,
1167                    adj_dot_a,
1168                    w,
1169                    wd,
1170                    vd,
1171                    dot_a,
1172                } => {
1173                    let ev = scratch[vd as usize];
1174                    let w_v = scratch[w as usize];
1175                    let wd_v = scratch[wd as usize];
1176                    scratch[adj_a as usize] += w_v * ev;
1177                    scratch[adj_dot_a as usize] += wd_v * ev + w_v * ev * scratch[dot_a as usize];
1178                }
1179                HOp::RevLog {
1180                    adj_a,
1181                    adj_dot_a,
1182                    w,
1183                    wd,
1184                    va,
1185                    dot_a,
1186                } => {
1187                    let u = scratch[va as usize];
1188                    let w_v = scratch[w as usize];
1189                    let wd_v = scratch[wd as usize];
1190                    scratch[adj_a as usize] += w_v / u;
1191                    scratch[adj_dot_a as usize] +=
1192                        wd_v / u + w_v * (-1.0 / (u * u)) * scratch[dot_a as usize];
1193                }
1194                HOp::RevLog10 {
1195                    adj_a,
1196                    adj_dot_a,
1197                    w,
1198                    wd,
1199                    va,
1200                    dot_a,
1201                } => {
1202                    let u = scratch[va as usize];
1203                    let c = std::f64::consts::LN_10;
1204                    let w_v = scratch[w as usize];
1205                    let wd_v = scratch[wd as usize];
1206                    scratch[adj_a as usize] += w_v / (u * c);
1207                    scratch[adj_dot_a as usize] +=
1208                        wd_v / (u * c) + w_v * (-1.0 / (u * u * c)) * scratch[dot_a as usize];
1209                }
1210                HOp::RevSin {
1211                    adj_a,
1212                    adj_dot_a,
1213                    w,
1214                    wd,
1215                    va,
1216                    dot_a,
1217                } => {
1218                    let u = scratch[va as usize];
1219                    let cu = u.cos();
1220                    let w_v = scratch[w as usize];
1221                    let wd_v = scratch[wd as usize];
1222                    scratch[adj_a as usize] += w_v * cu;
1223                    scratch[adj_dot_a as usize] +=
1224                        wd_v * cu + w_v * (-u.sin()) * scratch[dot_a as usize];
1225                }
1226                HOp::RevCos {
1227                    adj_a,
1228                    adj_dot_a,
1229                    w,
1230                    wd,
1231                    va,
1232                    dot_a,
1233                } => {
1234                    let u = scratch[va as usize];
1235                    let su = u.sin();
1236                    let w_v = scratch[w as usize];
1237                    let wd_v = scratch[wd as usize];
1238                    scratch[adj_a as usize] -= w_v * su;
1239                    scratch[adj_dot_a as usize] +=
1240                        wd_v * (-su) + w_v * (-u.cos()) * scratch[dot_a as usize];
1241                }
1242
1243                HOp::HessEmit {
1244                    hess_ptr,
1245                    adj_dot_slot,
1246                } => {
1247                    values[hess_ptr as usize] += weight * scratch[adj_dot_slot as usize];
1248                }
1249            }
1250        }
1251    }
1252}
1253
1254/// Whether [`HessianProgram::compile`] can lower a single opcode. The
1255/// program path covers smooth arithmetic plus `sin`/`cos`; every other
1256/// opcode — AMPL external `Funcall`, the remaining transcendentals
1257/// (`tan`/`atan`/`acos`/the hyperbolics/`asin`/`erf`…) and `atan2`, and the
1258/// `min`/`max`/conditional/logical family — is unsupported, so `compile`
1259/// returns `None` and the caller falls back to the `Tape`
1260/// (`build_with_externals`) interpreter path rather than panicking on user
1261/// input (code review L28). This is the single source of truth for the
1262/// supported set: the per-sweep match arms and the dependence/reachability
1263/// analyses are only ever reached with opcodes this predicate accepts, so
1264/// their unsupported branches are `unreachable!`.
1265fn program_supports_op(op: &TapeOp) -> bool {
1266    matches!(
1267        op,
1268        TapeOp::Const(_)
1269            | TapeOp::Var(_)
1270            | TapeOp::Add(_, _)
1271            | TapeOp::Sub(_, _)
1272            | TapeOp::Mul(_, _)
1273            | TapeOp::Div(_, _)
1274            | TapeOp::Pow(_, _)
1275            | TapeOp::Neg(_)
1276            | TapeOp::Abs(_)
1277            | TapeOp::Sqrt(_)
1278            | TapeOp::Exp(_)
1279            | TapeOp::Log(_)
1280            | TapeOp::Log10(_)
1281            | TapeOp::Sin(_)
1282            | TapeOp::Cos(_)
1283    )
1284}
1285
1286/// `out[i]` = does tape slot `i` contribute (transitively) to the
1287/// output slot `n-1`. Used to skip emitting reverse-pass ops for
1288/// dead slots.
1289fn reachable_to_output(tape: &Tape) -> Vec<bool> {
1290    let n = tape.ops.len();
1291    let mut r = vec![false; n];
1292    if n == 0 {
1293        return r;
1294    }
1295    r[n - 1] = true;
1296    for i in (0..n).rev() {
1297        if !r[i] {
1298            continue;
1299        }
1300        match tape.ops[i] {
1301            TapeOp::Const(_) | TapeOp::Var(_) => {}
1302            TapeOp::Add(a, b)
1303            | TapeOp::Sub(a, b)
1304            | TapeOp::Mul(a, b)
1305            | TapeOp::Div(a, b)
1306            | TapeOp::Pow(a, b)
1307            | TapeOp::Atan2(a, b)
1308            | TapeOp::CEntropy(a, b) => {
1309                r[a] = true;
1310                r[b] = true;
1311            }
1312            TapeOp::Neg(a)
1313            | TapeOp::Abs(a)
1314            | TapeOp::Sqrt(a)
1315            | TapeOp::Exp(a)
1316            | TapeOp::Log(a)
1317            | TapeOp::Log10(a)
1318            | TapeOp::Sin(a)
1319            | TapeOp::Cos(a)
1320            | TapeOp::Tan(a)
1321            | TapeOp::Atan(a)
1322            | TapeOp::Acos(a)
1323            | TapeOp::Sinh(a)
1324            | TapeOp::Cosh(a)
1325            | TapeOp::Tanh(a)
1326            | TapeOp::Asin(a)
1327            | TapeOp::Acosh(a)
1328            | TapeOp::Asinh(a)
1329            | TapeOp::Erf(a)
1330            | TapeOp::XLogX(a)
1331            | TapeOp::Atanh(a) => {
1332                r[a] = true;
1333            }
1334            TapeOp::Funcall(_) => unreachable!(
1335                "HessianProgram path does not support AMPL external functions; \
1336                 use the Tape (build_with_externals) path instead."
1337            ),
1338            TapeOp::Cmp(_, _, _)
1339            | TapeOp::And(_, _)
1340            | TapeOp::Or(_, _)
1341            | TapeOp::Not(_)
1342            | TapeOp::Select(_, _, _)
1343            | TapeOp::Min(_, _)
1344            | TapeOp::Max(_, _) => unreachable!(
1345                "HessianProgram path does not support conditional / logical / min-max \
1346                 opcodes; use the Tape (build_with_externals) path instead."
1347            ),
1348        }
1349    }
1350    r
1351}
1352
1353/// `out[i]` = does tape slot `i` transitively read from `Var(j)`.
1354/// Used to prune forward-tangent ops (slots with `out[i] = false`
1355/// have `dot[i] = 0` and the rest of the per-`j` pass can skip
1356/// them).
1357fn depends_on_var(tape: &Tape, j: usize) -> Vec<bool> {
1358    let n = tape.ops.len();
1359    let mut d = vec![false; n];
1360    for (i, op) in tape.ops.iter().enumerate() {
1361        d[i] = match *op {
1362            TapeOp::Const(_) => false,
1363            TapeOp::Var(k) => k == j,
1364            TapeOp::Add(a, b)
1365            | TapeOp::Sub(a, b)
1366            | TapeOp::Mul(a, b)
1367            | TapeOp::Div(a, b)
1368            | TapeOp::Pow(a, b)
1369            | TapeOp::CEntropy(a, b)
1370            | TapeOp::Atan2(a, b) => d[a] || d[b],
1371            TapeOp::Neg(a)
1372            | TapeOp::Abs(a)
1373            | TapeOp::Sqrt(a)
1374            | TapeOp::Exp(a)
1375            | TapeOp::Log(a)
1376            | TapeOp::Log10(a)
1377            | TapeOp::Sin(a)
1378            | TapeOp::Cos(a)
1379            | TapeOp::Tan(a)
1380            | TapeOp::Atan(a)
1381            | TapeOp::Acos(a)
1382            | TapeOp::Sinh(a)
1383            | TapeOp::Cosh(a)
1384            | TapeOp::Tanh(a)
1385            | TapeOp::Asin(a)
1386            | TapeOp::Acosh(a)
1387            | TapeOp::Asinh(a)
1388            | TapeOp::XLogX(a)
1389            | TapeOp::Erf(a)
1390            | TapeOp::Atanh(a) => d[a],
1391            TapeOp::Funcall(_) => unreachable!(
1392                "HessianProgram path does not support AMPL external functions; \
1393                 use the Tape (build_with_externals) path instead."
1394            ),
1395            TapeOp::Cmp(_, _, _)
1396            | TapeOp::And(_, _)
1397            | TapeOp::Or(_, _)
1398            | TapeOp::Not(_)
1399            | TapeOp::Select(_, _, _)
1400            | TapeOp::Min(_, _)
1401            | TapeOp::Max(_, _) => unreachable!(
1402                "HessianProgram path does not support conditional / logical / min-max \
1403                 opcodes; use the Tape (build_with_externals) path instead."
1404            ),
1405        };
1406    }
1407    d
1408}
1409
1410#[cfg(test)]
1411mod tests {
1412    use super::*;
1413    use crate::nl_reader::{BinOp, Expr, UnaryOp};
1414    use std::collections::BTreeSet;
1415    use std::sync::Arc;
1416
1417    fn cnst(c: f64) -> Expr {
1418        Expr::Const(c)
1419    }
1420    fn var(i: usize) -> Expr {
1421        Expr::Var(i)
1422    }
1423    fn add(a: Expr, b: Expr) -> Expr {
1424        Expr::Binary(BinOp::Add, Box::new(a), Box::new(b))
1425    }
1426    fn mul(a: Expr, b: Expr) -> Expr {
1427        Expr::Binary(BinOp::Mul, Box::new(a), Box::new(b))
1428    }
1429    fn pow(a: Expr, b: Expr) -> Expr {
1430        Expr::Binary(BinOp::Pow, Box::new(a), Box::new(b))
1431    }
1432    fn div(a: Expr, b: Expr) -> Expr {
1433        Expr::Binary(BinOp::Div, Box::new(a), Box::new(b))
1434    }
1435    fn sub(a: Expr, b: Expr) -> Expr {
1436        Expr::Binary(BinOp::Sub, Box::new(a), Box::new(b))
1437    }
1438    fn unary(op: UnaryOp, a: Expr) -> Expr {
1439        Expr::Unary(op, Box::new(a))
1440    }
1441
1442    /// Build the same shared (row, col) -> pos map both AD paths
1443    /// scatter into. Lower-triangle pairs, in tape.variables() order.
1444    fn build_hess_map(tape: &Tape) -> (HashMap<(usize, usize), usize>, Vec<(usize, usize)>) {
1445        let vars = tape.variables();
1446        let mut pairs: Vec<(usize, usize)> = Vec::new();
1447        let mut map: HashMap<(usize, usize), usize> = HashMap::new();
1448        for (ai, &vi) in vars.iter().enumerate() {
1449            for &vj in &vars[..=ai] {
1450                let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
1451                map.entry((r, c)).or_insert_with(|| {
1452                    let p = pairs.len();
1453                    pairs.push((r, c));
1454                    p
1455                });
1456            }
1457        }
1458        (map, pairs)
1459    }
1460
1461    /// Run both implementations against the same input and assert
1462    /// values match to a tight ULP-aligned tolerance.
1463    fn assert_program_matches_tape(tape: &Tape, x: &[f64], weight: f64) {
1464        let (hess_map, pairs) = build_hess_map(tape);
1465        let nnz = pairs.len();
1466
1467        let mut tape_vals = vec![0.0; nnz];
1468        tape.hessian_accumulate(x, weight, &hess_map, &mut tape_vals);
1469
1470        let program =
1471            HessianProgram::compile(tape, &hess_map).expect("tape uses only supported opcodes");
1472        let mut scratch = vec![0.0; program.n_slots()];
1473        let mut prog_vals = vec![0.0; nnz];
1474        program.execute(x, weight, &mut scratch, &mut prog_vals);
1475
1476        for (k, &(r, c)) in pairs.iter().enumerate() {
1477            let tol = tape_vals[k].abs().max(1.0) * 1e-12;
1478            assert!(
1479                (tape_vals[k] - prog_vals[k]).abs() < tol,
1480                "H[{},{}]: tape={:.6e} prog={:.6e}",
1481                r,
1482                c,
1483                tape_vals[k],
1484                prog_vals[k]
1485            );
1486        }
1487    }
1488
1489    #[test]
1490    fn matches_quadratic() {
1491        let e = add(
1492            add(
1493                mul(cnst(3.0), pow(var(0), cnst(2.0))),
1494                mul(cnst(2.0), mul(var(0), var(1))),
1495            ),
1496            pow(var(1), cnst(2.0)),
1497        );
1498        let tape = Tape::build(&e);
1499        assert_program_matches_tape(&tape, &[2.0, 3.0], 1.0);
1500        assert_program_matches_tape(&tape, &[-1.5, 0.7], 2.5);
1501    }
1502
1503    #[test]
1504    fn matches_transcendental() {
1505        let e = Expr::Sum(vec![
1506            unary(UnaryOp::Exp, var(0)),
1507            unary(UnaryOp::Sin, var(1)),
1508            unary(UnaryOp::Log, var(0)),
1509            unary(UnaryOp::Sqrt, var(1)),
1510            mul(var(0), var(1)),
1511            unary(UnaryOp::Cos, add(var(0), var(1))),
1512        ]);
1513        let tape = Tape::build(&e);
1514        assert_program_matches_tape(&tape, &[1.5, 2.0], 1.0);
1515        assert_program_matches_tape(&tape, &[0.3, 4.1], -0.4);
1516    }
1517
1518    #[test]
1519    fn matches_division() {
1520        let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
1521        let tape = Tape::build(&e);
1522        assert_program_matches_tape(&tape, &[0.5, 1.2], 1.0);
1523    }
1524
1525    #[test]
1526    fn matches_through_cse() {
1527        let body = Arc::new(add(var(0), var(1)));
1528        let e = add(
1529            pow(Expr::Cse(body.clone()), cnst(2.0)),
1530            Expr::Cse(body.clone()),
1531        );
1532        let tape = Tape::build(&e);
1533        assert_program_matches_tape(&tape, &[1.0, 2.0], 1.0);
1534        assert_program_matches_tape(&tape, &[-0.5, 3.3], 0.7);
1535    }
1536
1537    #[test]
1538    fn matches_pow_chain() {
1539        // After Tier 1 this lowers to a Mul chain; verify both
1540        // paths agree on the lowered form too.
1541        let e = add(pow(var(0), cnst(3.0)), pow(var(1), cnst(-2.0)));
1542        let tape = Tape::build(&e);
1543        assert_program_matches_tape(&tape, &[1.7, 0.8], 1.0);
1544    }
1545
1546    #[test]
1547    fn matches_residual_pow_with_var_exponent() {
1548        // Pow where the exponent is variable (not constant), so
1549        // it survives Tier 1 and exercises the RevPow / DotPow
1550        // compound branches.
1551        let e = pow(var(0), var(1));
1552        let tape = Tape::build(&e);
1553        assert_program_matches_tape(&tape, &[2.5, 1.4], 1.0);
1554        assert_program_matches_tape(&tape, &[0.6, 2.1], -1.0);
1555    }
1556
1557    #[test]
1558    fn matches_sub_neg_abs() {
1559        let e = sub(
1560            unary(UnaryOp::Neg, var(0)),
1561            unary(UnaryOp::Abs, sub(var(1), var(0))),
1562        );
1563        let tape = Tape::build(&e);
1564        assert_program_matches_tape(&tape, &[1.0, -2.0], 1.0);
1565        assert_program_matches_tape(&tape, &[-3.5, 4.0], 0.9);
1566    }
1567
1568    #[test]
1569    fn slots_layout_matches_design() {
1570        let e = mul(var(0), var(1));
1571        let tape = Tape::build(&e);
1572        let (hess_map, _) = build_hess_map(&tape);
1573        let prog = HessianProgram::compile(&tape, &hess_map).expect("mul tape is supported");
1574        assert_eq!(prog.n_slots(), 4 * tape.ops.len());
1575    }
1576
1577    /// Sanity: the pruning analyses are consistent with the slot
1578    /// structure exposed via `hessian_sparsity()`.
1579    #[test]
1580    fn dependence_matches_hessian_sparsity_for_simple_case() {
1581        let e = add(unary(UnaryOp::Sin, var(0)), mul(var(1), var(2)));
1582        let tape = Tape::build(&e);
1583        let s: BTreeSet<(usize, usize)> = tape.hessian_sparsity();
1584        // (0,0) from sin, (2,1) from x1*x2, (1,1)/(2,2) NOT there
1585        // because Mul(x1, x2) emits cross only.
1586        assert!(s.contains(&(0, 0)));
1587        assert!(s.contains(&(2, 1)));
1588        assert_program_matches_tape(&tape, &[0.7, 1.1, 2.2], 1.0);
1589    }
1590
1591    /// Code review L28: a tape using an opcode the program path cannot
1592    /// lower (here `tan`) must make `compile` return `None` — a graceful
1593    /// fall-back-to-the-`Tape`-path signal — rather than panic. Previously
1594    /// the per-sweep match arms `panic!`'d on such ops, which would crash
1595    /// the process on arbitrary user `.nl` input if this path were ever
1596    /// wired into dispatch.
1597    #[test]
1598    fn unsupported_opcode_returns_none_instead_of_panicking() {
1599        // `tan(x0)` lowers to `TapeOp::Tan`, which the HessianProgram
1600        // compiler does not support.
1601        let e = unary(UnaryOp::Tan, var(0));
1602        let tape = Tape::build(&e);
1603        let (hess_map, _) = build_hess_map(&tape);
1604        assert!(
1605            HessianProgram::compile(&tape, &hess_map).is_none(),
1606            "tan() tape must fall back (None), not compile"
1607        );
1608
1609        // A fully-supported tape still compiles to `Some`, so the guard
1610        // rejects only genuinely-unsupported ops.
1611        let ok = mul(var(0), var(1));
1612        let ok_tape = Tape::build(&ok);
1613        let (ok_map, _) = build_hess_map(&ok_tape);
1614        assert!(
1615            HessianProgram::compile(&ok_tape, &ok_map).is_some(),
1616            "a supported (mul) tape must still compile"
1617        );
1618    }
1619}
pounce_cli/nl_hessian_program.rs

pounce_cli/
nl_hessian_program.rs
