autodiff and allocation ahhhhh

Author: szy21

src/cache/precomputed_quantities.jl

@@ -438,82 +438,83 @@ function thermo_vars(::NonEquilMoistModel, ::Any, ᶜY, K, Φ)
     return (; e_int, q_pt = TD.PhasePartition(q_pt_args...))
 end
 
-"""
-    thermo_state_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ)
+# Individual getter functions for thermodynamic state variables.
+# These avoid struct construction which can cause allocations on GPU.
 
-Compute grid-scale thermodynamic state variables in one call.
+# Temperature getters
+function T_gs(thermo_params, ::DryModel, ::Any, ᶜY, K, Φ, ρ)
+    e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
+    return TD.air_temperature(thermo_params, e_int)
+end
 
-Returns a NamedTuple with:
-- `T`: Temperature
-- `q_tot_safe`: Total specific humidity (clipped to non-negative)
-- `q_liq_rai`: Liquid + rain specific humidity
-- `q_ice_sno`: Ice + snow specific humidity
+function T_gs(thermo_params, ::EquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
+    e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
+    q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)
+    sa_result = TD.saturation_adjustment(thermo_params, TD.ρe(), ρ, e_int, q_tot)
+    return sa_result.T
+end
 
-For DryModel: T from e_int, all q's are zero
-For EquilMoistModel: Uses `saturation_adjustment` once for all quantities
-For NonEquilMoistModel: T from e_int and prognostic q's, q's from prognostic values
-"""
-function thermo_state_gs(
-    thermo_params,
-    moisture_model::DryModel,
-    microphysics_model,
-    ᶜY,
-    K,
-    Φ,
-    ρ,
-)
+function T_gs(thermo_params, ::NonEquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
+    e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
+    q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)
+    q_liq_rai = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ)
+    q_ice_sno = specific(ᶜY.ρq_ice, ᶜY.ρ) + specific(ᶜY.ρq_sno, ᶜY.ρ)
+    return TD.air_temperature(thermo_params, e_int, q_tot, q_liq_rai, q_ice_sno)
+end
+
+# q_tot_safe getters
+function q_tot_safe_gs(thermo_params, ::DryModel, ::Any, ᶜY, K, Φ, ρ)
     e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
     T = TD.air_temperature(thermo_params, e_int)
-    FT = eltype(thermo_params)
-    return (;
-        T = T,
-        q_tot_safe = zero(FT),
-        q_liq_rai = zero(FT),
-        q_ice_sno = zero(FT),
-    )
+    return zero(T)  # Use zero(T) for autodiff compatibility
 end
 
-function thermo_state_gs(
-    thermo_params,
-    moisture_model::EquilMoistModel,
-    microphysics_model,
-    ᶜY,
-    K,
-    Φ,
-    ρ,
-)
+function q_tot_safe_gs(thermo_params, ::EquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
+    q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)
+    return max(0, q_tot)
+end
+
+function q_tot_safe_gs(thermo_params, ::NonEquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
+    q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)
+    return max(0, q_tot)
+end
+
+# q_liq_rai getters
+function q_liq_rai_gs(thermo_params, ::DryModel, ::Any, ᶜY, K, Φ, ρ)
+    e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
+    T = TD.air_temperature(thermo_params, e_int)
+    return zero(T)  # Use zero(T) for autodiff compatibility
+end
+
+function q_liq_rai_gs(thermo_params, ::EquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
     e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
     q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)
-    # Use saturation_adjustment once for all quantities
     sa_result = TD.saturation_adjustment(thermo_params, TD.ρe(), ρ, e_int, q_tot)
-    return (;
-        T = sa_result.T,
-        q_tot_safe = max(0, q_tot),
-        q_liq_rai = sa_result.q_liq,
-        q_ice_sno = sa_result.q_ice,
-    )
+    return sa_result.q_liq
 end
 
-function thermo_state_gs(
-    thermo_params,
-    moisture_model::NonEquilMoistModel,
-    microphysics_model,
-    ᶜY,
-    K,
-    Φ,
-    ρ,
-)
+function q_liq_rai_gs(thermo_params, ::NonEquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
+    q_liq_rai = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ)
+    return max(0, q_liq_rai)
+end
+
+# q_ice_sno getters
+function q_ice_sno_gs(thermo_params, ::DryModel, ::Any, ᶜY, K, Φ, ρ)
+    e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
+    T = TD.air_temperature(thermo_params, e_int)
+    return zero(T)  # Use zero(T) for autodiff compatibility
+end
+
+function q_ice_sno_gs(thermo_params, ::EquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
     e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ
     q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ)
-    q_liq_rai = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ)
+    sa_result = TD.saturation_adjustment(thermo_params, TD.ρe(), ρ, e_int, q_tot)
+    return sa_result.q_ice
+end
+
+function q_ice_sno_gs(thermo_params, ::NonEquilMoistModel, ::Any, ᶜY, K, Φ, ρ)
     q_ice_sno = specific(ᶜY.ρq_ice, ᶜY.ρ) + specific(ᶜY.ρq_sno, ᶜY.ρ)
-    T = TD.air_temperature(thermo_params, e_int, q_tot, q_liq_rai, q_ice_sno)
-    return (;
-        T = T,
-        q_tot_safe = max(0, q_tot),
-        q_liq_rai = max(0, q_liq_rai),
-        q_ice_sno = max(0, q_ice_sno),
-    )
+    return max(0, q_ice_sno)
 end
 
 function eddy_diffusivity_coefficient_H(D₀, H, z_sfc, z)
@@ -576,13 +577,13 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
         # @. ᶜK += Y.c.ρtke / Y.c.ρ
         # TODO: We should think more about these increments before we use them.
     end
-    # Compute all thermodynamic state variables in one call (avoids multiple saturation_adjustment calls)
-    ᶜthermo_state = p.scratch.ᶜtemp_thermo_state
-    @. ᶜthermo_state = thermo_state_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
-    @. ᶜT = ᶜthermo_state.T
-    @. ᶜq_tot_safe = ᶜthermo_state.q_tot_safe
-    @. ᶜq_liq_rai = ᶜthermo_state.q_liq_rai
-    @. ᶜq_ice_sno = ᶜthermo_state.q_ice_sno
+    # Compute thermodynamic state variables using individual getter functions.
+    # Note: For EquilMoistModel, this calls saturation_adjustment 3 times per grid point
+    # (T, q_liq_rai, q_ice_sno each call it; q_tot_safe doesn't need it).
+    @. ᶜT = T_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
+    @. ᶜq_tot_safe = q_tot_safe_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
+    @. ᶜq_liq_rai = q_liq_rai_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
+    @. ᶜq_ice_sno = q_ice_sno_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
     ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ))
     @. ᶜh_tot = TD.total_enthalpy(thermo_params, ᶜe_tot, ᶜT, ᶜq_tot_safe, ᶜq_liq_rai, ᶜq_ice_sno)
     @. ᶜp = TD.air_pressure(thermo_params, ᶜT, Y.c.ρ, ᶜq_tot_safe, ᶜq_liq_rai, ᶜq_ice_sno)

src/cache/temporary_quantities.jl

@@ -29,11 +29,6 @@ function temporary_quantities(Y, atmos)
 
     FT = Spaces.undertype(center_space)
     uvw_vec = UVW(FT(0), FT(0), FT(0))
-    # Type for thermo state: (T, q_tot_safe, q_liq_rai, q_ice_sno)
-    ThermoStateType = NamedTuple{
-        (:T, :q_tot_safe, :q_liq_rai, :q_ice_sno),
-        NTuple{4, FT},
-    }
     return (;
         ᶠtemp_scalar = Fields.Field(FT, face_space), # ᶠp, ᶠρK_h
         ᶠtemp_scalar_2 = Fields.Field(FT, face_space), # ᶠρK_u
@@ -44,7 +39,6 @@ function temporary_quantities(Y, atmos)
         ᶜtemp_scalar_5 = Fields.Field(FT, center_space),
         ᶜtemp_scalar_6 = Fields.Field(FT, center_space),
         ᶜtemp_scalar_7 = Fields.Field(FT, center_space),
-        ᶜtemp_thermo_state = Fields.Field(ThermoStateType, center_space), # thermo_state_gs result
         ᶠtemp_field_level = Fields.level(Fields.Field(FT, face_space), half),
         temp_field_level = Fields.level(Fields.Field(FT, center_space), 1),
         temp_field_level_2 = Fields.level(Fields.Field(FT, center_space), 1),

src/solver/types.jl

@@ -379,7 +379,6 @@ function MixingLength(master, wall, tke, buoy, l_grid)
     return MixingLength(promote(master, wall, tke, buoy, l_grid)...)
 end
 
-
 abstract type AbstractEDMF end
 
 """