Skip to content

Stochastic Physics CPU and GPU Optimizations - NGARCH #65

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.

Sign up for GitHub

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

Jump to bottom
Open
jasonjunweilyu wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Stochastic Physics CPU and GPU Optimizations - NGARCH #65

jasonjunweilyu wants to merge 9 commits into from

Conversation

@jasonjunweilyu
Copy link

@jasonjunweilyu jasonjunweilyu commented Dec 17, 2025 •
edited by MetBenjaminWent
Loading

PR Summary

Sci/Tech Reviewer: @mo-alistairp
Code Reviewer: @MetBenjaminWent
CC: @hiker

This completes the work of NGARCH stochastic physics optimization for CPU and GPU by migrating from fcm to git to be merged to main. Details of the completed work are documented in ticket:543.

Code Quality Checklist

(Some checks are automatically carried out via the CI pipeline)

  • I have performed a self-review of my own code
  • My code follows the project's style guidelines
    style guidelines
  • Comments have been included that aid undertanding and enhance the
    readability of the code
  • My changes generate no new warnings

Testing

  • I have tested this change locally, using the LFRic Apps rose-stem suite
  • If any tests fail (rose-stem or CI) the reason is understood and
    acceptable (eg. kgo changes)
  • I have added tests to cover new functionality as appropriate (eg. system
    tests, unit tests, etc.)
  • Any new tests have been assigned an appropriate amount of compute resource
    and have tests been allocated to an appropriate testing group (i.e. the
    developer tests are for jobs which use a small amount of compute resource
    and complete in a matter of minutes)

Since only the rose-stem climate config is now using stochastic physics, I have run the lfric_atm_clim_gal9_nomg-C12_gadi_intel_fast-debug-64bit available on NCI GADI. The build and run jobs finished successfully. For KGO verification, since KGOs on GADI have not been updated, I have compared the KGO with the one produced by running the trunk version I forked from and confirmed that they are identical.

trac.log

Security Considerations

  • I have reviewed my changes for potential security issues
  • Sensitive data is properly handled (if applicable)
  • Authentication and authorisation are properly implemented (if applicable)

Performance Impact

  • Performance of the code has been considered and, if applicable, suitable
    performance measurements have been conducted

AI Assistance and Attribution

  • Some of the content of this change has been produced with the assistance
    of Generative AI tool name (e.g., Met Office Github Copilot Enterprise,
    Github Copilot Personal, ChatGPT GPT-4, etc) and I have followed the
    Simulation Systems AI policy
    (including attribution labels)

Documentation

  • Where appropriate I have updated documentation related to this change and
    confirmed that it builds correctly

PSyclone Approval

  • If you have edited any psyclone related code (eg. PsyKAl-lite, Kernal
    inteface, optimisation scripts, LFRic data structure code) then please
    contact the
    tooscollabdevteam@metoffice.gov.uk

Sci/Tech Review

  • I understand this area of code and the changes being added
  • The proposed changes correspond to the pull request description
  • Documentation is sufficient (do documentation papers need updating)
  • Sufficient testing has been completed

Please alert the code reviewer via a tag when you have approved the SR

Code Review

  • All dependencies have been resolved
  • Related Issues have been properly linked and addressed
  • CLA compliance has been confirmed
  • Code quality standards have been met
  • Tests are adequate and have passed
  • Documentation is complete and accurate
  • Security considerations have been addressed
  • Performance impact is acceptable

@github-actions github-actions bot added the cla-required The CLA has not yet been signed by the author of this PR - added by GA label Dec 17, 2025
@jasonjunweilyu jasonjunweilyu mentioned this pull request Dec 17, 2025
Open
@github-actions github-actions bot added cla-signed This contributor has signed the CLA. and removed cla-required The CLA has not yet been signed by the author of this PR - added by GA labels Dec 17, 2025
@mo-alistairp mo-alistairp self-requested a review December 17, 2025 11:46
@mo-alistairp
Copy link
Contributor

mo-alistairp commented Dec 17, 2025

Hi Jason, thanks for moving this over to GitHub.

I've left a few things for you to address before I can approve the SciTech review. These are all style fixes -

  • the python transformation files are all failing the flake8 style checker. Could you run this and make those changes?
  • there are a few style changes in the Fortran files that I'll leave as inline comments

mo-alistairp
mo-alistairp requested changes Dec 17, 2025
View reviewed changes
Copy link
Contributor

@mo-alistairp mo-alistairp left a comment

Choose a reason for hiding this comment

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

The changes functionally look good. I've left inline comments requesting changes - they are all style changes (mostly to do with spacing and alignment). I also will need you to fix all the pep8 failures in the Python scripts, which can be seen by running the flake8 command.

@jasonjunweilyu jasonjunweilyu force-pushed the 64_port_stochastic_physics_gpu branch 2 times, most recently from 328a2d9 to 8cefbfb Compare December 18, 2025 06:25
@jasonjunweilyu
Copy link
Author

jasonjunweilyu commented Dec 18, 2025

Hi @mo-alistairp , thanks for the comments of the styling issues. I have addresssed them and so re-request your review now.

mo-alistairp
mo-alistairp approved these changes Jan 7, 2026
View reviewed changes
Copy link
Contributor

@mo-alistairp mo-alistairp left a comment

Choose a reason for hiding this comment

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

Thank you for making those changes it is indeed passing lint checks now. I'm happy with these changes so SciTech review: PASS. I'll pass this onto code review now with @MetBenjaminWent

@yaswant yaswant added this to the Spring 2026 milestone Jan 12, 2026
MetBenjaminWent
MetBenjaminWent approved these changes Jan 19, 2026
View reviewed changes
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent left a comment

Choose a reason for hiding this comment

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

Developer tests all pass, and a local check shows that core CPU OMP functionality for EXs remains unchanged, but also has GPU functionality for these 3x scripts.

Whilst I would have liked to see these scripts tidied up and broken into functions, I recognise that it would block this ticket.

I'm approving this ticket on technical groups, but noting a technical debt included in these scripts.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py
const = LFRicConstants()
cpu_parallel = OMPParallelTrans()

if OFFLOAD_DIRECTIVES == "omp":
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

This seems like a good way to ensure this is set, however, could these checks be nested inside a function to reduce the length of this script?
For PSYKAL LFRic API, these would be added here:
https://github.com/MetOffice/lfric_core/blob/main/infrastructure/build/psyclone/psyclone_tools.py

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, I can add this in a new PR if needed.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py

# Make setval_* compute redundantly to the level 1 halo if it
# is in its own loop
for loop in subroutine.loops():
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

Like above, could this functionality be nested inside a function to reduce the length of this script?
For PSYKAL LFRic API, these would be added here:
https://github.com/MetOffice/lfric_core/blob/main/infrastructure/build/psyclone/psyclone_tools.py

Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 19, 2026

Choose a reason for hiding this comment

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

This is redundant_computation__setval ?

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, I can add this in a new PR if needed. It is indeed doing the same thing as redundant_computation__setval.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py

# Colour loops over cells unless they are on discontinuous spaces
# (alternatively we could annotate the kernels with atomics)
for loop in subroutine.loops():
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

Like above, could this functionality be nested inside a function to reduce the length of this script?
For PSYKAL LFRic API, these would be added here:
https://github.com/MetOffice/lfric_core/blob/main/infrastructure/build/psyclone/psyclone_tools.py

Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 19, 2026

Choose a reason for hiding this comment

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

This is colour_loops ?

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, it is colour_loops essentially.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py

# Mark kernels inside the loops over cells as GPU-enabled
# and inline them.
for loop in subroutine.loops():
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

Like above, could this functionality be nested inside a function to reduce the length of this script?
For PSYKAL LFRic API, these would be added here:
https://github.com/MetOffice/lfric_core/blob/main/infrastructure/build/psyclone/psyclone_tools.py

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, I can add this in a new PR if needed.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py
)
print(f"GPU-annotated kernel '{kern.name}'")

try:
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

I'm not sure how much I trust nested try statements, but given the dependency between annotating the kernel, then in-lining it, this doesn't seem unreasonable.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py
)
for name in kernel_names
):
try:
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

However here, could the if's and the try wrapper around the choices be reversed?

Have the ifs for the options, and then trys for each option?
It just seems a bit safer.

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, I can verify this in the new PR.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py
# arrays and convert to code unsupported intrinsics.

# Add GPU offloading to loops unless they are over colours or are null.
for loop in subroutine.walk(Loop):
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

Like above, could this functionality be nested inside a function to reduce the length of this script?
For PSYKAL LFRic API, these would be added here:
https://github.com/MetOffice/lfric_core/blob/main/infrastructure/build/psyclone/psyclone_tools.py

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, I can add this in a new PR if needed.

...lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py

# Apply OpenMP thread parallelism for any kernels we've not been able
# to offload to GPU.
for loop in subroutine.walk(Loop):
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

Like above, could this functionality be nested inside a function to reduce the length of this script?
For PSYKAL LFRic API, these would be added here:
https://github.com/MetOffice/lfric_core/blob/main/infrastructure/build/psyclone/psyclone_tools.py

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, I can add this in a new PR if needed.

applications/lfric_atm/optimisation/meto-ex1a/psykal/algorithm/skeb_main_alg_mod.py
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

Comments from applications/lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py are also applicable here as both scripts appear very similar, indicating duplicated code?

Where there are minor differences, these could be function options?

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, this is correct.

applications/lfric_atm/optimisation/meto-ex1a/psykal/algorithm/spt_main_alg_mod.py
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent Jan 13, 2026

Choose a reason for hiding this comment

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

Comments from applications/lfric_atm/optimisation/meto-ex1a/psykal/algorithm/runtime_constants/physics_constants_mod.py are also applicable here as both scripts appear very similar, indicating duplicated code?

Where there are minor differences, these could be function options?

Copy link
Author

@jasonjunweilyu jasonjunweilyu Jan 20, 2026

Choose a reason for hiding this comment

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

Yes, this is correct.

@MetBenjaminWent
Copy link
Contributor

MetBenjaminWent commented Jan 19, 2026 •
edited
Loading

Hi @jasonjunweilyu would you be able to fully reverify all the commits for us so we can commit this? Thanks!

https://metoffice.github.io/simulation-systems/WorkingPractices/gh_authorisation.html
One can retrospectively change commits in a repository, if required, with
git rebase -i HEAD~<N>
where is the number of commits back to change and setting each commit to edit or reword (see git documentation for more details).

image

@jasonjunweilyu jasonjunweilyu force-pushed the 64_port_stochastic_physics_gpu branch from c1dd404 to 70f44a1 Compare January 20, 2026 07:19
@jasonjunweilyu
Copy link
Author

jasonjunweilyu commented Jan 20, 2026 •
edited
Loading

Hi @MetBenjaminWent , thanks a lot for your review. I have reverified my first two commits with git rebase -i --rebase-merges HEAD~7 and git commit --amend --no-edit -S.
However, this seems to make your two merge commits partially verified as rebase rewrites the history and you are the original author of those two merge commits. Would you mind trying to rebase these two commits with your signature? Thanks. I think the only conflict might be just the name orders in CONTRIBUTORS.md. Sorry for the inconvenience.

@MetBenjaminWent MetBenjaminWent self-requested a review January 20, 2026 10:18
MetBenjaminWent
MetBenjaminWent approved these changes Jan 20, 2026
View reviewed changes
Copy link
Contributor

@MetBenjaminWent MetBenjaminWent left a comment

Choose a reason for hiding this comment

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

Re approving on request of James Bruten to test workflow retrigger.

@jasonjunweilyu jasonjunweilyu mentioned this pull request Jan 21, 2026
Open
@jasonjunweilyu
Copy link
Author

jasonjunweilyu commented Jan 21, 2026

Hi @MetBenjaminWent , I have just resolve conflicts between main and my branch by committing a merge request. For solving the technical debt of improving PSyclone script modularity, I have created an issue here: #167 .

Please let me know if you need me to do anything else. Otherwise, I will wait for you to merge this branch in without re-requesting review from you. Thanks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@MetBenjaminWent MetBenjaminWent MetBenjaminWent approved these changes

@mo-alistairp mo-alistairp mo-alistairp approved these changes

Assignees

@jasonjunweilyu jasonjunweilyu

Labels

cla-signed This contributor has signed the CLA.

Projects

None yet

Milestone

Spring 2026

Development

Successfully merging this pull request may close these issues.

NG-ARCH: Porting Stochastic Physics to GPU

4 participants

@jasonjunweilyu @mo-alistairp @MetBenjaminWent @yaswant