perf(BMT): recover the primal tendency from the Rosenbrock Jacobian pass#742
Draft
haakon-e wants to merge 1 commit into
Draft
perf(BMT): recover the primal tendency from the Rosenbrock Jacobian pass#742haakon-e wants to merge 1 commit into
haakon-e wants to merge 1 commit into
Conversation
Member
Author
|
This change is part of the following stack:
Change managed by git-spice. |
This was referenced Jun 23, 2026
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 23, 2026 00:32
09ba907 to
e4595a3
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 23, 2026 00:32
632b432 to
f712481
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 23, 2026 21:32
e4595a3 to
2068d98
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 23, 2026 21:32
f712481 to
7dfaa2e
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 23, 2026 22:26
2068d98 to
49695f3
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 23, 2026 22:26
7dfaa2e to
183d055
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 24, 2026 21:56
49695f3 to
2aabe66
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 24, 2026 21:56
183d055 to
fb540f8
Compare
Codecov Report❌ Patch coverage is 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
🚀 New features to boost your workflow:
|
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 25, 2026 18:01
2aabe66 to
90f0aea
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 25, 2026 18:01
fb540f8 to
9c27634
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 26, 2026 00:38
90f0aea to
7e848a2
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 26, 2026 00:38
9c27634 to
a35a6f7
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 26, 2026 01:22
7e848a2 to
dfa0e14
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 26, 2026 01:22
a35a6f7 to
45815ea
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
June 26, 2026 01:31
dfa0e14 to
0604b2c
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
June 26, 2026 01:31
45815ea to
c443a95
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
July 1, 2026 05:35
0604b2c to
544ff37
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 1, 2026 05:35
c443a95 to
e4798a9
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
July 1, 2026 05:46
544ff37 to
e92fd4a
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 1, 2026 05:46
e4798a9 to
1882907
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
July 1, 2026 06:24
e92fd4a to
4da1efb
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 1, 2026 06:25
1882907 to
5788941
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
July 1, 2026 15:01
4da1efb to
3d066ff
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 1, 2026 15:01
5788941 to
c592fab
Compare
haakon-e
force-pushed
the
he/ice-selfcol-triangle
branch
from
July 1, 2026 21:17
3d066ff to
e573c80
Compare
haakon-e
changed the base branch from
he/ice-selfcol-triangle
to
he/rosenbrock-modes
July 1, 2026 21:29
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 1, 2026 21:29
c592fab to
b204ce7
Compare
haakon-e
force-pushed
the
he/rosenbrock-modes
branch
from
July 8, 2026 04:34
090fd49 to
043da95
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 8, 2026 04:34
b204ce7 to
00b8d5d
Compare
haakon-e
force-pushed
the
he/rosenbrock-modes
branch
from
July 10, 2026 02:14
043da95 to
c5c4783
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 10, 2026 02:14
00b8d5d to
174b6c1
Compare
haakon-e
force-pushed
the
he/rosenbrock-modes
branch
from
July 10, 2026 23:15
c5c4783 to
1c96cf0
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 10, 2026 23:15
174b6c1 to
e362cbc
Compare
haakon-e
force-pushed
the
he/rosenbrock-modes
branch
from
July 11, 2026 00:58
1c96cf0 to
f821b88
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 11, 2026 00:58
e362cbc to
9f61db9
Compare
haakon-e
force-pushed
the
he/rosenbrock-modes
branch
from
July 11, 2026 02:14
f821b88 to
3805142
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 11, 2026 02:14
9f61db9 to
d51f9b9
Compare
haakon-e
force-pushed
the
he/rosenbrock-modes
branch
from
July 11, 2026 02:31
3805142 to
ed1ddb6
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 11, 2026 02:31
d51f9b9 to
fb14b5f
Compare
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
force-pushed
the
he/rosenbrock-modes
branch
from
July 11, 2026 03:28
ed1ddb6 to
10acc85
Compare
haakon-e
force-pushed
the
he/diffresults-jacobian
branch
from
July 11, 2026 03:28
fb14b5f to
4a53652
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
In the Rosenbrock micro substep the primal tendency
f = g(x)and the JacobianForwardDiff.jacobian(g, x)were computed separately, but ForwardDiff's Jacobian pass already evaluatesfas the dual.valueand discards it. Recover it withDiffResults/jacobian!, eliminating one full (quadrature-dominated) tendency evaluation per substep.What changed (src/BMT_rosenbrock.jl, Project.toml, BulkMicrophysicsTendencies.jl)
_tendency_and_jacobian(jacobian, g, x)method:r = ForwardDiff.jacobian!(DiffResults.JacobianResult(x), g, x), thenf = typeof(x)(DiffResults.value(r)),J = DiffResults.jacobian(r). The method reconstructs the sameFieldVectorspecies type asx(MicroState1MorMicroState2MP3), removing the duplicated DiffResults code that previously lived inline in the 2M+P3 loop.fby-product, so their_tendency_and_jacobianmethods keep the separatef = g(x). The non-finite euler fallback also keepsf = g(x).DiffResultsto [deps] + [compat];import DiffResultsin 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 immutableJacobianResultdoes the identical single pass and additionally records the value, sofis recovered for free.JacobianResultis immutable for a StaticArray, so thejacobian!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
fand J bit-identical to the separate computation (verified===on the reconstructedMicroState1MandMicroState2MP3and on J); no extra allocations (@allocated == 0for both helper paths and the full 2M+P3bulk_microphysics_tendencies); type-stable (@inferred).Caveats
DiffResultsbecomes a direct dep (it was transitive via ForwardDiff).