3D chunking for chemistry full-domain mode

[mn416]

PR Summary

Sci/Tech Reviewer: Joe Wallwork (@joewallwork)
Code Reviewer: Pierre Siddall (@Pierre-siddall)

UKCA provides various ways to invoke the chemistry solver: level-wise mode (with 1D chunking), column mode (with 1D chunking), and full-domain mode. This PR introduces 3D chunking for full-domain mode, generalising and subsuming all previous modes of operation while enabling new kinds of chunking that were not expressible before. For example, the new namelist entry

&tuning_segments
ukca_chem_full_chunk_size=8,4,2,

(assuming l_ukca_asad_full=.true.) specifies a chunk size of 8x4x2, i.e. 8 grid points per row, 4 rows, and 2 levels. A value <= 0 in any dimension will lead to a chunk size equal to the full-domain size in that dimension. This extra flexibility leads to performance improvements both with and without OpenMP enabled, as it is better able to exploit the spatial proximity of hard-to-solve grid points. Here are results comparing full-domain chunking (with a chunk size of 8x8x2) against the default UKCA configuration (column mode with a chunk size of a single column).

Full UM wall-clock run times for N48, 1 MPI task, 12 time steps:

  Threads | Column (default)       | Full Domain (8x8x2 chunk size)
  ------- | ---------------------- | ------------------------------
  1       | 205.5 + 591.8 = 797.3  | 168.9 + 569.3 = 738.2
  2       | 121.5 + 311.2 = 432.7  | 89.2 + 300.5  = 389.7
  4       |  83.0 + 172.6 = 255.6  | 50.2 + 164.5  = 214.7

Timings are of the form T_first + T_rest = T_all where T_first, T_rest, and T_all are the times taken for the first two UM time steps, the subsequent time steps, and all time steps respectively.

In summary, 3D chunking gives an overall improvement of 7.4-16.0%, broken down into a 17.8-39.5% improvement for the first two time steps, and a 3.8-4.7% improvement for subsequent time steps.

• linked MetOffice/um#96
• linked Update PSyclone scripts for UKCA chemistry 3D chunking lfric_apps#593

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
• 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 UKCA rose-stem suite
• If shared files have been modified, I have run the UM and LFRic Apps rose
  stem suites
• 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.)

The UM has been run locally on the N48 configuration for 1 MPI task and 6 time steps, with the following main results from mule-cumf:

1. Running full-domain mode with a chunk size equal to the full domain gives a perfect match when compared against the existing full-domain mode without chunking.

2. Running full-domain mode with a chunk size of 1x1x-1, i.e. 1 column, gives a perfect match when compared against column mode with a chunk size of a single column.

The UM has been run locally with various values for the new namelist parameter, including: (1) the use of -1 to signify no chunking in particular dimension; (2) a chunk size dimension that exceeds the full-domain size of a dimension, in which case the chunk size is truncated; (3) a chunk size that divides perfectly into the full domain size; and (3) a chunk size that does not divide perfectly into the full domain size, in which case the chunk size in non uniform. In all cases, chunking is performed as intended.

trac.log

Testing branch: https://github.com/RobWatersMet/um/tree/testing-mn416-3d-chunking

Macro has been applied and dependencies.sh updated with related UKCA branch

Test Suite Results - um - testing_mn416_3d_chunking/run1

Suite Information

Item	Value
Suite Name	testing_mn416_3d_chunking/run1
Suite User	robert.waters
Workflow Start	2026-08-20T14:01:49
Groups Run	all

Dependency	Reference	Main Like
casim	MetOffice/casim@2026.07.1	True
jules	MetOffice/jules@2026.07.1	True
moci	MetOffice/moci@2026.07.1	True
mule	MetOffice/mule@2026.07.1	True
shumlib	MetOffice/shumlib@2026.07.1	True
socrates	MetOffice/socrates@2026.07.1	True
SimSys_Scripts	MetOffice/SimSys_Scripts@2026.07.1	True
ukca	mn416/ukca@mn416-3d-chunking	True
um	RobWatersMet/um@testing-mn416-3d-chunking	False
um_aux	MetOffice/um_aux@2026.07.1	True
um_meta	MetOffice/um_meta@2026.07.1	True

Approvals

Code Owners

Config Owners

No UM Config Owners Required

Task Information

✅ succeeded tasks - 2770

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

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

[mn416]

Mohit Dalvi (@mcdalvi) Might you be SciTech reviewer for this one as you've looked at our chemsitry chunking work in the past?