Preserve the copy rule boundary in the trtrs! pullback

[sunxd3]

Closes #1246.

Forward-over-reverse through a triangular solve failed because pre-forward optimisation inlined the copy(B) snapshot in the trtrs! reverse rule through unsafe_copyto! to raw memmove. At that point the element type and copy count needed for differentiation have been erased, so Mooncake intentionally has no rule for the foreign call.

Mark this snapshot call Base.@noinline so the typed copy primitive remains visible to forward mode, and add a focused Hessian regression. This remains necessary with #1215: its width-generic Lifted/NDualArray copy rule handles every tangent lane once the call boundary survives, while the earlier inlining phase and unsupported raw memmove boundary are unchanged.

Validated with the rules/high_order_derivative_patches and rules/lapack test groups, plus the focused reproducer on Julia 1.10–1.12.

Prepared with Codex.

CI Summary — GitHub Actions

Documentation Preview

Mooncake.jl documentation for PR #1247 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR1247/

Performance

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 160.0 ns │     1.69 │        1.81 │    0.75 │        3.69 │   7.14 │
│                  _sum_1000 │ 962.0 ns │     7.01 │        1.04 │  3140.0 │        44.9 │   1.07 │
│               sum_sin_1000 │  6.96 μs │     3.35 │        1.35 │    1.53 │        11.4 │   1.83 │
│              _sum_sin_1000 │  6.13 μs │     3.42 │        1.85 │   237.0 │        12.7 │   2.13 │
│                   kron_sum │ 198.0 μs │     13.6 │        3.24 │    9.12 │       381.0 │   20.1 │
│              kron_view_sum │ 287.0 μs │     11.2 │        5.01 │    13.7 │       332.0 │    8.8 │
│      naive_map_sin_cos_exp │  2.42 μs │      2.9 │        1.52 │ missing │        7.31 │   2.06 │
│            map_sin_cos_exp │  2.24 μs │     3.59 │        1.66 │    1.51 │        6.75 │   2.75 │
│      broadcast_sin_cos_exp │  2.28 μs │     3.22 │        1.58 │    4.35 │        1.48 │   2.14 │
│                 simple_mlp │ 369.0 μs │     4.69 │        2.53 │     2.5 │        8.55 │   2.72 │
│                     gp_lml │ 380.0 μs │     5.72 │        1.76 │    3.01 │     missing │    3.0 │
│ turing_broadcast_benchmark │  1.67 ms │      6.9 │        3.51 │ missing │        37.4 │   2.38 │
│         large_single_block │ 410.0 ns │     5.57 │         2.0 │  4850.0 │        34.4 │    2.1 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[yebai]

Two minor comments, good catches!

[sunxd3]

thanks!