Make extract_atom kokkos-aware

[rbberger]

This is to document the discussion I had with Stan about the current state of extract_atom in the library interface and how it interacts with data that lives on device. Currently when using Kokkos LAMMPS does a sync during each output step. However, we do already allow people to write Python scripts that run at arbitrary intervals that are not necessarily connected to output.

Thus, if those intervals are not aligned with output steps, they will likely access outdated data on the host. Since we don't want to add a burden to users of extract_atom, the suggestion is to override extract in AtomKokkos, use a std::map<std::string, int> to map the name to the mask value and do a sync on that mask, prior to returning the result of the base Atom::extract()

I'll work on this, but do other people from @lammps/core have any comments on this?

[stanmoore1]

@rbberger good catch, I think this is a bug since it will give incorrect results without the proposed changes here.

[rbberger]

@stanmoore1 after reviewing the code, I'm concluding that at least from the Python features we have at the moment we are in the clear, since those are all fixes that are kokkosable=0, which means their existence will cause a full sync before and after their calls. In a way that is the safe thing to do, since you don't know what things will access from within Python. However, it might be interesting to turn off this aggressive auto_sync to avoid copies in some cases and be more selective by hand.

[stanmoore1]

I talked to @rbberger and there is also a sync to host at the end of a run, so this should be safe unless library calls are being made during a run, which I don't think is possible.

[akohlmey]

unless library calls are being made during a run, which I don't think is possible.

LAMMPS GUI makes library interface calls during LAMMPS runs. This is possible by running the LAMMPS simulation as a separate thread. You can also access the library interface from within a "python" command or other PYTHON package calls to Python from LAMMPS.

There are limits as to what can be done, though. For example, any LAMMPS command will trigger an error when issued during a run. But "extract" type functions should be possible. MPI calls can be a problem, when called from a separate there but won't be supported by LAMMPS GUI anyway. MPI calls should be safe during "python" commands, since all MPI ranks will be entering Python.

[stanmoore1]

But "extract" type functions should be possible.

OK so it seems like this is a bug and it should be fixed.

[akohlmey]

But "extract" type functions should be possible.

OK so it seems like this is a bug and it should be fixed.

Probably, this is all new stuff and for example with LAMMPS GUI and intercepting errors through exception handling, this is exploring code paths that were previously not possible. The standalone LAMMPS executable just stop with an error, but LAMMPS gui will recover, issue a "clear" command or recreate the LAMMPS instance and start over. I am regularly discovering new ways to crash LAMMPS GUI because of having some inconsistent internal data and thus requiring changes to the core code. E.g. see the changes to Force and Update in PR #3950

A wild card in this process are the code paths used by the Cython based code. I have no idea how they deal with KOKKOS data.