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

Specialize on diagonal fieldvector broadcasts #1615

Merged
merged 1 commit into from
Feb 21, 2024

Conversation

charleskawczynski
Copy link
Member

@charleskawczynski charleskawczynski commented Feb 20, 2024

This PR specializes on "diagonal" broadcast expressions with fieldvectors. Specifically, if a fieldvector broadcast expression (one with a FieldVectorStyle) has one and only one fieldvector type, then we don't check the broadcast axes (the combined axes is the same). The benefit of skipping this code path is that checking the axes is type-unstable (JuliaArrays/BlockArrays.jl#310). From what I can tell, there is not a simple path to fixing JuliaArrays/BlockArrays.jl#310 because:

  • We want to use Tuples, so that this is stack-allocated
  • Combining axes for Tuples is type unstable because the result depends on the values in the tuples.

So, I think that the simplest solution is to specialize when we know that the axes will not change.

This PR adds a test to show that is_diagonal_bc will return true/false when expected at compile time (@code_typed is_diagonal_bc(bc) returns true or false), and tests that (what was breaking) fieldvector broadcast expressions with multiple fields does not allocate.

Closes #1465.

One other benefit of this solution is that we can at some point reuse is_diagonal_bc to specialize on these broadcast expressions in copyto!, so that we perform a single cuda kernel launch (as opposed to however many fields are in the fieldvector).

The only thing remaining is that I'd like to run ClimaAtmos CI on this branch to make sure that this doesn't blow up compile-times. Update: good news, it looks fine xref: CliMA/ClimaAtmos.jl#2700 (comment)

@Sbozzolo
Copy link
Member

I think @dennisYatunin is in a better position to review this, so I'll leave this to him.

Copy link
Member

@dennisYatunin dennisYatunin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall, I think this looks good. You could simplify is_diagonal_bc by rewriting it in terms of Utilities.UnrolledFunctions.unrolled_all, but what you have is fine. There is one change I'd like to see for catching mismatched spaces errors; see my comment.

src/Fields/fieldvector.jl Show resolved Hide resolved
@charleskawczynski
Copy link
Member Author

@dennisYatunin, looks like we can't error: https://buildkite.com/clima/climacore-ci/builds/3150#018dc877-7cd8-4de7-928d-de25532c360d. I'm going to merge this now, and we can follow up later if there's some improvement we can make.

@charleskawczynski charleskawczynski merged commit da32a43 into main Feb 21, 2024
16 checks passed
@charleskawczynski charleskawczynski deleted the ck/fv_specialize branch February 21, 2024 00:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

FieldVector broadcast inference failure
3 participants