Skip to content

Add missing @propagate_inbounds and optimized form of setindex#2488

Merged
dennisYatunin merged 1 commit intomainfrom
dy/inbounds_fixes
Apr 10, 2026
Merged

Add missing @propagate_inbounds and optimized form of setindex#2488
dennisYatunin merged 1 commit intomainfrom
dy/inbounds_fixes

Conversation

@dennisYatunin
Copy link
Copy Markdown
Member

Purpose

This PR fixes two issues I discovered while working on #2417:

  • When recursive unrolling was replaced with unrolled_map in Group/unrolling #2413, several @propagate_inbounds annotations where accidentally dropped, since neither unrolled_map nor the anonymous functions passed to it were annotated with @propagate_inbounds. This has led to unnecessary bounds checking during certain broadcast operations, which can be seen in the CI flame graphs for ClimaCore and ClimaAtmos.
  • In Add support for nonuniform data structures #2464, I implemented several functions in terms of Base.setindex, including the performance-critical struct_field_view and struct_indices. It turns out that Base.setindex does not always inline the index value, leading to unnecessary index comparisons and bounds checks during runtime. These were also visible in the CI flame graphs, though they were not as significant as the first issue.

Content

Three new functions are added to a sub-module of Utilities: unrolled_setindex, unrolled_insert, and unrolled_map_with_inbounds. These will eventually be generalized and moved to UnrolledUtilities.jl, but for now they can be explicitly imported from ClimaCore.Utilities.Unrolled. I expect we'll need to define other similar functions in the process of refactoring DataLayouts/Operators, after which we can update the UnrolledUtilities API to include them.


  • 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.

Copy link
Copy Markdown
Member

@imreddyTeja imreddyTeja left a comment

Choose a reason for hiding this comment

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

Do you have evidence that these changes make inbounds propagate as expected? Same question for the inlining of the index

Comment thread src/interface.jl
Comment thread src/Utilities/Utilities.jl
@dennisYatunin
Copy link
Copy Markdown
Member Author

@imreddyTeja Here are screenshots of the "Analysis: Flamegraph Profile" job from ClimaCore's CI that illustrate the difference; the first screenshot is from the latest main branch, while the second screenshot is from this branch:
Screenshot 2026-04-09 at 7 23 27 PM
Screenshot 2026-04-09 at 7 23 57 PM

@dennisYatunin dennisYatunin enabled auto-merge April 10, 2026 03:21
@dennisYatunin dennisYatunin merged commit db7c909 into main Apr 10, 2026
37 checks passed
@dennisYatunin dennisYatunin deleted the dy/inbounds_fixes branch April 10, 2026 03:23
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