Re-design AxisTensors

[charleskawczynski]

Purpose

The purpose of this SDI is to improve the design of AxisTensors. There are several issues with the existing design:

1. A single "mega-type" (a type-heavy struct) is used to allow for many different types of axis tensors-- (some previous tests create 3486 different kinds, to be exact).
2. Due to the "mega-type" approach, we end up stuck having to generate many methods for performant runtime code, rather than using far fewer methods and relying on dispatch. One price we pay with that is significantly higher latency. #853 makes some ClimaAtmos examples go from 2 min to 7 min. But other, larger jobs, reduce in runtime by ~10%.
3. Documentation. We can't easily document this type because there are so many potential type parameters. We sort of mask this by creating many user-facing aliases, but those are not easily documentable, and users often only see this in source code, and not in stack traces, because stack traces will show typeof(::AxisTensor) which will not show the alias, but the concrete type of AxisTensor, creating a visual mismatch for people debugging.
4. There are many wrappers in Geometry which perform nested operations on specific types and simply promote types in order to function, at the cost of doing all operations in a dense manner. This results in a huge number of flops relative to how many may actually be needed. One key aspect of the re-design will be to offer functions that do operations with a minimal number of flops.

Link to any relevant PRs/issues.

Cost/benefits/risks

• The cost of this is likely in the ~1-2 months of work for 1-2 people, as the AxisTensors is a fundamental component, is pretty , and will likely require breaking changes.
• The benefit of fixing this is mostly mentioned in the purpose: better performance, latency, and documentation
• The risk of not fixing this is that it's a fundamental building block for ClimaCore users, and the compilation latency issue cannot be solved upstream, so it's imperative that this is redesigned

Producers

• I know that @simonbyrne has ideas on how to improve the design, I'd be happy to help, too.

Components

• AxisTensors

Inputs

There are many relations (e.g., conversions) that will need to be re-written, but most of this is related to the type design

Results and deliverables

A better design to AxisTensors

Task breakdown

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

• [ ] Decide how to decompose the 3486 AxisTensors into composeable parts that can result in a better design for conversions
• [ ] Implement and test ...

Reviewers

• @simonbyrne, @charleskawczynski, @akshaysridhar, @valeriabarra

[tapios]

I'd like to see a clear description of the minimum functionality that we need (e.g., transformation covariant -> contravariant and inverse, in the horizontal and in 3d). That is, put down a set of requirements that need to be fulfilled, and then discuss the code structure that will fulfill them.

From the description here, I have the impression that what we need is not so much an improved design of AxisTensors (how did we ever get to 3486 types?), but getting rid of it and replacing it by much simpler code. I'd be hesitant to invest resources in improving something just because we have sunken costs in it.

[charleskawczynski]

I'd like to see a clear description of the minimum functionality that we need (e.g., transformation covariant -> contravariant and inverse, in the horizontal and in 3d). That is, put down a set of requirements that need to be fulfilled, and then discuss the code structure that will fulfill them.

Good point. Since @simonbyrne has some ideas about this, let's wait for when he returns.

From the description here, I have the impression that what we need is not so much an improved design of AxisTensors (how did we ever get to 3486 types?), but getting rid of it and replacing it by much simpler code. I'd be hesitant to invest resources in improving something just because we have sunken costs in it.

In some ways yes. And it's not that we have 3486 types, it's that we can make 3486 (distinguishable) types. Some of these transformations are indistinguishable by changing coordinate systems. For example, transformations in a x-y plane are distinguishable from transformations in a y-z plane.

I think @simonbyrne mentioned that we may be able to continue using what we have and just forward generic transformations to more generic methods (to avoid defining many methods). This would fix the many method issue, however, it wouldn't really address the documentation issue.

On another note, timings in point 2 should be confirmed. I've since run other examples that didn't result in slowdown. It's possible that more inlining is needed.

[simonbyrne]

There are 3 main representations we make use of: covariant, contravariant, local (aka uvw). This means that there are 6 possible conversion directions.

What really causes the blow up is dealing with subvectors (e.g. Covariant12Vector, WVector, etc), and subspaces (e.g. shallow water equations operate in a "1,2" space, inertial gravity wave in a "1,3" space). Adding in all these combinations causes a large explosion in the number of methods we use.

Additionally, we also currently define transform and project: these are basically the same, but transform will throw an error if the conversion is not exact (i.e. if you drop a non-zero dimension).

@charleskawczynski do you have a list of the methods we generate? Can you post them in a gist so I can see what we're using? That might help narrow down the scope

[charleskawczynski]

@charleskawczynski do you have a list of the methods we generate? Can you post them in a gist so I can see what we're using? That might help narrow down the scope

Yes, here is the gist. It includes:

• function name symbol or string (string when it's a functor)
• number of arguments (the project function has 2 and 3 argument methods), I added this mostly for scraping convenience
• tuple of argument types

There are a lot, but as you've said, it still narrows down the scope

[charleskawczynski]

There are actually more methods that I missed, which call project in a nested fashion, those should probably be written to minimize loads/stores/flops, too.

[charleskawczynski]

For example:

transform(
    ax::ContravariantAxis,
    v::CovariantTensor,
    local_geometry::LocalGeometry,
) = project(
    ax,
    local_geometry.∂ξ∂x *
    local_geometry.∂ξ∂x' *
    project(dual(axes(local_geometry.∂ξ∂x, 1)), v),
)

This is probably very slow.