Debugging GPU issue in saturation adjustment

Author: dennisYatunin

src/cache/precomputed_quantities.jl

@@ -356,6 +356,38 @@ function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰, ᶠu₃⁰, turbconv_model)
     return nothing
 end
 
+# TODO: Remove this.
+@inline function PhaseEquil_ρeq_temporary_overwrite(
+    param_set::TD.APS,
+    ρ::FT,
+    e_int::FT,
+    q_tot::FT,
+    maxiter::IT = nothing,
+    relative_temperature_tol::TT = nothing,
+    ::Type{sat_adjust_method} = TD.RS.NewtonsMethod,
+    T_guess::Union{FT, Nothing} = nothing,
+) where {FT <: Real, sat_adjust_method, IT <: TD.ITERTYPE, TT <: TD.TOLTYPE}
+    maxiter === nothing && (maxiter = 8)
+    relative_temperature_tol === nothing &&
+        (relative_temperature_tol = FT(1e-4))
+    phase_type = TD.PhaseEquil{FT}
+    q_tot_safe = clamp(q_tot, 0, 1) # clamp(q_tot, FT(0), FT(1))
+    T = TD.saturation_adjustment(
+        sat_adjust_method,
+        param_set,
+        e_int,
+        ρ,
+        q_tot_safe,
+        phase_type,
+        maxiter,
+        relative_temperature_tol,
+        T_guess,
+    )
+    q_pt = TD.PhasePartition_equil(param_set, T, ρ, q_tot_safe, phase_type)
+    p = TD.air_pressure(param_set, T, ρ, q_pt)
+    return TD.PhaseEquil{FT}(ρ, p, e_int, q_tot_safe, T)
+end
+
 function thermo_state(
     thermo_params;
     ρ = nothing,
@@ -374,7 +406,7 @@ function thermo_state(
     get_ts(ρ::Real, ::Nothing, ::Nothing, e_int::Real, ::Nothing, ::Nothing) =
         TD.PhaseDry_ρe(thermo_params, ρ, e_int)
     get_ts(ρ::Real, ::Nothing, ::Nothing, e_int::Real, q_tot::Real, ::Nothing) =
-        TD.PhaseEquil_ρeq(
+        PhaseEquil_ρeq_temporary_overwrite( # TD.PhaseEquil_ρeq(
             thermo_params,
             ρ,
             e_int,
@@ -490,6 +522,7 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
         # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ
         # TODO: We should think more about these increments before we use them.
     end
+    iszero(t) && @show(eltype(Y.c.ρ)) # TODO: Remove this.
     @. ᶜts = ts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
     @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
     @. ᶜh_tot = TD.total_specific_enthalpy(