Some improvements to interfaces and performance

[valeriabarra]

Purpose

This SDI proposes some clean-up of some interfaces in the current ClimaCoupler API so to improve on some of its aspects, such as type-safety and overall compiler optimization and performance.

Cost/Benefits/Risks

Costs
The costs could be shared across multiple co-authors that can concurrently work on separate files following the same improvement patterns or on the same pattern across different files.

Benefits
Benefits include:

• an easier to read and follow interface
• type-safety
• reduced allocations
• performance speed-up
• dependencies stability
• identification of missing interfaces in upstream packages (e.g., ClimaCore)

Risks
None identified at the moment.

People and Personnel

• Lead or Co-Lead of some parts of this SDI: @valeriabarra
• Potential Collaborators: @LenkaNovak , @juliasloan25 ?
• Reviewers: @charleskawczynski

Components

Examples of components (not an exhaustive list) include:

In src/Utilities.jl:

• Avoid abstract types in structs #224
• Fix broken docstrings #220
• Type-stability of functions #250
• Make swap_space not allocate. (Inline it and avoid parent). Make swap_space as not allocating and more GPU compatible #248

In Regridder.jl:

• Use relative module qualifiers for namespaces #222
• FT_dot is allocating and referencing to an undefined var. Can we remove this function? #252
• nan_to_zero should throw a @warn to warn the user that NaNs were detected #254
• Make comms_ctx a required argument (not optional) to avoid possible conflicts #256

Specifically, an example of how to avoid abstract types in structs in src/Utilities.jl is, instead of having

struct CoupledSimulation{FT, B, ...}
a::NamedTuple
...
end

we would have:

struct CoupledSimulation{FT <: Real, NT <: NamedTuple, B, ...}
a::NT
...
end

Also, it's useful to add helpers for extracting type information for types that are commonly needed in functions (e.g., the float type), for example:

float_type(::CoupledSimulation{FT}) where {FT} = FT
FT = float_type(coupler_sim)

Inputs

These inputs were kindly provided by @charleskawczynski and can be refined/modified further as we identify more points of improvement.

Results and deliverables

Better, type-safe, and more optimized interfaces in ClimaCoupler.
Update: So far the improvements have led to about 60% less allocations. See these CI perf jobs log.

Task Breakdown And Tentative Due Date

(A preliminary list of PRs and a preliminary timeline of PRs, milestones, and key results)

• Documentation improvements fixing broken docstrings (week 1)
• Improve on dependencies inclusions to avoid conflicted versions (week 1)
• Remove abstract types in structs (week 2)
• Improve type-stability of functions (week 3)
• Reduction of allocations and identification of missing upstream interfaces (week 3 and 4)

Proposed Delivery Date

Since I am not going to work on this full time and there can be collaborative components it is a bit hard for me to propose a specific delivery date, but I could say it could take 4+ weeks.

Some potential future improvements

• Something to keep in mind is the potential excessive requests when using @inline . This was observed in ClimaCore for instance (see this issue #1064). But it can still be beneficial for big functions or functions called multiple times.

• Everywhere we are using directly ClimaCore internals, such as parent it means that ClimaCore is not providing that interface that users need. Identify those missing interfaces and report them upstream. Avoiding the use of parent throughout the ClimaCoupler code (which not only reduces allocation, but makes the code GPU compatible) should be a separate SDI, requiring more work.

• QA and deliverables checked.

[LenkaNovak]

Looks good, Valeria. Your timelines seem good. It would be good to pick one driver (e.g., experiments/AMIP/moist_mpi_earth and the modules in src/) and keep track of how the performance improves with each PR (even if just manually using @Btime for now). I will be implementing better performance logs this month, so when we move onto the less obvious performance improvements it will be more automated.

[LenkaNovak]

To answer your questions: FT_dot can be removed, and swap_space should be designed (or substituted) so it does not allocate.

[LenkaNovak]

As for timelines 4-6 days in total (as you suggest) seems reasonable. This estimate accounts for tricky dependency clash issues we've been encountering during the interface SDI.

[valeriabarra]

Cumulatively, the work to reduce allocations and to improve type-safety in this SDI have reduced allocations of about 60% (as it can also be seen from this build log)