Kokkos interface

cmake/Modules/Packages/KOKKOS.cmake

@@ -182,6 +182,11 @@ if(PKG_ML-IAP)
   endif()
 endif()
 
+if(PKG_ML-METATENSOR)
+  list(APPEND KOKKOS_PKG_SOURCES ${KOKKOS_PKG_SOURCES_DIR}/metatensor_system_kokkos.cpp)
+
+endif()
+
 if(PKG_PHONON)
   list(APPEND KOKKOS_PKG_SOURCES ${KOKKOS_PKG_SOURCES_DIR}/dynamical_matrix_kokkos.cpp)
   list(APPEND KOKKOS_PKG_SOURCES ${KOKKOS_PKG_SOURCES_DIR}/third_order_kokkos.cpp)

doc/src/Build_extras.rst

@@ -1049,6 +1049,17 @@ https://pytorch.org/get-started/locally/.
          make yes-metatensor
          make <machine>
 
+   .. tab:: Metatensor and Kokkos
+
+      The metatensor-kokkos interface should be compiled as
+
+      .. code-block:: bash
+
+         cmake ../cmake/ -DPKG_KOKKOS=ON -DKokkos_ENABLE_CUDA=ON -DPKG_ML-METATENSOR=ON -DCMAKE_PREFIX_PATH=/.../libtorch/share/cmake/
+
+      where ``/.../libtorch/`` is the path to a libtorch C++11 ABI distribution (which can be downloaded from https://pytorch.org/get-started/locally/).
+      The OpenMP version (as opposed to the CUDA version) can be enabled with -DKokkos_ENABLE_OPENMP=ON instead of -DKokkos_ENABLE_CUDA=ON
+
 ----------
 
 .. _opt:

doc/src/pair_metatensor.rst

@@ -3,6 +3,8 @@
 pair_style metatensor command
 =============================
 
+Accelerator Variants: *metatensor/kk*
+
 Syntax
 """"""
 

examples/PACKAGES/metatensor/in.kokkos.metatensor

@@ -0,0 +1,28 @@
+units metal
+boundary p p p
+
+atom_style atomic/kk
+lattice fcc 3.6
+region box block 0 2 0 2 0 2
+create_box 1 box
+create_atoms 1 box
+
+# labelmap atom 1 Ni  # lammps-kokkos doesn't like this
+mass 1 58.693
+
+velocity all create 123 42
+
+run_style verlet/kk
+
+pair_style metatensor/kk nickel-lj.pt device cuda check_consistency off
+pair_coeff * * 28
+
+timestep 0.001
+fix 1 all npt temp 123 123 $(100 * dt) iso 0 0 $(1000 * dt) drag 1.0
+
+thermo 10
+thermo_style custom step temp pe etotal press vol
+
+# dump 1 all atom 10 dump.metatensor
+
+run 100

examples/PACKAGES/metatensor/readme.txt

@@ -0,0 +1,14 @@
+The base package can be compiled as
+cmake ../cmake -DPKG_ML-METATENSOR=ON -DCMAKE_PREFIX_PATH=/.../site-packages/torch/share/cmake/
+where /.../site-packages/torch/ is the path to a pip installation of torch
+
+The kokkos version should be compiled as
+cmake ../cmake/ -DPKG_KOKKOS=ON -DKokkos_ENABLE_CUDA=ON -DPKG_ML-METATENSOR=ON -DCMAKE_PREFIX_PATH=/.../libtorch/share/cmake/
+where /.../libtorch/ is the path to a libtorch C++11 ABI distribution (which can be downloaded from https://pytorch.org/get-started/locally/).
+The OpenMP version (as opposed to the CUDA version) can be enabled with -DKokkos_ENABLE_OPENMP=ON instead of -DKokkos_ENABLE_CUDA=ON
+
+The consistency between the two interfaces can be checked with
+../../../build/lmp -k on g 1 -pk kokkos newton on -in in.kokkos.metatensor
+(or `t Nt` instead of `g 1` for an OpenMP run with Nt threads) 
+and the output can be compared with that of the plain metatensor interface
+../../../build/lmp -in in.metatensor