Vectors are a little slow to compile

[Sbozzolo]

I was investigating the cause of latency in ClimaAtmos by looking at config/model_configs/diagnostic_edmfx_test_box.yml.

The cache for this configuration takes 70 seconds to compile on my computer. I looked into this and found that the implicit cache takes 16 seconds. Of these 16 seconds, 10 are to compile a single function compute_kinetic.

So, I made a ClimaCore reproducer:

using ClimaCore.CommonSpaces
import ClimaCore
import ClimaCore: Fields, Geometry, Operators, Spaces
import LinearAlgebra: dot

space = ExtrudedCubedSphereSpace(; z_elem = 10, z_min = 0, z_max = 1, radius = 10, h_elem = 10, n_quad_points = 4, staggering = CellCenter(), )

ᶜuₕ = similar(zeros(space), Geometry.Covariant12Vector{Float64})
ᶠu₃ = similar(zeros(Spaces.face_space(space)), Geometry.Covariant3Vector{Float64})
ᶠu³ = similar(zeros(Spaces.face_space(space)), Geometry.Contravariant3Vector{Float64})
ᶠuₕ³ = similar(zeros(Spaces.face_space(space)), Geometry.Contravariant3Vector{Float64})
ᶜu = similar(zeros(space), Geometry.Covariant123Vector{Float64})
fill!(parent(ᶜuₕ), 0)
fill!(parent(ᶠu₃), 0)
fill!(parent(ᶠuₕ³), 0)
fill!(parent(ᶜu), 0)
fill!(parent(ᶠu³), 0)
ᶜK = zeros(space)

# Warm InterpolateF2C up
_ = @. Operators.InterpolateF2C()(Geometry.Covariant123Vector(ᶠu₃))

function mytest(ᶜu, ᶠuₕ³, ᶠu³, ᶜuₕ)
@time @. ᶜu = Geometry.Covariant123Vector(ᶜuₕ) + Operators.InterpolateF2C()(Geometry.Covariant123Vector(ᶠu₃))
@time @. ᶠu³ = ᶠuₕ³ + Geometry.Contravariant3Vector(ᶠu₃)
@time @. ᶜK = 1 / 2 * (
        dot(Geometry.Covariant123Vector(ᶜuₕ), Geometry.Contravariant123Vector(ᶜuₕ)) +
        Operators.InterpolateF2C()(dot(Geometry.Covariant123Vector(ᶠu₃), Geometry.Contravariant123Vector(ᶠu₃))) +
        2 * dot(Geometry.Contravariant123Vector(ᶜuₕ), Operators.InterpolateF2C()(Geometry.Covariant123Vector(ᶠu₃)))
    )
end

@time mytest(ᶜu, ᶠuₕ³, ᶠu³, ᶜuₕ)

This results in:

  1.578459 seconds (18.89 M allocations: 908.468 MiB, 33.42% gc time, 99.98% compilation time)
  1.419649 seconds (11.95 M allocations: 617.694 MiB, 21.31% gc time, 99.97% compilation time)
  6.432258 seconds (86.22 M allocations: 3.801 GiB, 27.18% gc time, 99.96% compilation time)

This

ᶜu = Geometry.Covariant123Vector(ᶜuₕ) + Operators.InterpolateF2C()(Geometry.Covariant123Vector(ᶠu₃))

takes 1.5 seconds to compile and leads to 1 GB of inference allocations. Note that I have already called the interpolation routine in the line before, so the second term should be already compiled. If I substitute the interpolate call with the result of the previous line, I get

  1.030550 seconds (11.70 M allocations: 605.150 MiB, 29.07% gc time, 99.95% compilation time)

Which tells me that having to infer the additional operator cost 50 % more time and inference allocations.

Compiling the full expression for the kinetic energy takes 6.5 seconds and has almost 4 GB of inference allocation.

This seems excessive for these relatively simple operations.

Note also this difference:

stored = @. Operators.InterpolateF2C()(Geometry.Covariant123Vector(ᶠu₃))
stored2 = @. Geometry.Covariant123Vector(ᶜuₕ)

@time @. ᶜu = stored2 + stored

The result is

  0.180751 seconds (1.30 M allocations: 71.706 MiB, 99.73% compilation time)

But

stored = @. Operators.InterpolateF2C()(Geometry.Covariant123Vector(ᶠu₃))
stored2 = @. Geometry.Covariant123Vector(ᶜuₕ)

@time @. ᶜu = Geometry.Covariant123Vector(ᶜuₕ) + stored

is

  0.915434 seconds (10.90 M allocations: 566.447 MiB, 26.39% gc time, 99.95% compilation time)

[charleskawczynski]

The original example missed passing a couple variables, and I think it's good to ensure that there are zero inference failures. Slightly adjusting this, I have:

using ClimaCore.CommonSpaces
import ClimaCore
import ClimaCore: Fields, Geometry, Operators, Spaces
import LinearAlgebra: dot

space = ExtrudedCubedSphereSpace(; z_elem = 10, z_min = 0, z_max = 1, radius = 10, h_elem = 10, n_quad_points = 4, staggering = CellCenter(), );

ᶜuₕ = similar(zeros(space), Geometry.Covariant12Vector{Float64});
ᶠu₃ = similar(zeros(Spaces.face_space(space)), Geometry.Covariant3Vector{Float64});
ᶠu³ = similar(zeros(Spaces.face_space(space)), Geometry.Contravariant3Vector{Float64});
ᶠuₕ³ = similar(zeros(Spaces.face_space(space)), Geometry.Contravariant3Vector{Float64});
ᶜu = similar(zeros(space), Geometry.Covariant123Vector{Float64});
fill!(parent(ᶜuₕ), 0);
fill!(parent(ᶠu₃), 0);
fill!(parent(ᶠuₕ³), 0);
fill!(parent(ᶜu), 0);
fill!(parent(ᶠu³), 0);
ᶜK = zeros(space);

# Warm InterpolateF2C up
_ = @. Operators.InterpolateF2C()(Geometry.Covariant123Vector(ᶠu₃));

function mytest(ᶜK, ᶜu, ᶠuₕ³, ᶠu³, ᶜuₕ, ᶠu₃)
    interp_f2c = Operators.InterpolateF2C()
    @. ᶜu = Geometry.Covariant123Vector(ᶜuₕ) + interp_f2c(Geometry.Covariant123Vector(ᶠu₃))
    @. ᶠu³ = ᶠuₕ³ + Geometry.Contravariant3Vector(ᶠu₃)
    @. ᶜK = 1 / 2 * (
        dot(Geometry.Covariant123Vector(ᶜuₕ), Geometry.Contravariant123Vector(ᶜuₕ)) +
        interp_f2c(dot(Geometry.Covariant123Vector(ᶠu₃), Geometry.Contravariant123Vector(ᶠu₃))) +
        2 * dot(Geometry.Contravariant123Vector(ᶜuₕ), interp_f2c(Geometry.Covariant123Vector(ᶠu₃)))
    )
    return nothing
end

@time mytest(ᶜK, ᶜu, ᶠuₕ³, ᶠu³, ᶜuₕ, ᶠu₃)
using JET
@test_opt mytest(ᶜK, ᶜu, ᶠuₕ³, ᶠu³, ᶜuₕ, ᶠu₃)

Which gives:

julia> @time mytest(ᶜK, ᶜu, ᶠuₕ³, ᶠu³, ᶜuₕ, ᶠu₃)
  1.543241 seconds (21.46 M allocations: 1.169 GiB, 16.39% gc time, 99.58% compilation time)

on Julia 1.10.

So, it's still kind of expensive in the sense that this is 3 broadcast expressions, and we have many of them.

[Sbozzolo]

This is what I get running the script in the post above with Julia 1.11:

julia> include("mytest.jl")
  8.779434 seconds (111.39 M allocations: 5.010 GiB, 23.90% gc time, 99.96% compilation time)
Test Passed

Same result if I manually evaluate the code directly in the repl instead of including the file.

This is what I get when I run with 1.10:

julia> include("mytest.jl")
  2.639313 seconds (22.27 M allocations: 1.236 GiB, 20.13% gc time, 99.90% compilation time)
Test Passed

[Sbozzolo]

I looked at this in more details and found that:

• time is dominated by inference time and not compile time
• inference time is dominated by getidx

using SnoopCompileCore

[...]

tinf = @snoop_inference mytest(ᶜK, ᶜu, ᶠuₕ³, ᶠu³, ᶜuₕ, ᶠu₃);
using SnoopCompile
fg = flamegraph(tinf);
using ProfileView
ProfileView.view(fg);

[charleskawczynski]

#2284 had a reasonable impact on this issue.

[charleskawczynski]

I now see:

julia 1.10

julia> using Revise; include("../latency_repro.jl")
Precompiling ClimaCore finished.
  1 dependency successfully precompiled in 8 seconds. 113 already precompiled.
  0.854593 seconds (7.18 M allocations: 427.729 MiB, 15.51% gc time, 99.26% compilation time)
  1.710 ms (0 allocations: 0 bytes)
Test Passed

julia 1.11

julia> using Revise; include("../latency_repro.jl")
[ Info: Precompiling Revise [295af30f-e4ad-537b-8983-00126c2a3abe] (cache misses: incompatible header (16))
Precompiling ClimaCore...
  21 dependencies successfully precompiled in 17 seconds. 119 already precompiled.
  2.366420 seconds (36.98 M allocations: 1.630 GiB, 20.92% gc time, 99.75% compilation time)
  1.649 ms (0 allocations: 0 bytes)
Precompiling JET...
  1 dependency successfully precompiled in 3 seconds. 34 already precompiled.
[ Info: Precompiling JETReviseExt [f03d6c7f-2be2-5ab6-8acb-706271f84fa0] (cache misses: wrong dep version loaded (2), incompatible header (6))
Test Passed

[charleskawczynski]

What OS are you running on, @Sbozzolo?

[Sbozzolo]

What OS are you running on, @Sbozzolo?

This is on the latest ClimaCore (on a different Linux machine compared to the previous results)

Julia 1.11

  3.019175 seconds (36.91 M allocations: 1.625 GiB, 28.41% gc time, 99.93% compilation time)

Julia 1.10

  1.123543 seconds (7.18 M allocations: 429.109 MiB, 14.47% gc time, 99.80% compilation time)

[Sbozzolo]

Julia 1.12beta2

  3.186160 seconds (41.39 M allocations: 1.586 GiB, 16.90% gc time, 99.94% compilation time)

Marginally worse than 1.11 🫠