Ledir IM=0 DP matrix caching

[CarlosPenaDePedro]

Description

This PR adds persistent double-precision caches for the KM=0 DGEMM path in LEDIR.

Specifically, it introduces persistent allocatable arrays:

RPNMA_DGEMM
RPNMS_DGEMM

to avoid repeated allocations/recomputations in the DGEMM execution path for the KM=0 case. This is intended to reduce overhead and help improve load imbalance observed in this path.

More information in issue #392

The change passes the GitHub Actions test suite.

Contributor Declaration

By opening this pull request, I affirm the following:

• All authors agree to the Contributor License Agreement.
• The code follows the project's coding standards.
• I have performed self-review and added comments where needed.
• I have added or updated tests to verify that my changes are effective and functional.
• I have run all existing tests and confirmed they pass.

[samhatfield]

Thanks for this amazing work @CarlosPenaDePedro - we will start the review when we have time, but definitely well before the October deadline you mention 🙏🏻

[Copilot]

Pull request overview

This PR optimizes the single-precision KM=0 LEDIR DGEMM path by caching double-precision Legendre polynomial matrices at setup time, reducing repeated allocation and SP→DP conversion overhead.

Changes:

• Adds persistent RPNMA_DGEMM and RPNMS_DGEMM caches to the CPU FLT resolution wrapper.
• Populates the caches in SULEG for IM == 0 when running single precision.
• Updates LEDIR to use the cached double-precision matrices in the KM=0 DGEMM path.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
src/trans/cpu/internal/tpm_flt.F90	Adds persistent double-precision cache arrays to FLT state.
src/trans/cpu/internal/suleg_mod.F90	Allocates and fills the KM/IM=0 DGEMM caches during Legendre setup.
src/trans/cpu/internal/ledir_mod.F90	Replaces per-call matrix allocation/conversion with cached double-precision matrices.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[samhatfield]

Oof, and that's why we use Copilot. I will take a look at this.

[samhatfield]

Don't worry about the failing tests. Once #407 is merged, you can then rebase this again against develop and that should resolve them.

[samhatfield]

Could you rebase against develop again @CarlosPenaDePedro? That should fix the failing tests.

[samhatfield]

Todo:

• Rebase against develop
• Approve ECMWF contributions to Ledir_IM0_caching_rebased CarlosPenaDePedro/ectrans#1 which tries to mitigate repeated code and also some style fixes to LEDIR
• Check all comments above
• Verify that all checks pass
• For @samhatfield, check potential clash with READ_LEGPOL code path

[samhatfield]

All tests passed. Well done @CarlosPenaDePedro, you are off the critical path and can now relax 😝

I still need to take a closer look at the codepath for reading / writing the Legendre polynomials to see if there's a clash.

Also I'd like to repeat my benchmarking from earlier as a sanity check. Load balance is one thing, but it would be good to check that there is no negative unexpected impact on overall wall times etc. I'm travelling at the moment but hopefully next week I'll find some time.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

[samhatfield]

The issue Copilot identifies above is the following:

ecTrans can read pre-computed Legendre polynomials from file or a memory buffer. In this case, the computation of the polynomials is skipped. That includes the allocation and initialisation of S%RPNMA_DGEMM and S%RPNMS_DGEMM. In that case, LEDIR will segfault because those arrays are still accessed.

The proper fix is to make sure the cached arrays are initialised even when polynomials are read from file or buffer. Rather than have the double-precision IM=0 arrays written to and read from disk, I think we should simply copy the single-precision polynomials which are read into those cached arrays. This will incur a single -> double upcast which is different to what we have now, where the double-precision source arrays for the polynomials are stored directly in the cache. Then, we only have to make a slight modification in READ_LEGPOL_MOD.

It's actually closer to what we do in the develop branch. There in LEDIR we copy the stored single-precision polynomials into a double-precision buffer (so again, a single -> double upcast). This works, and implies that the problem with the IM=0 mode is not that the polynomials are too low in precision but rather the GEMM itself is too low in precision.

[CarlosPenaDePedro]

In that case we could to keep the behaviour consistent across setup paths, we could also populate the SULEG S%RPNMA_DGEMM and S%RPNMS_DGEMM by upcasting the already stored single-precision RPNMA/RPNMS arrays, instead of filling it directly from the double-precision source values. This would make the SULEG and READ_LEGPOL paths behave the same way, and should also be closer to the previous implementation, where LEDIR promoted the single-precision matrix to double at runtime.

The issue Copilot identifies above is the following:

ecTrans can read pre-computed Legendre polynomials from file or a memory buffer. In this case, the computation of the polynomials is skipped. That includes the allocation and initialisation of S%RPNMA_DGEMM and S%RPNMS_DGEMM. In that case, LEDIR will segfault because those arrays are still accessed.

The proper fix is to make sure the cached arrays are initialised even when polynomials are read from file or buffer. Rather than have the double-precision IM=0 arrays written to and read from disk, I think we should simply copy the single-precision polynomials which are read into those cached arrays. This will incur a single -> double upcast which is different to what we have now, where the double-precision source arrays for the polynomials are stored directly in the cache. Then, we only have to make a slight modification in READ_LEGPOL_MOD.

It's actually closer to what we do in the develop branch. There in LEDIR we copy the stored single-precision polynomials into a double-precision buffer (so again, a single -> double upcast). This works, and implies that the problem with the IM=0 mode is not that the polynomials are too low in precision but rather the GEMM itself is too low in precision.

[samhatfield]

In that case we could to keep the behaviour consistent across setup paths, we could also populate the SULEG S%RPNMA_DGEMM and S%RPNMS_DGEMM by upcasting the already stored single-precision RPNMA/RPNMS arrays, instead of filling it directly from the double-precision source values. This would make the SULEG and READ_LEGPOL paths behave the same way, and should also be closer to the previous implementation, where LEDIR promoted the single-precision matrix to double at runtime.

Yes, that's exactly what I was thinking.