Add ClimaCore perf patch

CliMA · Nov 10, 2021 · c41ec7c · c41ec7c
1 parent 38d0ebb
commit c41ec7c
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/src/Fields.jl b/src/Fields.jl
@@ -54,3 +54,16 @@ function FieldFromNamedTuple(space, nt::NamedTuple)
     cmv(z) = nt
     return cmv.(CC.Fields.coordinate_field(space))
 end
+
+# https://github.com/CliMA/ClimaCore.jl/issues/275
+transform_broadcasted(bc::Base.Broadcast.Broadcasted{CC.Fields.FieldVectorStyle}, symb, axes) =
+    Base.Broadcast.Broadcasted(bc.f, map(arg -> transform_broadcasted(arg, symb, axes), bc.args), axes)
+transform_broadcasted(fv::CC.Fields.FieldVector, symb, axes) = parent(getproperty(fv, symb))
+transform_broadcasted(x, symb, axes) = x
+@inline function Base.copyto!(dest::CC.Fields.FieldVector, bc::Base.Broadcast.Broadcasted{CC.Fields.FieldVectorStyle})
+    for symb in propertynames(dest)
+        p = parent(getproperty(dest, symb))
+        Base.copyto!(p, transform_broadcasted(bc, symb, axes(p)))
+    end
+    return dest
+end