Merge #945

945: Inline `copyto!` to fix Ref allocations r=charleskawczynski a=charleskawczynski

This PR applies ``@inline`` to `copyto!`, which was found to fix allocations with `Ref`s inside broadcast expressions, and adds tests. I'm going to try this out with ClimaAtmos to see if one more minor change is needed, too.

This fixes our `bycolumn` allocations with `Ref`s from 24576 bytes to 0.

Closes #946.

Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

.buildkite/pipeline.yml

@@ -91,6 +91,16 @@ steps:
       slurm_mem: 64GB
       slurm_ntasks: 1
 
+  - label: ":computer: Field broadcast performance"
+    key: "cpu_field_broadcast_perf"
+    command:
+      - "julia --color=yes --project=test test/Fields/field_bc.jl"
+    agents:
+      config: cpu
+      queue: central
+      slurm_mem: 64GB
+      slurm_ntasks: 1
+
   - label: ":flower_playing_cards: unit tests"
     key: "gpu_unittests"
     command:

src/DataLayouts/broadcast.jl

@@ -426,7 +426,7 @@ end
     return dest
 end
 
-function Base.copyto!(
+@inline function Base.copyto!(
     dest::IJFH{S, Nij},
     bc::Union{IJFH{S, Nij}, Base.Broadcast.Broadcasted{<:IJFHStyle{Nij}}},
 ) where {S, Nij}
@@ -439,7 +439,7 @@ function Base.copyto!(
     return dest
 end
 
-function Base.copyto!(
+@inline function Base.copyto!(
     dest::IFH{S, Ni},
     bc::Union{IFH{S, Ni}, Base.Broadcast.Broadcasted{<:IFHStyle{Ni}}},
 ) where {S, Ni}
@@ -489,7 +489,7 @@ end
     return dest
 end
 
-function _serial_copyto!(
+@inline function _serial_copyto!(
     dest::VIFH{S, Ni},
     bc::Union{VIFH{S, Ni, A}, Base.Broadcast.Broadcasted{VIFHStyle{Ni, A}}},
 ) where {S, Ni, A}
@@ -503,7 +503,7 @@ function _serial_copyto!(
     return dest
 end
 
-function _threaded_copyto!(
+@inline function _threaded_copyto!(
     dest::VIFH{S, Ni},
     bc::Base.Broadcast.Broadcasted{VIFHStyle{Ni, A}},
 ) where {S, Ni, A}
@@ -522,14 +522,14 @@ function _threaded_copyto!(
     return dest
 end
 
-function Base.copyto!(
+@inline function Base.copyto!(
     dest::VIFH{S, Ni},
     source::VIFH{S, Ni, A},
 ) where {S, Ni, A}
     return _serial_copyto!(dest, source)
 end
 
-function Base.copyto!(
+@inline function Base.copyto!(
     dest::VIFH{S, Ni},
     bc::Base.Broadcast.Broadcasted{VIFHStyle{Ni, A}},
 ) where {S, Ni, A}
@@ -539,7 +539,7 @@ function Base.copyto!(
     return _serial_copyto!(dest, bc)
 end
 
-function _serial_copyto!(
+@inline function _serial_copyto!(
     dest::VIJFH{S, Nij},
     bc::Union{VIJFH{S, Nij, A}, Base.Broadcast.Broadcasted{VIJFHStyle{Nij, A}}},
 ) where {S, Nij, A}
@@ -553,7 +553,7 @@ function _serial_copyto!(
     return dest
 end
 
-function _threaded_copyto!(
+@inline function _threaded_copyto!(
     dest::VIJFH{S, Nij},
     bc::Base.Broadcast.Broadcasted{VIJFHStyle{Nij, A}},
 ) where {S, Nij, A}
@@ -572,14 +572,14 @@ function _threaded_copyto!(
     return dest
 end
 
-function Base.copyto!(
+@inline function Base.copyto!(
     dest::VIJFH{S, Nij},
     source::VIJFH{S, Nij, A},
 ) where {S, Nij, A}
     return _serial_copyto!(dest, source)
 end
 
-function Base.copyto!(
+@inline function Base.copyto!(
     dest::VIJFH{S, Nij},
     bc::Base.Broadcast.Broadcasted{VIJFHStyle{Nij, A}},
 ) where {S, Nij, A}

src/Fields/broadcast.jl

@@ -129,7 +129,7 @@ function Base.similar(
     return Field(similar(todata(bc), Eltype), axes(bc))
 end
 
-function Base.copyto!(
+@inline function Base.copyto!(
     dest::Field,
     bc::Base.Broadcast.Broadcasted{<:AbstractFieldStyle},
 )

test/Fields/field_bc.jl

@@ -0,0 +1,53 @@
+using Test
+using StaticArrays, IntervalSets
+import ClimaCore
+import ClimaCore.Utilities: PlusHalf
+import ClimaCore.DataLayouts: IJFH
+import ClimaCore:
+    Fields, slab, Domains, Topologies, Meshes, Operators, Spaces, Geometry
+
+using LinearAlgebra: norm
+using Statistics: mean
+using ForwardDiff
+
+function FieldFromNamedTuple(space, nt::NamedTuple)
+    cmv(z) = nt
+    return cmv.(Fields.coordinate_field(space))
+end
+
+include(joinpath(@__DIR__, "util_spaces.jl"))
+
+# https://github.com/CliMA/ClimaCore.jl/issues/946
+@testset "Allocations with broadcasting Refs" begin
+    FT = Float64
+    function foo!(Yx::Fields.Field)
+        Yx .= Ref(1) .+ Yx
+        return nothing
+    end
+    function foocolumn!(Yx::Fields.Field)
+        Fields.bycolumn(axes(Yx)) do colidx
+            Yx[colidx] .= Ref(1) .+ Yx[colidx]
+            nothing
+        end
+        return nothing
+    end
+    for space in all_spaces(FT)
+        (
+            space isa Spaces.ExtrudedFiniteDifferenceSpace ||
+            space isa Spaces.SpectralElementSpace1D ||
+            space isa Spaces.SpectralElementSpace2D
+        ) || continue
+        Y = FieldFromNamedTuple(space, (; x = FT(2)))
+
+        # Plain broadcast
+        Yx = Y.x
+        foo!(Yx) # compile first
+        p = @allocated foo!(Yx)
+        @test p == 0
+
+        # bycolumn
+        foocolumn!(Yx) # compile first
+        p = @allocated foocolumn!(Yx)
+        @test p == 0
+    end
+end