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

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 the inputs
with static `Dual`s, evaluate g once, 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.

Static seeding replaces `DiffResults.JacobianResult`, whose mutable result
buffer heap-allocates on GPU; the `DiffResults` dependency is dropped. Both
paths go through the shared `_tendency_and_jacobian(jacobian, g, x)` method,
removing the duplicated code in the 2M+P3 substep loop.

Reduces the per-substep cost by roughly 27%. The tendency and Jacobian are
unchanged, the path stays type-stable, and it composes with an outer AD pass.

Author: haakon-e

Project.toml

@@ -5,6 +5,7 @@ authors = ["Climate Modeling Alliance"]
 
 [deps]
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -27,6 +28,7 @@ EmulatorModelsExt = ["DataFrames", "MLJ"]
 [compat]
 ClimaParams = "1.0.18"
 DataFrames = "1.6"
+DiffResults = "1"
 DocStringExtensions = "0.8, 0.9"
 FastGaussQuadrature = "1"
 ForwardDiff = "0.10, 1"

src/BMT_rosenbrock.jl

@@ -429,10 +429,10 @@ fields as the `Instantaneous` entry (without the activation diagnostic).
     Tsub = T
     for _ in 1:nsub_eff
         g = Instantaneous2MP3Tendency(mp, tps, ρ, Tsub, q_tot, logλ)
-        f = g(x)
         x_prev = x
         if all(isfinite, x)
-            J = _apply_growth(mode.growth, FD.jacobian(g, x))
+            f, J_raw = _tendency_and_jacobian(mode.jacobian, g, x)
+            J = _apply_growth(mode.growth, J_raw)
             z = _species_mask(mode.jacobian, mode.growth)(x)
             d = if all(isfinite, J)
                 _rosenbrock_update(x, f, J, z, h) - x
@@ -442,6 +442,7 @@ fields as the `Instantaneous` entry (without the activation diagnostic).
             d = _apply_limiter(mode.limiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps)
             x = max.(x .+ d, 0)
         else
+            f = g(x)
             x = _euler_update(x, f, h)
         end
         Δ = x - x_prev
@@ -749,6 +750,36 @@ end
 "Exact ForwardDiff Jacobian provider for [`_rosenbrock_average_1m`](@ref)."
 @inline _ad_jacobian_1m(g, x) = FD.jacobian(g, x)
 
+"""
+    _tendency_and_jacobian(jacobian, g, x)
+
+The raw substep tendency `f = g(x)` and the substep Jacobian (before the growth
+treatment) for a [`Jacobian`](@ref) option, returned as `(f, J)`.
+
+For [`ExactJacobian`](@ref) both the value and the Jacobian are taken from a
+single `ForwardDiff` pass: `x` is seeded with unit partials and `g` is evaluated
+once, giving the primal `f` and the `N×N` Jacobian as static arrays. The
+donor-based matrices ([`DonorJacobian`](@ref), [`CoupledDonorJacobian`](@ref))
+produce no tendency by-product, so `f = g(x)` is evaluated separately.
+"""
+@inline function _tendency_and_jacobian(::ExactJacobian, g, x::SA.FieldVector{N, FT}) where {N, FT}
+    Tag = typeof(FD.Tag(g, FT))
+    dx = SA.SVector(ntuple(
+        i -> FD.Dual{Tag}(x[i], ntuple(j -> ifelse(i == j, one(FT), zero(FT)), Val(N))...),
+        Val(N),
+    ))
+    y = g(dx)
+    f = typeof(x)(ntuple(i -> FD.value(y[i]), Val(N))...)
+    J = SA.SMatrix{N, N, FT}(ntuple(
+        k -> FD.partials(y[(k - 1) % N + 1], (k - 1) ÷ N + 1),
+        Val(N * N),
+    ))
+    return f, J
+end
+@inline _tendency_and_jacobian(::DonorJacobian, g, x) = (g(x), _jacobian_1m_linearized(g, x))
+@inline _tendency_and_jacobian(::CoupledDonorJacobian, g, x) =
+    (g(x), _jacobian_1m_relinearized(g, x))
+
 """
     _full_species_mask(x)
 
@@ -777,18 +808,17 @@ and the increment limiter through [`_jacobian_provider`](@ref),
     h = Δt / FT(nsub_eff)
     Lv_over_cp = TDI.TD.Parameters.LH_v0(tps) / TDI.TD.Parameters.cp_d(tps)
     Ls_over_cp = TDI.TD.Parameters.LH_s0(tps) / TDI.TD.Parameters.cp_d(tps)
-    jacobian = _jacobian_provider(mode.jacobian)
     mask = _species_mask(mode.jacobian, mode.growth)
 
     x = MicroState1M{FT}(q_lcl, q_icl, q_rai, q_sno)
     x₀ = x
     Tsub = T
     for _ in 1:nsub_eff
         g = Raw1MTendency(mp, tps, ρ, Tsub, q_tot)
-        f = g(x)
         x_prev = x
         if all(isfinite, x)
-            J = _apply_growth(mode.growth, jacobian(g, x))
+            f, J_raw = _tendency_and_jacobian(mode.jacobian, g, x)
+            J = _apply_growth(mode.growth, J_raw)
             z = mask(x)
             d = if all(isfinite, J)
                 _rosenbrock_update(x, f, J, z, h) - x
@@ -798,6 +828,7 @@ and the increment limiter through [`_jacobian_provider`](@ref),
             d = _apply_limiter(mode.limiter, x, d, ρ, Tsub, q_tot, Lv_over_cp, Ls_over_cp, tps)
             x = max.(x .+ d, 0)
         else
+            f = g(x)
             x = _euler_update(x, f, h)
         end
         Δ = x - x_prev