fix(P3): Manually unroll LogExpFunctions.logsumexp

- To improve compilation of some GPU kernels,
provide an unrolled implementation of logsumexp

Author: haakon-e

src/P3.jl

@@ -11,6 +11,7 @@ import SpecialFunctions as SF
 import RootSolvers as RS
 import LogExpFunctions
 import StaticArrays as SA
+import UnrolledUtilities as UU
 
 import ClimaParams as CP
 
@@ -20,6 +21,7 @@ import CloudMicrophysics.Common as CO
 import CloudMicrophysics.DistributionTools as DT
 import CloudMicrophysics.HetIceNucleation as CM_HetIce
 import CloudMicrophysics.Microphysics2M as CM2
+import CloudMicrophysics.Utilities as UT
 import CloudMicrophysics: ShowMethods
 
 include("P3_particle_properties.jl")

src/P3_size_distribution.jl

@@ -139,11 +139,12 @@ Compute `log(∫_0^∞ Dⁿ m(D) N′(D) dD)` given the `state` and `logλ`.
 """
 function logmass_gamma_moment(state::P3State, μ, logλ; n = 0)
     segments = get_segments(state)
-    return LogExpFunctions.logsumexp(
-        let (D_min, D_max) = segment, (a, b) = ice_mass_coeffs(state, (D_min + D_max) / 2)
-            loggamma_inc_moment(D_min, D_max, μ, logλ, b + n, a)
-        end for segment in segments
-    )
+    moments = UU.unrolled_map(segments) do segment
+        (D_min, D_max) = segment
+        (a, b) = ice_mass_coeffs(state, (D_min + D_max) / 2)
+        loggamma_inc_moment(D_min, D_max, μ, logλ, b + n, a)
+    end
+    return UT.unrolled_logsumexp(moments)
 end
 
 """

src/Utilities.jl

@@ -6,7 +6,10 @@ Contains pure numerical operations with no physics dependencies.
 """
 module Utilities
 
+import UnrolledUtilities as UU
+
 export clamp_to_nonneg, ϵ_numerics, ϵ_numerics_2M_M, ϵ_numerics_2M_N
+export unrolled_logsumexp
 
 """
     clamp_to_nonneg(x)
@@ -44,4 +47,77 @@ Numerical epsilon for 2-moment number calculations.
 """
 @inline ϵ_numerics_2M_N(FT) = eps(FT)
 
+
+"""
+    unrolled_logsumexp(x)
+
+Compute `log(sum(exp, x))` in a statically unrolled fashion.
+
+This method uses [`UnrolledUtilities`](https://github.com/CliMA/UnrolledUtilities.jl)
+to produce fully unrolled code with no dynamic dispatch or reductions,
+making it transparent to GPU compilers.
+
+The standard shift-by-max trick is used for numerical stability.
+
+Note: This code is two-pass (find max, then sum shifted exponentials). 
+LogExpFunctions.jl implements a one-pass version, but is not unrolled,
+so may result in more complicated GPU code.
+
+## Extended help
+
+Other implementation options were considered, detailed below, and may be revisited in the future.
+For now, this implementation is sufficient.
+
+### Naive implementation
+
+This is the most straightforward implementation, but it is not numerically stable.
+
+```julia
+log(UU.unrolled_sum(exp, x))
+```
+
+### One-pass unrolled implementation
+
+This is reaches into LogExpFunctions.jl internals,
+
+```julia
+FT = eltype(x)
+return LogExpFunctions._logsumexp_onepass_result(
+    UU.unrolled_reduce(LogExpFunctions._logsumexp_onepass_op, x, (FT(-Inf), zero(FT)))
+)
+```
+
+### Dispatch-wrapper for reduce
+
+Pass a wrapper to compile to unrolled reduce
+```julia
+# Note: This is a sketch, not tested
+struct UnrolledWrapper{T}
+    x::T
+end
+Base.iterate(w::UnrolledWrapper) = iterate(w.x)
+Base.iterate(w::UnrolledWrapper, state) = iterate(w.x, state)
+Base.length(w::UnrolledWrapper) = length(w.x)
+Base.eltype(w::UnrolledWrapper) = eltype(w.x)
+Base.reduce(op, w::UnrolledWrapper) = UU.unrolled_reduce(op, w.x)  # use unrolled reduce
+# ... then call:
+LogExpFunctions.logsumexp(UnrolledWrapper(x))
+```
+"""
+function unrolled_logsumexp(x)
+    # Find the maximum (ps: if any element is NaN, then xmax = NaN)
+    xmax = UU.unrolled_maximum(x)
+
+    # Handle non-finite values: if xmax is +Inf or -Inf or NaN, return it directly
+    # (avoids Inf - Inf = NaN and x - NaN = NaN in the shifted exponentials below)
+    isfinite(xmax) || return xmax
+
+    # Sum shifted exponentials
+    shifted_exp(xi) = exp(xi - xmax)
+    s = UU.unrolled_sum(shifted_exp, x)
+
+    return xmax + log(s)
+end
+
+
 end # module