Improve performance

[sirmarcel]

This is to track progress on removing "obvious" bottlenecks to improve performance.

Timing before starting work (for cutoff=10, 96 atoms ZrO2, script attached -- timings with profiler, so slower than real):

CPU (M3 Max)
cpu time: 109.082ms

GPU (H100)
cpu time: 1.724s
cuda time: 262.400ms

profile.zip

---

📚 Documentation preview 📚: https://meshlode--28.org.readthedocs.build/en/28/

[sirmarcel]

Timing after removing a for loop in the short-range part:

CPU (M3 Max):
cpu time: 42.564ms

CUDA (H100)
cpu time: 3.976ms
cuda time: 817.129us

[PicoCentauri]

Nice! Should we put a profiling example showing some breakdowns? Similar to https://luthaf.fr/rascaline/latest/examples/profiling.html maybe?

[sirmarcel]

I have the outputs on hand, but they're rather unwieldy, and not particularly informative -- maybe we can try to do something polished once this is done.

[sirmarcel]

Okay, optimised generate_kvectors a bit.

before (H100, N=4116):
Self CPU time total: 21.154ms
Self CUDA time total: 19.132ms

after (H100, N=4116):

Self CPU time total: 17.907ms
Self CUDA time total: 15.858ms

Now it seems that we're basically dominated by the FFTs and the convolution, which seems reasonable. I think the division by volume in pmepotential.py, line 125 is the next target, that's 1ms that seems avoidable... but not sure if it's worth doing. I think the most obvious things have been done now.

[sirmarcel]

Just for future reference, here are the timings for energy + forces. "Before" is main, "after" is the current state of this branch.

Before, CPU, M3 Max, 96 atoms
CPU time: 2.16s

(didn't bother with GPU)

After, CPU, M3 Max, 96 atoms
CPU time: 77.6ms

After, GPU, H100, 96 atoms
CPU time: 8.432ms
CUDA time: 1.680ms

After, GPU, H100, 4116 atoms
CPU time: 42.291ms
CUDA time: 39.356ms

After, GPU, H100, 8748 atoms
CPU time: 43.439ms
CUDA time: 40.820ms

[sirmarcel]

Test pass on kuma, removing draft status. Someone please check + merge. 🙏