Add Kokkos support for atom sorting on device

[stanmoore1]

Summary

Currently atom sorting is done on host CPU when using Kokkos. This is expensive since it must move all the data, and is performed in serial. This PR adds Kokkos support for atom sorting on device.

Related Issue(s)

None

Author(s)

Stan Moore (SNL)

Licensing

By submitting this pull request, I agree, that my contribution will be included in LAMMPS and redistributed under either the GNU General Public License version 2 (GPL v2) or the GNU Lesser General Public License version 2.1 (LGPL v2.1).

Backward Compatibility

Yes

[stanmoore1]

For 8M LJ atoms running for 1000 timesteps on a single V100 GPU, this PR gives ~25% speedup when sorting every 100 timesteps with the new code vs the old code and default sorting every 1000 timesteps. Part of the speedup is due to a better sorting binsize, and the other part is due to sorting more often on device which improves cache access.

[stanmoore1]

This gives a marginal speedup for 256k Rhodo benchmark, but only when using the old default binsize for sorting = 6 Angstrom instead of 12, otherwise it gives a slowdown. Not sure how to get a better heuristic for sorting binsize.

[stanmoore1]

@weinbe2 @arghdos

[stanmoore1]

Looks like I need to run performance tests to see if the majority of styles get a speedup or slowdown with the change in default sorting binsize.

[stanmoore1]

3 regression tests failing on the GPU, will debug those too

[sjplimp]

@stanmoore1 My recollection is that the optimal sort frequency was also different on CPU vs GPU? Is that a default
setting that should be different between the two modes?

[stanmoore1]

@sjplimp good point, we should consider changing the default from 1000 to 100, will do some tests.

[stanmoore1]

Regressions appear to be false positives related to changing binsize affecting the pRNG in fix langevin and fix gle.

[stanmoore1]

The sorting binsize change seems to help large LJ but hurts most other benchmarks like ReaxFF, Tersoff, EAM, and Rhodo so I will revert.

[stanmoore1]

Looks like sorting every 100 helps LJ but slows down other benchmarks too.

[stanmoore1]

This is ready to merge from my POV. Noted that LJ can get a speedup from using a different binsize and sorting every 100 timesteps, but this won't be the default behavior since it hurts other benchmarks.

[stanmoore1]

Actually seeing a slowdown for ReaxFF, need to double check

[sjplimp]

@stanmoore1 I don't understand if the atom_modify sort setttings affect this new on-GPU sorting. If they do, it seems like there should be info added to the atom_modify doc page to explain how to adjust these settings for both CPUs and GPUs, and also what the defaults are for both. Note that the atom_modify settings are for bin size and frequency.

[stanmoore1]

@sjplimp no I reverted all the changes to the atom_modify sort settings since they cause performance regressions for some benchmarks, keeping the old defaults for now. So the only change in this PR is to sort atoms on the GPU instead of CPU.

[sjplimp]

@stanmoore1 Hi Stan - my point is that the atom_modify sort doc page has no discussion of CPU vs GPU, including for the defaults. If you expect the user to do something different for GPUs, then that should be explained, even if it is just guidance on what to try for GPUs. Are you saying the user should not do anything different in their input script for GPUs (for sorting), and that the code should now simply run faster (for sorting) b/c it is being done on the GPU? However you also seem to be saying that sometimes sorting on the GPU is slower than on the CPU? Which doesn't make that much sense to me.

[stanmoore1]

The atom_modify sort defaults are sortfreq = 1000 and binsize = cutneighmax/2.0. For the LJ benchmark on GPUs, using sortfreq = 100 and binsize = cutneighmax is better but that hurts performance of other benchmarks like Tersoff on the GPU so we can't make that the default for GPUs. I could mention this in the docs.

[sjplimp]

If you mean atom_modify sort, then yes, I think that is a good thing to do, and that is best place to do it.
You should probably also mention it on the Kokkos package page under 7.4 Accelerator packages in the manual.

[stanmoore1]

@sjplimp will do

[stanmoore1]

The performance regression I saw with ReaxFF is related to #3756.

[stanmoore1]

This PR is ready to merge from my POV, but should be merged after #3756 since I already merged those changes into this PR.

[stanmoore1]

Actually forgot to update the docs, will do that now.

[athomps]

I approve.