Fix unrolling over field broadcast arguments

Author: dennisYatunin

src/Fields/Fields.jl

@@ -30,7 +30,7 @@ import ..Utilities: PlusHalf, half
 using ..RecursiveApply
 using ClimaComms
 import Adapt
-import UnrolledUtilities: unrolled_map, unrolled_findfirst
+import UnrolledUtilities: unrolled_map, unrolled_reduce, unrolled_findfirst
 
 import StaticArrays, LinearAlgebra, Statistics, InteractiveUtils
 

src/Fields/broadcast.jl

@@ -40,6 +40,12 @@ Base.Broadcast.BroadcastStyle(
     ::FieldStyle{DS2},
 ) where {DS1, DS2} = FieldStyle(Base.Broadcast.BroadcastStyle(DS1(), DS2()))
 
+# Override the recursive unrolling used in combine_styles (which can lead to
+# inference failures in broadcast expressions with more than 10 arguments) with
+# manual unrolling (which sacrifices latency but is always inferrable).
+Base.Broadcast.combine_styles(style1::AbstractFieldStyle, style2, styles...) =
+    unrolled_reduce(Base.Broadcast.combine_styles, (style1, style2, styles...))
+
 Base.Broadcast.broadcastable(field::Field) = field
 
 Base.eltype(bc::Base.Broadcast.Broadcasted{<:AbstractFieldStyle}) =