Rewrite ApplySymmetryToTensor3 with NumPy/BLAS (11x speedup) - #125
Rewrite ApplySymmetryToTensor3 with NumPy/BLAS (11x speedup)#125psavery wants to merge 1 commit into
Conversation
Profiling the free-energy Hessian of a 128-atom PdH supercell in python-sscha showed 82% of the run inside this method: the Fortran symph routines apply the symmetries element by element in serial. The three steps are redone with numpy array operations: the permutation symmetry as an average over the 6 index orderings, the translation average computed once per group of equivalent atoms, and each symmetry at gamma applied as one matrix product over the cartesian components. On the PdH case: 181 s -> 16 s, matching the old result to 1e-16. Also checked on two lower-symmetry cells (120 and 216 atoms). Peak memory grows because of numpy temporaries: on the 216-atom cell about 11 GB against 6.5 GB, with the time going 211 s -> 34 s. The matrix products use numpy's threaded BLAS but do not depend on it: on the 120-atom cell the Fortran path takes 20.4 s and the numpy one 6.5 s on a single thread (3.8 s with threads on a 24-core machine). On shared nodes without core pinning, cap the threads with OPENBLAS_NUM_THREADS. The Fortran routines (permute_v3, trans_v3, sym_v3) no longer have callers here and can be removed in a follow-up if preferred. Signed-off-by: Patrick Avery <patrick.avery@kitware.com>
|
Hi @psavery , thanks for the PR. |
|
I'll take a look at processing them in blocks to reduce peak memory usage, and also report back how much performance is affected! |
|
I'm at a conference this week though, so I may need to get to it next week. |
|
Hi yes, no hurry at all. Also maybe wait a moment to implement the batch, as we have two other pending PRs on the free energy hessian, so probably merging will generate new conflicts to solve. One of which is focused in speeding up and reducing the memory footprint of the fortan code, so for sure there are conflicts there. |
2 participants
Profiling the free-energy Hessian of a 128-atom PdH supercell in python-sscha showed 82% of the run inside this method: the Fortran symph routines apply the symmetries element by element in serial.
The three steps are redone with numpy array operations: the permutation symmetry as an average over the 6 index orderings, the translation average computed once per group of equivalent atoms, and each symmetry at gamma applied as one matrix product over the cartesian components.
On the PdH case: 181 s -> 16 s, matching the old result to 1e-16. Also checked on two lower-symmetry cells (120 and 216 atoms). Peak memory grows because of numpy temporaries: on the 216-atom cell about 11 GB against 6.5 GB, with the time going 211 s -> 34 s.
The matrix products use numpy's threaded BLAS but do not depend on it: on the 120-atom cell the Fortran path takes 20.4 s and the numpy one 6.5 s on a single thread (3.8 s with threads on a 24-core machine). On shared nodes without core pinning, cap the threads with OPENBLAS_NUM_THREADS.
The Fortran routines (permute_v3, trans_v3, sym_v3) no longer have callers here and can be removed in a follow-up if preferred.
The other half of this is SSCHAcode/python-sscha#427