Performance roadmap

[charleskawczynski]

This issue is a continuation of #635, but I'm excluding some items (some addressed, and others which I've explained in #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

Preview Give feedback

1. 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. Fuse dss calls #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 Fuse dynamic and tracer tendency dss calls #2758
2. 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
3. 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.
4. Use * over / if not being optimized low impact, very little effort Use rmul over rdiv ClimaCore.jl#1496
5. Custom writing highly expensive kernels. We could try writing custom kernels to improve performance of specific kernels.
6. Improve RRTMGP performance (Allow for computing multiple fluxes RRTMGP.jl#389 seems important)
7. Add option for asynchronous f! and j! computation ClimaTimeSteppers.jl#233
   Performance +1
8. Add support for async T_exp! and T_lim! ClimaTimeSteppers.jl#247
   +0
9. Collect kernel stats for all cuda kernels on A100 Add kernel statistics collection to auto_launch ClimaCore.jl#1729, Add benchmark utility extension ClimaTimeSteppers.jl#277
10. Create prototype example to demo performance benefit of @fuse (without implementing the macro)
11. Reduce number of DSS calls Fuse dss calls #2689
12. Group dss calls in dss_hyperdiffusion_tendency! and dss_tracer_hyperdiffusion_tendency! by fusing T_exp! and T_lim!. Done in Fuse dynamic and tracer tendency dss calls #2758.
13. Always inline in ClimaCore's broadcast kernels Always inline ClimaCore.jl#1647
14. Implement similar-space point-wise fused kernels Add support for MultiBroadcastFusion ClimaCore.jl#1641
15. Implement point-wise fused kernels with both center and face space compatibility ClimaCore.jl#1739
    +0
    
16. Implement similar-space Finite-Difference fused kernels ClimaCore.jl#1740
    enhancement +1
    
17. Implement similar-space FD+pointwise fused kernels ClimaCore.jl#1741
    enhancement +1
    
18. Implement similar-space pointwise fused parallel kernels ClimaCore.jl#1742
    enhancement +1
    
19. Implement mixed-space pointwise fused parallel ClimaCore.jl#1743
    enhancement +1
    
20. Implement mixed-space FD fused parallel kernels ClimaCore.jl#1744
    enhancement +1
    
21. Implement mixed-space FD+pointwise fused parallel kernels ClimaCore.jl#1745
    enhancement +1
    
22. Use local memory in cuda band-matrix solve ClimaCore.jl#1738
    +0
    
23. Use shared/local memory in FD stencils kernels ClimaCore.jl#1746
    enhancement +1
    
24. Rename and split set_precomputed_quantities! ClimaTimeSteppers.jl#270
    Performance +1
    
25. Implement a CachedField type ClimaCore.jl#1747
    enhancement +1
26. Remove field dimension in DataLayouts ClimaCore.jl#1748
    enhancement +1
27. Add size to parameter space of data layout ClimaCore.jl#11
    +0
28. Implement parallel band matrix solve ClimaCore.jl#1753
    enhancement +1
29. Remove methods that append local_geometry_field to broadcasted objects ClimaCore.jl#1754
    enhancement +1
30. Fix uncoalesced memory reads ClimaCore.jl#1910
    1 of 1
    enhancement performance +2
    

[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!