CliMA / CliMA/ClimaCore.jl

The hidden cost of broadcasting

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

Description

Our CI runners spend up to 90% of their time compiling code, and latency has been a major pain point.

This issue collects some experiments that quantify the overhead due to using broadcasting in terms of compile time.

I am not sure what lesson to learn from here exactly, except that compile time scales more than linearly with the number of terms in a broadcasted expression.

So, reducing the size of expressions reduces compile time. This, of course comes at cost of increased runtime.

### The base case
```julia
using ClimaCore.CommonSpaces
import ClimaCore: Operators

function mytest(grad, div, ofie, fie)
return nothing
end

space = ExtrudedCubedSphereSpace(; z_elem = 10, z_min = 0, z_max = 1, radius = 10, h_elem = 10, n_quad_points = 4, staggering = CellCenter(), )
field = zeros(space)
ofield = zeros(space)
mydiv = Operators.Divergence()
mygrad = Operators.Gradient()

@time mytest(mygrad, mydiv, ofield, field);
@time mytest(mygrad, mydiv, ofield, field);
```

This script sets the stage up and does nothing. Save as `mytest.jl` and run with
```
time julia --project=.buildkite -e 'include("mytest.jl")'
```

On my computer, this outputs
```
0.000059 seconds (109 allocations: 5.469 KiB, 67.57% compilation time)
0.000004 seconds

________________________________________________________
Executed in 12.20 secs fish external
usr time 12.14 secs 260.00 micros 12.14 secs
sys time 0.27 secs 86.00 micros 0.27 secs
```
Total runtime is around 12 seconds, compiling and running `mytest` takes no time.

### A minimal test

Now, let's do a base `copyto`

```julia
function mytest(grad, div, ofie, fie)
@. ofield = fie
end
```

That outputs
```
0.098856 seconds (568.05 k allocations: 29.769 MiB, 99.87% compilation time)
0.000139 seconds (1 allocation: 32 bytes)
```

Compiling `copyto` takes approximately 0.1 seconds.

Let's skip the broadcasting:
```julia
function mytest(grad, div, p2, p1)
@inbounds for i in eachindex(p1)
p2[i] = p1[i]
end
end
@time mytest(mygrad, mydiv, parent(ofield), parent(field));
@time mytest(mygrad, mydiv, parent(ofield), parent(field));
```

Skipping broadcasting with fields is about 8 times faster:
```
0.014571 seconds (5.52 k allocations: 274.203 KiB, 99.31% compilation time)
0.000080 seconds
```

### A more complex case (10 sins)

Now, let's repeat the test but have a broadcasted expression with 10 terms

```julia
function mytest(grad, div, ofie, fie)
@. ofie = fie |> sin |> sin |> sin |> sin |> sin |> sin |> sin |> sin |> sin |> sin
end
```

```
0.140623 seconds (873.41 k allocations: 46.521 MiB, 95.51% compilation time)
0.006265 seconds (10 allocations: 80 bytes)
```

But if we don't go through broadcasting:

```julia
function mytest(grad, div, p2, p1)
@inbounds for i in eachindex(p1)
p2[i] = p1[i] |> sin |> sin |> sin |> sin |> sin |> sin |> sin |> sin |> sin |> sin
end
end
```

Which is 2x as fast

```
0.063442 seconds (130.20 k allocations: 5.775 MiB, 93.39% compilation time)
0.004117 seconds
```

### More sins

Now, let's add more `|> sin`

20 sins: `0.289147 seconds (1.64 M allocations: 97.530 MiB, 14.11% gc time, 95.32% compilation time)`
30 sins: `0.375251 seconds (2.63 M allocations: 167.602 MiB, 9.92% gc time, 95.75% compilation time)`
40 sins: `0.675250 seconds (3.86 M allocations: 287.396 MiB, 25.75% gc time, 96.15% compilation time)`
50 sins: `0.874510 seconds (5.33 M allocations: 455.638 MiB, 22.08% gc time, 96.42% compilation time)`
60 sins: `1.123960 seconds (7.05 M allocations: 672.522 MiB, 18.72% gc time, 96.66% compilation time)`
70 sins: `1.583039 seconds (9.03 M allocations: 964.154 MiB, 22.80% gc time, 96.91% compilation time)`
80 sins: `1.931034 seconds (11.26 M allocations: 1.268 GiB, 19.88% gc time, 97.15% compilation time)`
90 sins: `2.384610 seconds (13.75 M allocations: 1.635 GiB, 12.39% gc time, 97.64% compilation time)`
100 sins: `2.760357 seconds (16.48 M allocations: 2.040 GiB, 12.34% gc time, 97.80% compilation time)`

20 sins, Non-broadcated: `0.076848 seconds (132.06 k allocations: 5.870 MiB, 81.73% compilation time)`
100 sins, Non-broadcasted: `0.125252 seconds (147.02 k allocations: 6.531 MiB, 50.63% compilation time)`

How much is the ClimaCore overhead? Comparing with just using Arrays instead of fields
100 sins, Broadcasted with arrays: `1.171034 seconds (10.14 M allocations: 576.629 MiB, 19.21% gc time, 99.38% compilation time)`

So, it seems that ClimaCore adds a factor of 2 overhead with respect to just using arrays.

### Operators

Now, let's try using operators. We will call multiple`div(grad())`.

```julia
function mytest(grad, div, ofie, fie)
@. ofield = div(grad(fie))
# More divgrads as in
# @. ofield = div(grad(div(grad(fie))))
end
```

1 `divgrad`: `1.654600 seconds (13.41 M allocations: 658.225 MiB, 17.64% gc time, 99.88% compilation time)`
2 `divgrad`s: `2.397050 seconds (20.42 M allocations: 977.172 MiB, 23.34% gc time, 99.86% compilation time)`
3 `divgrad`s: `3.279898 seconds (30.58 M allocations: 1.400 GiB, 26.31% gc time, 99.85% compilation time)`
6 `divgrads`s: `6.930279 seconds (80.85 M allocations: 3.558 GiB, 25.71% gc time, 99.87% compilation time)`
12 `divgrads`s: `22.216704 seconds (270.54 M allocations: 11.681 GiB, 25.57% gc time, 99.92% compilation time)`
24 `divgrads`: ` 79.178018 seconds (1.01 G allocations: 43.405 GiB, 26.82% gc time, 99.95% compilation time)`

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by saving the supplied experiments as mytest.jl and running them with `time julia --project=.buildkite -e 'include("mytest.jl")'`. Compare broadcasted field expressions, explicit loops, array operations, and nested div/grad cases using the reported timings. The issue does not define a specific change or completion target, so a concrete optimization goal would need to be established first.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.