CPU SP performance imbalance (IM=0 path) linked to allocations and SP→DP copies in LEDIR

[CarlosPenaDePedro]

Context

• Organisation: Barcelona Supercomputing Center (BSC)
• Project: DestinE ClimateDT

Summary

We observe a significant OpenMP thread imbalance in the single-precision CPU version of ectrans when running with IM=0.

The issue was first detected on MN5 GPP and also appears on LUMI-C, but the reproducer below is for MN5.

The imbalance becomes visible when GEMM is fast enough, e.g. with:

export MKL_CBWR=AUTO

On MN5 this selects AVX512 behavior and makes the SGEMM calls sufficiently fast that the overhead from temporary allocation and single-to-double copies in ledir_mod.F90 becomes a relevant fraction of runtime.

Versions

ecbuild: 3.8.5
fiat:    1.4.1
ectrans: 1.8.0

Platform

System: MN5 GPP partition
Nodes: 100
Cores per node: 112
Compiler: Intel ifort
MPI: Intel MPI 2021.10.0
MKL: 2023.2.0
FFTW: 3.3.10

Modules:

module load intel/2023.2.0 impi/2021.10.0 mkl/2023.2.0 ucx/1.16.0 fftw/3.3.10 cmake
export FC=ifort

Reproducer

Submission script:

#!/bin/bash
#SBATCH -J ectrans-tco2559
#SBATCH --qos=gp_bsces
#SBATCH --account=bsc32
#SBATCH -N 100
#SBATCH --ntasks-per-node=14
#SBATCH --cpus-per-task=8
#SBATCH --hint=nomultithread
#SBATCH --time=00:30:00

export OMP_NUM_THREADS=8
export OMP_PLACES=cores
export OMP_PROC_BIND=close

module load intel/2023.2.0 impi/2021.10.0 mkl/2023.2.0 ucx/1.16.0 fftw/3.3.10

export DR_HOOK=1
export DR_HOOK_OPT=prof

# Important: needed for GEMM to be fast enough for the imbalance to show clearly.
# On MN5 this selects AVX512 behavior.
export MKL_CBWR=AUTO

srun /gpfs/scratch/ehpc01/bsc032799/ectrans-1.8/ectrans/bin/ectrans-benchmark-cpu-sp \
  --truncation 1279 \
  --vordiv \
  --scders \
  --nlev 137 \
  --nfld 16 \
  --nproma 16 \
  --niter 50

Observed behavior

With DR_HOOK profiling, rank 0 and similar affected ranks show strong imbalance in LEDIR_SGEMM_1 and LEDIR_SGEMM_2 across OpenMP threads, being thread 0 much slower as it takes the DP path.

Example excerpt:

#  % Time         Cumul         Self        Total     # of calls        Self       Total    Routine@<thread-id>
                                                                         ms/call     ms/call
63     0.99        5.747        0.057        0.057             53        1.08        1.08   *LEDIR_SGEMM_1@1
66     0.33        5.787        0.019        0.019             53        0.37        0.37   *LEDIR_SGEMM_2@1
72     0.20        5.787        0.012        0.012            106        0.11        0.11    LEDIR_SGEMM_1@8
73     0.19        5.787        0.011        0.011             53        0.21        0.21    LEDIR_SGEMM_1@2
74     0.17        5.787        0.010        0.010            106        0.09        0.09    LEDIR_SGEMM_2@8
75     0.17        5.787        0.010        0.010            106        0.09        0.09    LEDIR_SGEMM_1@7
76     0.16        5.787        0.009        0.009             53        0.17        0.17    LEDIR_SGEMM_2@2
77     0.14        5.787        0.008        0.008            106        0.07        0.07    LEDIR_SGEMM_2@7
79     0.12        5.787        0.007        0.007             53        0.14        0.14    LEDIR_SGEMM_1@3
80     0.12        5.787        0.007        0.007             85        0.08        0.08    LEDIR_SGEMM_1@6
81     0.11        5.787        0.006        0.006             53        0.12        0.12    LEDIR_SGEMM_1@4
82     0.11        5.787        0.006        0.006             53        0.12        0.12    LEDIR_SGEMM_2@3
83     0.11        5.787        0.006        0.006             85        0.07        0.07    LEDIR_SGEMM_2@6
84     0.11        5.787        0.006        0.006             74        0.08        0.08    LEDIR_SGEMM_1@5
86     0.10        5.793        0.006        0.006             53        0.11        0.11    LEDIR_SGEMM_2@4
87     0.09        5.793        0.005        0.005             74        0.07        0.07    LEDIR_SGEMM_2@5

Additional profiling / suspected source

After adding extra DR_HOOK regions in:

src/trans/cpu/internal/ledir_mod.F90

The imbalance is associated with the IM=0 path in the single-precision version. As we understand it, this path is required to preserve conservation properties and maintain scientific accuracy.

However, in large-scale runs we observe that this imbalance becomes more pronounced. In particular, there are cases where a single DGEMM-based path becomes slower than executing multiple SGEMM paths.

Additional timers added around this region suggest that part of this extra cost may come from:

• Allocation/deallocation of temporary double-precision arrays
• Copy and conversion of the matrix slice from single to double precision before the GEMM call

The intention of this issue is not to question the IM=0 logic itself, but to highlight that the allocation and copy overheads could potentially be reduced or avoided (e.g. via caching or reuse), which may help mitigate the observed imbalance.

Relevant code pattern:

ELSE
  BLOCK
     REAL(KIND=JPRD), allocatable :: ZB_D(:,:), ZCS_D(:,:), ZRPNMS(:,:)
     INTEGER(KIND=JPIM) :: I1, I2, I3, I4

     I1 = size(S%FA(KMLOC)%RPNMS(:,1))
     I2 = size(S%FA(KMLOC)%RPNMS(1,:))
     ALLOCATE(ZRPNMS(I1,I2))
     ALLOCATE(ZB_D(KDGLU,KIFC))
     ALLOCATE(ZCS_D((R%NTMAX-KM+3)/2,KIFC))

     IFLD=0
     DO JK=1,KFC,ISKIP
        IFLD=IFLD+1
        DO J=1,KDGLU
           ZB_D(J,IFLD)=PSIA(JK,ISL+J-1)*REAL(PW(ISL+J-1),JPRB)
        ENDDO
     ENDDO

     DO I3=1,I1
        DO I4=1,I2
           ZRPNMS(I3,I4) = S%FA(KMLOC)%RPNMS(I3,I4)
        END DO
     END DO

     CALL GEMM('T','N',ILS,KIFC,KDGLU,1.0_JPRD,ZRPNMS,KDGLU,&
          &ZB_D,KDGLU,0._JPRD,ZCS_D,ILS)

     IFLD=0
     DO JK=1,KFC,ISKIP
        IFLD=IFLD+1
        DO J=1,ILS
           ZCS(J,IFLD) = ZCS_D(J,IFLD)
        ENDDO
     ENDDO

     DEALLOCATE(ZRPNMS)
     DEALLOCATE(ZB_D)
     DEALLOCATE(ZCS_D)
  END BLOCK
END IF

When the GEMM duration is very small, the allocation/deallocation and the repeated copy/conversion of S%FA(KMLOC)%RPNMS into ZRPNMS become a noticeable part of the cost and may contribute to the observed imbalance.

Expected behavior

More balanced OpenMP thread execution in the single-precision CPU path when IM=0.

Actual behavior

Clear thread imbalance in LEDIR_SGEMM_* regions, especially at scale and when GEMM is fast.

Possible improvement

A potential optimization would be to trade memory for performance by caching the double-precision version of the IM=0 matrix slice in the single-precision path.

For example:

• Build the double-precision slice once (e.g. in SULEG)
• Reuse it in LEDIR
• Avoid repeated allocation, copy, and conversion

This would remove:

• Repeated allocation/deallocation of ZRPNMS
• Repeated single → double conversion of the matrix slice

[samhatfield]

Thanks for the detailed write up @CarlosPenaDePedro, this is very impressive.

The fix for the zeroth mode in the direct transform came about 7 years ago after extensive testing of the single-precision variant of the IFS. At that time the model was subject to substantial testing for its suitability in operational forecasting. A conservation issue stemming from the direct Legendre transform was identified, in which rounding errors in this mode can accumulate and lead to spurious zonal bands which were very pronounced in error plots, especially after aggregating over many forecast dates. The solution at the time (admittedly in a bit of a rush) was to just promote this mode to double precision. Probably there is a more elegant solution, such as using a compensated summation SGEMM or equivalent. Unfortunately, especially given that this hack has "sort of worked" up until now, nobody has had the free time to return to this issue and propose something better.

If the load imbalance you detected really starts to hurt us then this may change. I would love to come up with a better solution than "just use double precision", but pending that, I'm very open to suggestions to improve the hack.

What you suggest is a good one, and we already use a cached work array approach elsewhere, e.g. here.

One other idea is to modify the distribution of zonal wavenumbers so that the MPI tasks handling the zero mode get fewer modes overall, and some of their modes are reassigned to other MPI tasks. This would also improve load balancing. However I'm reluctant to start with this approach as I would have to play around with the decomposition logic which isn't such a nice job.

So for now let's try to cache ZRPNMA, ZB_D, ZCA_D etc. and see if this improves the load imbalance. I'm happy to take this on but it will have to wait for some time as I'm quite busy at the moment. If you want to have a go yourself @CarlosPenaDePedro you're more than welcome!

[CarlosPenaDePedro]

Thanks a lot for the detailed explanation and the extra context, it makes much more sense now where this workaround came from and why it has stayed in place for so long.

I can imagine that moving towards something like compensated summation would not be straightforward either. Most optimized BLAS/GEMM libraries do not really support that kind of approach directly, so doing it efficiently would probably require custom work, plus quite a lot of additional validation given how sensitive the transforms are.

I'm also relatively busy at the moment, but I’d still be happy to take a look at the caching approach and share any results or observations if I manage to make some progress on it.

[CarlosPenaDePedro]

Hello,

I have a first working implementation that provides improved timing results.

Ledir calls on the problematic threads have gone from 10–12 ms per call down to 3 ms per call. Compared to the 1–1.7 ms range of the threads that do not process IM=0, the imbalance is now much smaller.

The implementation is available in the following branch, which is based on the 1.8.0 tag:

BRANCH

I have not tested bit-identicality, but I expect it may no longer be preserved due to changes in rounding and the use of double precision. The values are now computed directly in double precision during initialization, rather than being stored first in single precision and cast to double precision later when calling DGEMM. This removes an intermediate rounding step, so small numerical differences are expected.

[samhatfield]

Great work @CarlosPenaDePedro! Many thanks for this. I'll take a look later, but in the mean time, could you rebase your branch on top of develop? develop contains a few fixes which should improve stability of the CI.

[CarlosPenaDePedro]

Thanks for taking a look! No worries at all, there's no hurry.

I've already rebased the changes on top of develop and pushed them to a new branch:

Rebased_Branch

[samhatfield]

I've been trying to reproduce the load imbalance in a simpler setup. Here's what I did. Firstly I built an ecTrans bundle on ECMWF / HPC2020 (using the bundle merged in a recent PR):

./package/bundle/ectrans-bundle create --bundle ./package/bundle/bundle.yml --src-dir=./package/bundle/source
./package/bundle/ectrans-bundle build --build-dir=./package/bundle/build --src-dir=./package/bundle/source --arch=ecmwf/hpc2020/intel/2023.2.0/hpcx-openmpi/2.9.0

Then I ran a weak-scaling scenario using a similar job script to @CarlosPenaDePedro:

#!/bin/bash
#SBATCH --nodes=<nodes>
#SBATCH --qos=np
#SBATCH --ntasks-per-node=16
#SBATCH --cpus-per-task=8
#SBATCH --hint=nomultithread
#SBATCH --time=00:30:00

export OMP_NUM_THREADS=8
export OMP_PLACES=cores
export OMP_PROC_BIND=close
export OMP_STACKSIZE=256M

ROOT=/path/to/ectrans/package/bundle/build
source $ROOT/env.sh

export DR_HOOK=1
export DR_HOOK_OPT=prof

export MKL_CBWR=AUTO

srun $ROOT/bin/ectrans-benchmark-cpu-sp --truncation <truncation> --nlev <nlev> --nprtrv 1

Because I have --nprtrv 1, I can reproduce a similar pressure on the Legendre transform from running on the normal number of nodes with much fewer nodes. For example, I can run 1279 on 1 node with 10 levels. This gives roughly the same kind of GEMM sizes as running 1279 on 10 nodes with 137 levels. Similarly, 2559 on 4 nodes with 5 levels, 3999 on 10 nodes with 3 levels, and 5199 on 16 nodes with 2 levels. Very rough.

I ran this weak scaling set for the develop branch, and for a slightly modified branch in which there is no special treatment for the m = 0 mode, i.e. all modes use SGEMM.

Region	1279 (develop)	2559 (develop)	3999 (develop)	5199 (develop)	1279 (all sgemm)	2559 (all sgemm)	3999 (all sgemm)	5199 (all sgemm)
LEDIR_SGEMM_1@1	4.85	37.72	99.42	201.21	2.66	5.96	12.99	32.43
LEDIR_SGEMM_1@2	2.1	6.62	11.64	17.32	2.24	6.19	12.78	18.37
LEDIR_SGEMM_1@3	1.7	4.97	9.78	15.79	1.89	5.2	11.06	17.22
LEDIR_SGEMM_1@4	1.19	3.78	7.13	9.51	1.37	3.86	8	10.44
LEDIR_SGEMM_1@5	0.97	3.03	6.48	9.63	1.04	2.86	6.03	9.66
LEDIR_SGEMM_1@6	0.92	2.77	6.46	8.66	0.95	2.31	6.52	8.85
LEDIR_SGEMM_1@7	0.86	2.7	5.25	10.45	0.88	2.35	5.41	10.42
LEDIR_SGEMM_1@8	0.81	2.51	4.47	9.86	0.85	2.19	4.2	9.81
Max	4.85	37.72	99.42	201.21	2.66	6.19	12.99	32.43
Min	0.81	2.51	4.47	8.66	0.85	2.19	4.20	8.85
Ratio	5.99	15.03	22.24	23.23	3.13	2.83	3.09	3.66

When everything is single precision, the imbalance between the fastest and slowest threads is always about 3.

When the m = 0 mode is double precision, this load imbalance increases. In fact for T5199 it goes up to 23x.

Next I'll try your branch @CarlosPenaDePedro.

[samhatfield]

Here are the results from Ledir_IM0_caching_rebased:

Region	1279 (m0 caching)	2559 (m0 caching)	3999 (m0 caching)	5199 (m0 caching)
LEDIR_SGEMM_1@1	2.74	9.71	32.09	59.65
LEDIR_SGEMM_1@2	2.12	6.42	14.04	18.7
LEDIR_SGEMM_1@3	1.81	4.88	11.56	18.07
LEDIR_SGEMM_1@4	1.31	3.71	8.33	10.85
LEDIR_SGEMM_1@5	1	3.02	7.03	9.77
LEDIR_SGEMM_1@6	0.95	2.66	6.89	8.87
LEDIR_SGEMM_1@7	0.88	2.53	5.61	10.66
LEDIR_SGEMM_1@8	0.83	2.29	4.89	10.02
Max	2.74	9.71	32.09	59.65
Min	0.83	2.29	4.89	8.87
Ratio	3.30	4.24	6.56	6.72

That's pretty good! The imbalance for T5199 is reduced from 23x down to 6x. Also some big improvements at other resolutions. Is this what you expected to see @CarlosPenaDePedro?

[CarlosPenaDePedro]

Very nice work. It's impressive how many resolutions you've been able to test using such a modest amount of resources. I hadn't realised that the problem could be reproduced so effectively in this reduced setup.

It's very interesting to see how the imbalance grows with resolution. I observed the same trend in the coupled TCo-eORCA12 runs: at TCo1279 the imbalance represented only about 2% of the total runtime, while at TCo2559 it increased to roughly 8%.

It's great to see that the patch performs so well. Reducing the TCo5199 imbalance from 23× to around 6× is a substantial improvement! I'm still surprised that the allocation and type conversion in the m = 0 path could take so much time compared to the DGEMM itself.

One thing I find particularly interesting is that the "all SGEMM" build still shows an imbalance of about 3×. Are the DR_HOOK numbers you reported average times per call, or self total times in seconds?

If they are average times per call, would it be possible to also share the self total times? I think those would give a clearer picture of the overall balance, since one thread could simply be executing fewer calls with a higher average duration rather than doing more total work.

If they are already self total times, then maybe the remaining imbalance comes from the task granularity: not enough sets/work units to be evenly distributed across threads, especially at the higher truncations?

[CarlosPenaDePedro]

Hello, The DestinE Phase 3 kick-off meeting is taking place right now, and they have started discussing the timeline for the next operational cycle. The current target for the model freeze appears to be around mid-October 2026.

To be completely honest, do you think an ectrans version including this fix could realistically be released in time for that schedule, or should we instead plan to target the following operational cycle?

[samhatfield]

Hi @CarlosPenaDePedro, sorry responses will be bit patchy for the next month or so. Yes, mid-October gives plenty of time. Right now this seems like a straightforward fix. Why don't you go ahead and open a PR with your branch, and then we can take the discussion there and iterate on your first draft?

As for your other questions, all good questions yes. I'll try to pick this up again sometime and get back to you.

[CarlosPenaDePedro]

Perfect! I’ll open a PR with my branch and we can continue the discussion there.

No worries about the delayed responses, and thanks for taking a look!