Skip to content

perf(BMT): recover the primal tendency from the Rosenbrock Jacobian pass#742

Draft
haakon-e wants to merge 1 commit into
he/rosenbrock-modesfrom
he/diffresults-jacobian
Draft

perf(BMT): recover the primal tendency from the Rosenbrock Jacobian pass#742
haakon-e wants to merge 1 commit into
he/rosenbrock-modesfrom
he/diffresults-jacobian

Conversation

@haakon-e

Copy link
Copy Markdown
Member

Summary

In the Rosenbrock micro substep the primal tendency f = g(x) and the Jacobian ForwardDiff.jacobian(g, x) were computed separately, but ForwardDiff's Jacobian pass already evaluates f as the dual .value and discards it. Recover it with DiffResults/jacobian!, eliminating one full (quadrature-dominated) tendency evaluation per substep.

What changed (src/BMT_rosenbrock.jl, Project.toml, BulkMicrophysicsTendencies.jl)

  • Both the 2M+P3 and the 1M ExactJacobian substep paths now go through a shared _tendency_and_jacobian(jacobian, g, x) method: r = ForwardDiff.jacobian!(DiffResults.JacobianResult(x), g, x), then f = typeof(x)(DiffResults.value(r)), J = DiffResults.jacobian(r). The method reconstructs the same FieldVector species type as x (MicroState1M or MicroState2MP3), removing the duplicated DiffResults code that previously lived inline in the 2M+P3 loop.
  • The analytic Donor / CoupledDonor 1M Jacobians have no f by-product, so their _tendency_and_jacobian methods keep the separate f = g(x). The non-finite euler fallback also keeps f = g(x).
  • Add DiffResults to [deps] + [compat]; import DiffResults in BulkMicrophysicsTendencies.jl.

Mechanism / why

ForwardDiff.jacobian(g, x::StaticArray) runs one dual evaluation and extracts the Jacobian, throwing away the dual .value (= f). jacobian! into an immutable JacobianResult does the identical single pass and additionally records the value, so f is recovered for free. JacobianResult is immutable for a StaticArray, so the jacobian! return value must be captured (it is not mutated in place). Two opaque call sites otherwise prevent common-subexpression elimination, so the saving is real.

Validation

  • f and J bit-identical to the separate computation (verified === on the reconstructed MicroState1M and MicroState2MP3 and on J); no extra allocations (@allocated == 0 for both helper paths and the full 2M+P3 bulk_microphysics_tendencies); type-stable (@inferred).
  • ~27% faster per substep (order-independent since J ≈ 2.68× f).
  • rosenbrock_1m / mode / framework / verbose tests pass (F32 + F64), including the allocation and "vs fine explicit reference" checks.

Caveats

  • ExactJacobian paths only (by design). DiffResults becomes a direct dep (it was transitive via ForwardDiff).

@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 09ba907 to e4595a3 Compare June 23, 2026 00:32
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 632b432 to f712481 Compare June 23, 2026 00:32
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from e4595a3 to 2068d98 Compare June 23, 2026 21:32
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from f712481 to 7dfaa2e Compare June 23, 2026 21:32
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 2068d98 to 49695f3 Compare June 23, 2026 22:26
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 7dfaa2e to 183d055 Compare June 23, 2026 22:26
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 49695f3 to 2aabe66 Compare June 24, 2026 21:56
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 183d055 to fb540f8 Compare June 24, 2026 21:56
@codecov

codecov Bot commented Jun 24, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 88.23529% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.99%. Comparing base (10acc85) to head (4a53652).

Additional details and impacted files
@@                   Coverage Diff                   @@
##           he/rosenbrock-modes     #742      +/-   ##
=======================================================
- Coverage                93.11%   92.99%   -0.13%     
=======================================================
  Files                       57       57              
  Lines                     3298     3310      +12     
=======================================================
+ Hits                      3071     3078       +7     
- Misses                     227      232       +5     
Components Coverage Δ
src 93.68% <88.23%> (-0.14%) ⬇️
ext 69.47% <ø> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 2aabe66 to 90f0aea Compare June 25, 2026 18:01
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from fb540f8 to 9c27634 Compare June 25, 2026 18:01
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 90f0aea to 7e848a2 Compare June 26, 2026 00:38
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 9c27634 to a35a6f7 Compare June 26, 2026 00:38
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 7e848a2 to dfa0e14 Compare June 26, 2026 01:22
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from a35a6f7 to 45815ea Compare June 26, 2026 01:22
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from dfa0e14 to 0604b2c Compare June 26, 2026 01:31
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 45815ea to c443a95 Compare June 26, 2026 01:31
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 0604b2c to 544ff37 Compare July 1, 2026 05:35
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from c443a95 to e4798a9 Compare July 1, 2026 05:35
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 544ff37 to e92fd4a Compare July 1, 2026 05:46
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from e4798a9 to 1882907 Compare July 1, 2026 05:46
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from e92fd4a to 4da1efb Compare July 1, 2026 06:24
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 1882907 to 5788941 Compare July 1, 2026 06:25
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 4da1efb to 3d066ff Compare July 1, 2026 15:01
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 5788941 to c592fab Compare July 1, 2026 15:01
@haakon-e
haakon-e force-pushed the he/ice-selfcol-triangle branch from 3d066ff to e573c80 Compare July 1, 2026 21:17
@haakon-e
haakon-e changed the base branch from he/ice-selfcol-triangle to he/rosenbrock-modes July 1, 2026 21:29
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from c592fab to b204ce7 Compare July 1, 2026 21:29
@haakon-e
haakon-e force-pushed the he/rosenbrock-modes branch from 090fd49 to 043da95 Compare July 8, 2026 04:34
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from b204ce7 to 00b8d5d Compare July 8, 2026 04:34
@haakon-e
haakon-e force-pushed the he/rosenbrock-modes branch from 043da95 to c5c4783 Compare July 10, 2026 02:14
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 00b8d5d to 174b6c1 Compare July 10, 2026 02:14
@haakon-e
haakon-e force-pushed the he/rosenbrock-modes branch from c5c4783 to 1c96cf0 Compare July 10, 2026 23:15
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 174b6c1 to e362cbc Compare July 10, 2026 23:15
@haakon-e
haakon-e force-pushed the he/rosenbrock-modes branch from 1c96cf0 to f821b88 Compare July 11, 2026 00:58
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from e362cbc to 9f61db9 Compare July 11, 2026 00:58
@haakon-e
haakon-e force-pushed the he/rosenbrock-modes branch from f821b88 to 3805142 Compare July 11, 2026 02:14
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from 9f61db9 to d51f9b9 Compare July 11, 2026 02:14
@haakon-e
haakon-e force-pushed the he/rosenbrock-modes branch from 3805142 to ed1ddb6 Compare July 11, 2026 02:31
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from d51f9b9 to fb14b5f Compare July 11, 2026 02:31
ForwardDiff's Jacobian pass already evaluates the primal tendency f as the dual
.value and discards it, so a separate f = g(x) repeats a full (quadrature-
dominated) tendency evaluation. Recover f from the same pass: seed all N
partials of the state in a single call to g with static Duals, and read the
values and partials directly into the same FieldVector species type as x
(MicroState1M or MicroState2MP3). Applies to the 2M+P3 and the 1M
ExactJacobian paths; the analytic Donor/CoupledDonor 1M Jacobians (no f
by-product) keep their separate evaluation.

The DiffResults dependency is dropped: DiffResults.JacobianResult's mutable
result buffer heap-allocates on GPU, so the seeding is done directly with
StaticArrays instead. Both paths go through the shared
_tendency_and_jacobian(jacobian, g, x) method, removing the duplicated code
in the 2M+P3 substep loop.

Measured on MicroState2MP3 (N=8, quadrature-dominated) and MicroState1M
(N=4): the single full-width call is ~31% faster than the previous two-call
f = g(x); J = FD.jacobian(g, x) for 2M+P3, and ~39% faster for 1M, with zero
allocations in both cases. f and J are bitwise identical to the two-call
reference across representative states. The tendency and Jacobian are
unchanged, the path stays type-stable, and it composes with an outer AD pass.
@haakon-e
haakon-e force-pushed the he/rosenbrock-modes branch from ed1ddb6 to 10acc85 Compare July 11, 2026 03:28
@haakon-e
haakon-e force-pushed the he/diffresults-jacobian branch from fb14b5f to 4a53652 Compare July 11, 2026 03:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant