Improve tracer performance in jacobian update #4070
+25
−5
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
haakon-e
reviewed
Oct 24, 2025
Comment on lines
572
to
578
| @. p.scratch.ᶠtemp_scalar = ClimaAtmos.ᶠinterp(ᶜρ * ᶜJ) / ᶠJ | ||
| @. p.scratch.ᶜtemp_scalar_5 = -(ᶜwₚ) / ᶜρ | ||
| @. ∂ᶜρχₚ_err_∂ᶜρχₚ = | ||
| dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ | ||
| DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ | ||
| DiagonalMatrixRow(p.scratch.ᶠtemp_scalar) ⋅ | ||
| ᶠright_bias_matrix() ⋅ | ||
| DiagonalMatrixRow(-Geometry.WVector(ᶜwₚ) / ᶜρ) - (I,) | ||
| DiagonalMatrixRow(Geometry.WVector(p.scratch.ᶜtemp_scalar_5)) - (I,) |
Member
There was a problem hiding this comment.
Suggested change
| @. p.scratch.ᶠtemp_scalar = ClimaAtmos.ᶠinterp(ᶜρ * ᶜJ) / ᶠJ | |
| @. p.scratch.ᶜtemp_scalar_5 = -(ᶜwₚ) / ᶜρ | |
| @. ∂ᶜρχₚ_err_∂ᶜρχₚ = | |
| dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ | |
| DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ | |
| DiagonalMatrixRow(p.scratch.ᶠtemp_scalar) ⋅ | |
| ᶠright_bias_matrix() ⋅ | |
| DiagonalMatrixRow(-Geometry.WVector(ᶜwₚ) / ᶜρ) - (I,) | |
| DiagonalMatrixRow(Geometry.WVector(p.scratch.ᶜtemp_scalar_5)) - (I,) | |
| ᶠρ = @. p.scratch.ᶠtemp_scalar = ᶠinterp(ᶜρ * ᶜJ) / ᶠJ | |
| ᶜminus_wₚ_div_ρ = @. p.scratch.ᶜtemp_scalar_5 = -(ᶜwₚ) / ᶜρ | |
| @. ∂ᶜρχₚ_err_∂ᶜρχₚ = | |
| dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ | |
| DiagonalMatrixRow(ᶠρ) ⋅ | |
| ᶠright_bias_matrix() ⋅ | |
| DiagonalMatrixRow(Geometry.WVector(ᶜminus_wₚ_div_ρ)) - (I,) |
I think giving the temp fields reasonable names improves readability.
Member
Author
There was a problem hiding this comment.
I updated the intermediate computations to this:
@. p.scratch.ᶜbidiagonal_adjoint_matrix_c3 = -(ClimaAtmos.ᶜprecipdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ClimaAtmos.ᶠinterp(ᶜρ * ᶜJ) / ᶠJ)
@. p.scratch.ᶠband_matrix_wvec = ClimaAtmos.ᶠright_bias_matrix() ⋅ DiagonalMatrixRow(ClimaCore.Geometry.WVector(-(ᶜwₚ) / ᶜρ))Do you have a suggestion for names that make sense here?
imreddyTeja
commented
Oct 24, 2025
Split up a high register pressure broadcast. It appears that just the broadcast is nearly seven times as fast after the change. SYPD for 1m with 30 he improves by ~11%
imreddyTeja
force-pushed
the
tr/tracer-jac-brap
branch
from
106ef3d to
43410b9
Compare
szy21
approved these changes
Nov 7, 2025
Member
szy21
left a comment
szy21
left a comment
There was a problem hiding this comment.
I don't think there would be a good name for these temporary fields...
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Purpose
Split up a high register pressure broadcast.
closes #4064
To-do
Content
It appears that just the broadcast is nearly seven times as fast after the change.
Sypd of numerics_sphere_he30ze63.yml and longrun_aquaplanet_allsky_1M.yml:
before: 0.429
after: 0.476
~11% speedup
BEFORE:
AFTER:
