Performance roadmap

[charleskawczynski]

This issue is a continuation of CliMA/ClimaAtmos.jl#635, but I'm excluding some items (some addressed, and others which I've explained in CliMA/ClimaAtmos.jl#635) to reduce the noise.

Memory access patterns

We should make sure that we inline all kernels, use shared/local memory when possible, and ensure we have coalesced reads/writes.

Reducing loads and stores

The primary point of improving performance beyond our current state is by reducing the number of memory loads and stores. One way to do that is by fusing operations, which can allow the compiler to hoist (and eliminate) memory loads/stores. Another way is to explicitly pass less data through broadcast expressions (where possible).

There are a few different options / paths to capturing some of this performance that we've left on the table, and each approach has its limitations, pros and cons:

• Optimize / fuse simple fieldvector operations in timestepper
  • Pros: should be somewhat simple / straightforward, no changes to ClimaAtmos
  • Cons: impact may be about 10%
• FD operators read data redundantly (i.e. the +half and -half values), we could improve this by first reading into shared memory, similar to spectral element operators.
  • Pros: could improve vertical kernels by as much as 2x, no ClimaAtmos changes. This is probably a good idea to implement at some point.
  • Cons: Does not impact horizontal kernels.
• Fusing operations (e.g., @fuse begin @. a = b; @. c = d end)
  • Pros:
    • Will improve CPU/GPU performance
    • It's an optimization that can be done incrementally (lower risk of failure, easy to prototype) and apply this to ClimaAtmos / other repos
  • Cons:
    • Requires changing ClimaAtmos (we have many broadcast expressions)
    • This will (likely) be limited by similar BC expressions (cannot mix cell center / face, or horizontal / vertical, union-splitting), which will limit the number of fuses we can perform
  • Use lazy evaluation of BC expressions
  • Pros:
    • There's potential for fusing many operations this way (maybe even beyond cell center / face)
    • This could also help us reduce code duplication (e.g., reusing functions/expressions for tendencies and diagnostics)
    • The code could become more modular, and this may allow us to more easily unit test individual kernels
  • Cons:
    • There is risk of exploding compilation times
    • We'll need some way to recover useable stack traces (e.g., enabling eager execution)
    • There is a risk that profiling the code may becomes more complicated
    • A significant portion of ClimaAtmos will need to be re-written to a functional-style approach
  • Reduce reads / writes from LocalGeometry (if possible)
    • Pros
    • (possibly) no changes to ClimaAtmos required
    • Cons
    • We need a prototype to prove that there is performance available
    • This optimization will only improve kernels that use the LocalGeometry
    • Optimizations beyond this will nullify this effort $^1$

$^1$ It's important to note that one can nullify the other. That is, if we perform two optimizations: 1) eliminate loading X from kernel A and eliminate loading Y from kernel B 2) fuse kernels A and B

we could end up with the same number of loads and stores if we had only performed optimization 1) or 2) alone.

Removing unnecessary work

We can remove unnecessary work, e.g., in precomputed quantities, or using a caching system

Parallelism

There are other optimizations we can perform, which can also have a notable impact. For example, parallelizing work, reducing allocations to reduce the frequency of GC, reducing MPI communication, and emitting more efficient low-level code. Below is a list of some of these items:

Scaling

Minimize number of dss calls, and gc calls.

Misc

There are other miscilaneous items, specified in the task list.

### Tasks
- [x] Minimize the number of DSS calls. There are some ungrouped weighted DSS calls that can easily be grouped together. This is extremely low hanging fruit, and should probably be done asap. https://github.com/CliMA/ClimaAtmos.jl/pull/2689. We can still group dss calls in `dss_hyperdiffusion_tendency!` and `dss_tracer_hyperdiffusion_tendency!` if we fuse `T_exp!` and `T_lim!`. Done in https://github.com/CliMA/ClimaAtmos.jl/pull/2758
- [ ] GC pauses can happen at different places on different processes, which will hamper scaling efficiency (as processes will wait for other processes running GC): once we've reduced allocations we can disable the GC and trigger it manually intermittently. `GC.enable(false)` and `GC.gc()` to trigger manually
- [ ] Measure performance of new diagnostics. We've added a callbacks flame graph, which shows time spent during diagnostics, but we should use more realistic parameters, like resolution and frequency of different diagnostics.
- [x] Use `*` over `/` if not being optimized **low impact, very little effort** https://github.com/CliMA/ClimaCore.jl/pull/1496
- [ ] Custom writing highly expensive kernels. We could try writing custom kernels to improve performance of specific kernels.
- [ ] Improve RRTMGP performance (https://github.com/CliMA/RRTMGP.jl/issues/389 seems important)
- [ ] https://github.com/CliMA/ClimaTimeSteppers.jl/issues/233
- [ ] https://github.com/CliMA/ClimaTimeSteppers.jl/issues/247
- [x] Collect kernel stats for all cuda kernels on A100 https://github.com/CliMA/ClimaCore.jl/pull/1729, https://github.com/CliMA/ClimaTimeSteppers.jl/pull/277
- [x] Create prototype example to demo performance benefit of `@fuse` (without implementing the macro)
- [x] Reduce number of DSS calls https://github.com/CliMA/ClimaAtmos.jl/pull/2689
- [x] Group dss calls in `dss_hyperdiffusion_tendency!` and `dss_tracer_hyperdiffusion_tendency!` by fusing `T_exp!` and `T_lim!`. Done in https://github.com/CliMA/ClimaAtmos.jl/pull/2758.
- [x] Always inline in ClimaCore's broadcast kernels https://github.com/CliMA/ClimaCore.jl/pull/1647
- [x] Implement similar-space point-wise fused kernels https://github.com/CliMA/ClimaCore.jl/pull/1641
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1739
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1740
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1741
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1742
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1743
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1744
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1745
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1738
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1746
- [ ] https://github.com/CliMA/ClimaTimeSteppers.jl/issues/270
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1747
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1748
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/11
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1753
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1754
- [ ] https://github.com/CliMA/ClimaCore.jl/issues/1910

[charleskawczynski]

I've removed the prototype (as we already have developed https://github.com/CliMA/MultiBroadcastFusion.jl, which has performance tests) to reduce the noise in this issue.

I'm pleasantly surprised that the generic/recursive pattern appears (somehow) more performant than the hard-coded one, but I'll take it!

[tapios]

Really nice and helpful. Thank you!