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 implicit solver compilation issues #1506

Merged
merged 1 commit into from
Dec 12, 2023
Merged

Conversation

dennisYatunin
Copy link
Member

@dennisYatunin dennisYatunin commented Oct 19, 2023

This PR fixes a number of compilation issues in the implicit solver infrastructure, which have been causing inference failures and allocations in the ClimaAtmos-based unit tests (and in ClimaAtmos.jl itself).

Content

  • Replace all of the recursively unrolled functions used in the MatrixFields module with generatively unrolled functions, which gets rid of the stack overflow in type inference compiler errors. Some of these functions are taken from Unrolled.jl, while others build upon the functions defined there. Move all unrolled functions into Utilities.UnrolledFunctions and add some documentation. Update the Manifest.toml files to add Unrolled.jl.
  • Rewrite FieldNameDict so that there are no constant Union types being passed around, which gets rid of the allocations in the prognostic EDMF unit test. This involves merging FieldMatrix and FieldMatrixBroadcasted into a single type, and also merging FieldVectorView and FieldVectorViewBroadcasted. Add the is_lazy function to distinguish between materialized and un-materialized FieldNameDicts. Update all relevant documentation.
  • Modify the definitions of norm_sqr(::FieldVector) and array_type(::FieldVector) to make them type-stable for nested FieldVectors, which fixes the remaining type instabilities in a few unit tests. Nested FieldVectors can be generated while debugging a StationaryIterativeSolve that is part of a larger FieldMatrixSolverAlgorithm.
  • Change some of the FieldNameSet method definitions to avoid specializations of the form where {T}, as was suggested by @charleskawczynski. This does not appear to affect compilation, but it does make the code a bit easier to read.

Invalidations

Adding Unrolled.jl changes 4 rows in the invalidations table and brings the total number of invalidations up from 5529 to 5554. On the main branch, these 4 rows are

┌─────────────────────────────────────────────────┬───────────────┬───────────────┬─────────────────┐
│ <file name>:<line number>                       │ Function Name │ Invalidations │ Invalidations % │
│                                                 │               │               │     (xᵢ/∑x)     │
├─────────────────────────────────────────────────┼───────────────┼───────────────┼─────────────────┤
│ ColorTypes/1dGw6/src/operations.jl:24           │    isequal    │      50       │        1        │
│ InlineStrings/rlLZO/src/InlineStrings.jl:606    │     string    │      41       │        0        │
│ StaticArrays/yXGNL/src/abstractarray.jl:145     │    similar    │      35       │        0        │
│ BlockArrays/4FBcw/src/abstractblockarray.jl:221 │    getindex   │       8       │        0        │
└─────────────────────────────────────────────────┴───────────────┴───────────────┴─────────────────┘

On this branch, these 4 rows are

┌─────────────────────────────────────────────────┬───────────────┬───────────────┬─────────────────┐
│ <file name>:<line number>                       │ Function Name │ Invalidations │ Invalidations % │
│                                                 │               │               │     (xᵢ/∑x)     │
├─────────────────────────────────────────────────┼───────────────┼───────────────┼─────────────────┤
│ InlineStrings/rlLZO/src/InlineStrings.jl:606    │     string    │      77       │        1        │
│ StaticArrays/yXGNL/src/abstractarray.jl:145     │    similar    │      33       │        0        │
│ Unrolled/KtWta/src/range.jl:24                  │    getindex   │      33       │        0        │
│ ColorTypes/1dGw6/src/operations.jl:24           │    isequal    │      27       │        0        │
└─────────────────────────────────────────────────┴───────────────┴───────────────┴─────────────────┘

Since the same function names appear in both tables, the increase in invalidations is probably not related to the other changes in this PR.

@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Oct 19, 2023
@bors
Copy link
Contributor

bors bot commented Oct 19, 2023

try

Build failed:

@dennisYatunin dennisYatunin force-pushed the dy/unrolled_functions branch 2 times, most recently from f907215 to ae46d25 Compare October 20, 2023 23:48
@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Oct 20, 2023
@bors
Copy link
Contributor

bors bot commented Oct 21, 2023

try

Build failed:

@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Oct 21, 2023
@bors
Copy link
Contributor

bors bot commented Oct 21, 2023

try

Build failed:

@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Oct 27, 2023
@bors
Copy link
Contributor

bors bot commented Oct 27, 2023

try

Build succeeded!

The publicly hosted instance of bors-ng is deprecated and will go away soon.

If you want to self-host your own instance, instructions are here.
For more help, visit the forum.

If you want to switch to GitHub's built-in merge queue, visit their help page.

@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Oct 27, 2023
@bors
Copy link
Contributor

bors bot commented Oct 27, 2023

try

Build succeeded!

The publicly hosted instance of bors-ng is deprecated and will go away soon.

If you want to self-host your own instance, instructions are here.
For more help, visit the forum.

If you want to switch to GitHub's built-in merge queue, visit their help page.

@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Oct 28, 2023
@bors
Copy link
Contributor

bors bot commented Oct 28, 2023

try

Build succeeded!

The publicly hosted instance of bors-ng is deprecated and will go away soon.

If you want to self-host your own instance, instructions are here.
For more help, visit the forum.

If you want to switch to GitHub's built-in merge queue, visit their help page.

@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Nov 1, 2023
Copy link
Contributor

bors bot commented Nov 1, 2023

try

Build failed:

@dennisYatunin
Copy link
Member Author

bors try

bors bot added a commit that referenced this pull request Nov 1, 2023
Copy link
Contributor

bors bot commented Nov 1, 2023

try

Build failed:

@dennisYatunin dennisYatunin force-pushed the dy/unrolled_functions branch 14 times, most recently from 901b05e to 62fb219 Compare December 6, 2023 01:27
@dennisYatunin dennisYatunin changed the title Switch to using Unrolled.jl Switch to using Unrolled.jl and fix other FieldName compilation issues Dec 6, 2023
@dennisYatunin dennisYatunin changed the title Switch to using Unrolled.jl and fix other FieldName compilation issues Fix Implicit Solver Compilation Issues Dec 6, 2023
@dennisYatunin dennisYatunin changed the title Fix Implicit Solver Compilation Issues Fix implicit solver compilation issues Dec 6, 2023
Copy link
Member

@charleskawczynski charleskawczynski left a comment

Choose a reason for hiding this comment

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

We may want to revisit this at some point to see if we can eliminate the unroll functions, but overall looks fine to me.

@charleskawczynski
Copy link
Member

Also, we can just ignore the invalidations issue (we can revisit this when we upgrade to julia 1.10)

@dennisYatunin dennisYatunin merged commit c593d09 into main Dec 12, 2023
15 of 16 checks passed
@dennisYatunin dennisYatunin deleted the dy/unrolled_functions branch December 12, 2023 22:28
@dennisYatunin dennisYatunin mentioned this pull request Dec 12, 2023
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.

3 participants