Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix integrals for broadcasts and update docstrings #1947

Merged
merged 2 commits into from
Aug 22, 2024
Merged

Conversation

dennisYatunin
Copy link
Member

This PR fixes the issue identified in CliMA/ClimaAtmos.jl#3254, allowing the definite and indefinite integral functions to accept AbstractBroadcasteds in addition to Fields. It also makes the function argument names more consistent, updates the docstrings, and modifies the unit tests to use AbstractBroadcasteds instead of Fields.

  • Code follows the style guidelines OR N/A.
  • Unit tests are included OR N/A.
  • Code is exercised in an integration test OR N/A.
  • Documentation has been added/updated OR N/A.

@Sbozzolo
Copy link
Member

I argue that we should keep all the usage of anything related to Broadcast internal to ClimaCore. We should not expect people to know about what Base.Broadcast is and using it dramatically increases the entry barrier to using and extending our codes.

If the issue is the extra allocation, we can use temporary Fields to reduce/remove that.

@dennisYatunin
Copy link
Member Author

@Sbozzolo, the goal here is not necessarily reducing allocations, but rather kernel fusion. We don't need to launch an additional kernel to populate the integrand before computing an integral, and the current design lets us avoid this by passing unmaterialized Broadcasted objects instead of Fields. I agree that manipulation of Broadcasted objects should be done internally, and this design is just an intermediate step toward making ColumnReduce and ColumnAccumulate into ClimaCore operators instead of functions. Once that's done in a future release, we won't need to deal with Broadcasted objects for integrals anymore.

@dennisYatunin dennisYatunin force-pushed the dy/integrals branch 2 times, most recently from 3d8f497 to 4e157ab Compare August 21, 2024 21:21
@Sbozzolo
Copy link
Member

@Sbozzolo, the goal here is not necessarily reducing allocations, but rather kernel fusion. We don't need to launch an additional kernel to populate the integrand before computing an integral, and the current design lets us avoid this by passing unmaterialized Broadcasted objects instead of Fields. I agree that manipulation of Broadcasted objects should be done internally, and this design is just an intermediate step toward making ColumnReduce and ColumnAccumulate into ClimaCore operators instead of functions. Once that's done in a future release, we won't need to deal with Broadcasted objects for integrals anymore.

If the issue is kernel fusion (or performance), how important is to enable this capability now? Can't we wait after you added the functions you want to add to ClimaCore?

I am worried that we are opening the door to code that only one or two people can work with. Broadcasted expressions have started appearing in packages and I don't want it to be another parent. Squeezing a tiny tiny increase of performance is not worth doing if the cost if having code that is hard to manipulate and increases the entry barrier.

Let's cross the bridge of kernel fusion when it starts to matter and focus on delivering code that people can understand and use.

@dennisYatunin
Copy link
Member Author

It is certainly possible to avoid making use of this capability now---you can just materialize the integrand and pass in the resulting Field. Perhaps I'll add this bugfix for now, but people can use the standard caching pattern instead of kernel fusion until the operator interface is fully fleshed out.

More generally, though, I think we will need to get used to seeing @lazy_broadcasted expressions (or just raw Broadcasted objects) in other packages. Unlike parent, which bypasses all of ClimaCore's Field infrastructure to return an Array, the broadcasting infrastructure is an intrinsic part of ClimaCore, and it is the primary mechanism for specifying lazy operations in Julia. Lazy operations allow us to build highly generalizable and modular interfaces without sacrificing performance (for example, see the implicit solver code in MatrixFields), and we will probably want to make greater use of them moving forward. In particular, the cache, tendencies, and diagnostics in ClimaAtmos could all be made significantly faster (probably several times faster, not a "tiny tiny increase") using such a design.

@dennisYatunin
Copy link
Member Author

After an offline discussion, we have decided that this can be merged in, but that we shouldn't use the kernel fusion capability in ClimaAtmos until it is better documented on the ClimaCore side. I will be adding documentation and converting the integral functions into operators in future PRs.

@dennisYatunin dennisYatunin merged commit 09d57ca into main Aug 22, 2024
21 of 22 checks passed
@dennisYatunin dennisYatunin deleted the dy/integrals branch August 22, 2024 18:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants