Use LazyBroadcast for tendencies

[charleskawczynski]

A step towards #3594.

This PR adds the use of LazyBroadcast.jl to fuse the sponge operations. This is hopefully the beginning of a long road of fusing many operations, and moving the codebase into a more functional, and more unit-tested codebase.

[Sbozzolo]

I'd prefer we didn't introduce more technical concepts in science packages without proper documentation and written guidence on what they do and how to use them, especially if this becomes more and more common.

The goal should be that people developing this package should be able to understand and use this feature without input from experts.

Examples of questions:

• Can I just call @lazy in front of any broadcasted expression and everything will work?
• What benefits should I see?
• How does this enable kernel fusion? What's the difference with ClimaCore.@fused?

One way to check that you provided enough guidance is with user stories and asking people if they can use and understand this feature without your input.

Example of a user story:

I am a reserach scientist and I am debugging a piece of ClimaAtmos. I notice that the return objects of some the functions I call are some Broadcast.Brocasted objects that I don't know how to manipulate. This is not explained anywhere within the ClimaAtmos repo. I dig deeper and see that it might be connected to @lazy. I come to this conclusion because I am familiar with the other operations and @lazy is the only thing I have never seen. Following the imports, I end up in the LazyBroadcast.jl package. The only documentation in that pacakge is a readme that is very technical for me (e.g., I've never heard of "materializing"; it's not even explained in the Julia manual). The link to the Julia manual doesn't really help me: it mosty talks about broadcasting as in using the dot operator. The readme doesn't even mention the @lazy macro. It's unclear to me how the three points that are mentioned in the readme relate to my problem. The only examples provided are tests and don't explain what's happening. At this point, I need to ask for help, if I haven't done it already.

It's up to you to where you add documentation and guidance. For example, you might decide that LazyBrodcast is intended for a technical audience. That would be fine, but then you have to provide documentation within ClimaAtmos (or ClimaCore and links to the relevenat documentation in ClimaAtmos).

I think that this is very important. Many Clima packages, including ClimaAtmos and ClimaCore, require a lot of tribal knowledge to be used and developed. This is a big barrier to adoption, developlment, and suistanability in general. I think we should really strive to make sure that all our new contributions are propertly documented, while also catching up with the missing documentation.

[charleskawczynski]

I'd prefer we didn't introduce more technical concepts in science packages without proper documentation and written guidence on what they do and how to use them, especially if this becomes more and more common.

We can add documentation. We can also add better unit tests with this approach, which I've also updated.

The goal should be that people developing this package should be able to understand and use this feature without input from experts.

I agree, but users will need to gain an idea of how to use things, and we can document how things work

Examples of questions:

• Can I just call @lazy in front of any broadcasted expression and everything will work?

The new functional-based tendencies are dispatched in a binary way: one returns a broadcasted object, via @lazy, and the other returns a single point value that should be valid when the tendency is not to be added. Alternatively, we could use a design where we return a Null object, and filter the broadcast expressions. I'm open to either design.

• What benefits should I see?

This has several benefits:

• We can actually write unit tests, and more granular unit tests
• We can compose ClimaCore operators, and there is no other way (that I know of)
• We can fuse operations. I've spoken with several experts in the field about improving performance, and it is clear that large-scale fusion of kernels has the largest potential performance benefit. We don't know how much because it's difficult to measure, but the next optimization is at most 35%, based on Amdahl's law. I outlined the performance items in the performance roadmap, if you'd like to read more.
• We might even be able to do lazy-caching for expensive quantities (I'm not yet sure what ingredients are needed for the compiler to perform common-sub-expression elimination, which may be useful for this approach).

• How does this enable kernel fusion? What's the difference with ClimaCore.@fused?

This is "manual" fusion. We're manually collecting broadcasted objects, and putting them into a single @. call (which calls materialize, which performs the execution of the expression). I'm starting to think there is one big benefit of using LazyBroadcast over MultiBroadcastFusion: we can leverage dispatch and composition better with LazyBroadcast (again, see this issue). Also, LazyBroadcast.jl's implementation is extremely simple, so it's easy for users to dig into the details if they want to.

One way to check that you provided enough guidance is with user stories and asking people if they can use and understand this feature without your input.

Example of a user story:

I am a reserach scientist and I am debugging a piece of ClimaAtmos. I notice that the return objects of some the functions I call are some Broadcast.Brocasted objects that I don't know how to manipulate. This is not explained anywhere within the ClimaAtmos repo. I dig deeper and see that it might be connected to @lazy. I come to this conclusion because I am familiar with the other operations and @lazy is the only thing I have never seen. Following the imports, I end up in the LazyBroadcast.jl package. The only documentation in that pacakge is a readme that is very technical for me (e.g., I've never heard of "materializing"; it's not even explained in the Julia manual). The link to the Julia manual doesn't really help me: it mosty talks about broadcasting as in using the dot operator. The readme doesn't even mention the @lazy macro. It's unclear to me how the three points that are mentioned in the readme relate to my problem. The only examples provided are tests and don't explain what's happening. At this point, I need to ask for help, if I haven't done it already.

I hear you. And I need help with this, as it's a big challenge. Please hear me: we should follow the most promising software design, and provide users with excellent documentation. I don't want to avoid the best path forward because we don't yet have documentation. I need your help writing it.

It's up to you to where you add documentation and guidance. For example, you might decide that LazyBrodcast is intended for a technical audience. That would be fine, but then you have to provide documentation within ClimaAtmos (or ClimaCore and links to the relevenat documentation in ClimaAtmos).

Yes. Right now, I'm doing a sort-of tracer-bullet PR, to see what this might look like.

I think that this is very important. Many Clima packages, including ClimaAtmos and ClimaCore, require a lot of tribal knowledge to be used and developed. This is a big barrier to adoption, developlment, and suistanability in general. I think we should really strive to make sure that all our new contributions are propertly documented, while also catching up with the missing documentation.

Agreed.

[Sbozzolo]

If this is the path we are taking, let's just take the time to document this now. This is part of the delivery of the feature. The questions I listed in the previous message were questions that users ought to be able to find the answer for somewhere, so your answers can be the starting point for this. Even if the final form of this will look different, the documentation written can be still adapted and reused.
If this is an experiment, let's not merge it into main until we figure things out.

We will never recover from the missing documentation if we keep adding more missing documentation, especially when documentation is never a priority.

[charleskawczynski]

Agreed, I think we should start by beefing up documentation at the bottom, and work our way up as things concretize. I just opened CliMA/LazyBroadcast.jl#16.

[charleskawczynski]

To avoid this PR from getting stale, I may add LazyBroadcast as a dependency in a separate PR, so that I don't need to continually rebase the manifest changes.

[charleskawczynski]

I'm thinking that we can start by using LazyBroadcast, and not changing the order of tendency computations. This has two benefits:

• We can microbenchmark fused vs unfused versions, in a separate and dedicated effort. I think that this is probably best since heuristics like register pressure can result in performance degradation-- it's not trivial to see that fusing two expressions will be faster or not.
• We still benefit from being able to write unit tests.

I'll update the PR.

[charleskawczynski]

I think we should address some edge cases first. In particular, I don't want any performance penalty in the case that the functional tendencies are absent.