CliMA / CliMA/ClimaCore.jl

Vectors are a little slow to compile

Open
#2,215 8 comments 0 reactions 0 assignees View on GitHub
bug Latency
Dominant language
Julia
Stars
117
Forks
19
Avg merge
3d 4h
Merged PRs (30d)
41

Description

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:
```julia
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
```julia
ᶜ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:
```julia
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
```julia
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)
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the reproducer in the issue and the `compute_kinetic` path identified through `config/model_configs/diagnostic_edmfx_test_box.yml`. Compare compilation time and inference allocations for `Operators.InterpolateF2C`, vector construction, and the full kinetic-energy expression. Done means identifying and resolving the excessive compilation cost while preserving the demonstrated operations.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.