[DRAFT] Contract the anharmonic matrix elements one index at a time

[utksi]

The PR in general is about certain speedups that seem to be on the table.

Building a three- or four-phonon matrix element is contracting a force constant tensor with one eigenvector per leg. In several places TDEP does this by forming the full outer product of the eigenvectors and then contracting the whole thing down in a single pass. The contraction is separable, so peeling one index at a time gets the same number for a lot less arithmetic. That is what most of the commits to this branch (in this PR) are about.

Per q-quartet the four-phonon element goes from 4n^7 to 4n^4 operations, and the three-phonon element from 4n^5 to 3n^3, where n is the number of modes. The same substitution goes in six places: the scattering rates, the third-and fourth-order free energies, the lineshape three- and four-phonon self energies, and the dielectric matrix elements. Two smaller commits also are there.

---

1. Speedup variance

How much this helps depends strongly on the cell. The matrix elements grow faster with mode count than the rest of the scattering loop does, so a six-mode cell like MgO spends most of its time outside the code being changed, while a sixty-mode one spends almost all of it inside. The tables below therefore give times rather than ratios: the ratio says as much about the cell as about the change.

2. On whether things are correct.

outfile.thermal_conductivity is identical to the current main across every case I ran: MgO from tests/ at 12^3, 20^3, 28^3, 36^3 and 44^3 for three phonons and 8^3, 12^3 for four, and a 24-mode Tl3VSe4 cell at 12^3 and 16^3.

3. Setup for timings

Wall clock in seconds around each mpirun call, so queue and startup time are excluded. All runs were on Arrhenius (NAISS), 64 MPI ranks (unless stated otherwise) on one exclusively allocated AMD EPYC 9755 (Zen 5) node, gcc 14.3 with -O3 -march=znver5, OpenMPI 5.0.8, FlexiBLAS on the AOCL backend, OMP_NUM_THREADS=1.

Everything here runs on the full q-grid. I do not quote subsampled numbers, but they are of course faster too.

The BLAS makes very little difference here, which is probably worth saying since these commits move work into zgemm. Sweeping FlexiBLAS across AOCL, BLIS, MKL, OpenBLAS and reference NETLIB on the same case moves the runtime by about 10%, and the ordering of the two branches does not change.

One thing about how to read these: --fourthorder is cumulative rather than exclusive, so the four-phonon rows include the three-phonon scattering as well. The lineshape rows are third-order only unless marked otherwise.

MgO, from tests/thermal_conductivity:

case	main	this
3ph -qg 12 12 12	2s	2s
3ph -qg 20 20 20	6s	4s
3ph -qg 28 28 28	27s	19s
3ph -qg 36 36 36	234s	74s
3ph -qg 44 44 44	348s	232s
4ph -qg 8 8 8	53s	14s
4ph -qg 12 12 12	1708s	415s
lineshape --path	6s	6s

• Take the 36^3 row with a pinch (or a fistful) of salt. Main spends 134 of its 234 seconds idle there against 6 of our 74, so most of that gap is a load imbalance.

Tl3VSe4, 24 modes:

case	main	this
3ph -qg 12 12 12	144s	12s
3ph -qg 16 16 16	738s	45s
4ph -qg 6 6 6	>34 h	30s
4ph -qg 8 8 8	not run	292s
anharmonic free energy	14s	2s
lineshape --path	365s	164s

The lineshape row is somewhat soft; that runtime varies by tens of percent between repeats, so I think it should be read as clearly faster rather than some number.

How much this helps also depends on how many ranks one gives it, because the work 'removed' is a part that parallelises well. Tl3VSe4 3ph 12^3 on one node:

ranks	main	this
16	530s	33s
64	145s	12s
128	85s	9s
256	69s	15s

Past128 ranks, get slower since what is left is startup and communication rather than matrix elements.

Everything in the MgO table runs from inputs already in tests/, so the correctness part is reproducible. Tl3VSe4 is 8 atoms in the primitive cell, 24 modes, and about 19 MB of force constants once third and fourth order are included, which is more than seems reasonable to add to a tests/ and examples/ that come to 35 MB between them. I put it in a branch of my fork:

git clone --branch fix/matrix-element-contraction-data --single-branch --depth 1 https://github.com/utksi/tdep.git tdep-bench

That branch is this branch plus a benchmarks/Tl3VSe4/ directory, so it builds and runs without fetching anything else. Tl3VSe4 force constants courtesy of Dr. Ashis Kundu (LiU).

The README there should give the cell, the cutoffs, and the exact command for each row above. If you would rather have the case in-tree it could be its own PR. In general it might be good to have a 'bigger' cell for thermal conductivity in examples?

In principle, I have some numbers for CsPbBr3 (20 atom orthorhombic primitive cell) as well; but could not do a fair comparison against current upstream:main when increasing grid-size. Calculated without SOC, the idea was to just observe the scale-up.
But, if someone is experienced with this cell-size (also it's polar), then the numbers are:

case	main	this branch
3ph -qg 8 8 8	16423s	27s
3ph -qg 12 12 12	not run	155s
3ph -qg 16 16 16	not run	763s
4ph -qg 6 6 6	not run	10908s

I did run one upstream point after all, the 8x8x8 three-phonon case, which gives a somewhat dramatic speedup. Output is identical. The rest I left alone for now.

4. Minor: ctr not initialized

(The changes are minuscule, but I can see some funky non-recommended compiler making a mess).

Initialise ctr in free_energy_thirdorder is a correctness fix. ctr is read before it is assigned, and the loop hands work out with mod(ctr, mw%n) .ne. mw%r, which only partitions anything if every rank is counting from the same place. Uninitialised they are not. Printing the value on entry across a 64 rank run gave 32764 on eighteen ranks, 32765 on seventeen, 32766 on seventeen and 32767 on twelve, so the map from q-pair to rank is neither onto nor one to one: some pairs are summed twice, others by nobody.

The third order free energy therefore depends on how many ranks you use, and two runs at the same rank count need not agree. On Tl3VSe4, free energy column of the second cumulant block in eV/atom:

ranks	main	this branch
1	-0.001495587666	-0.001495587666
2	-0.001497622239	-0.001495587666
8	-0.001496650567	-0.001495587666
64	-0.001495499089	-0.001495587666

A second 64 rank run of main gave -0.001495733420. One rank is right by accident because mod(anything, 1) is 0. Adding the single ctr = 0 fixes this. free_energy_fourthorder_secondorder has the same omission, though nothing calls it at present it seems.

Commits

They are independent and can be taken or dropped one at a time.

Sub-optimal runs?

It's entirely possible I ran tdep in a sub-optimal way. The thing in favor should be operation count/complexity dropping for identical outputs, so it should win. Someone more experienced with using tdep can perhaps test this (maybe also on a different compute resource).

[flokno]

Great work! Thank you Utkarsh