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)

cmake/Modules/Packages/ML-METATENSOR.cmake

@@ -4,14 +4,14 @@ if(CMAKE_CXX_STANDARD LESS 17)
 be set to at least C++17")
 endif()
 
-if (BUILD_OMP AND APPLE)
-    message(FATAL_ERROR
-        "Can not enable both BUILD_OMP and PGK_ML-METATENSOR on Apple systems, "
-        "since this results in two different versions of libiomp5.dylib (one "
-        "from the system and one from Torch) being linked to the final "
-        "executable, which then segfaults"
-    )
-endif()
+# if (BUILD_OMP AND APPLE)
+#     message(FATAL_ERROR
+#         "Can not enable both BUILD_OMP and PGK_ML-METATENSOR on Apple systems, "
+#         "since this results in two different versions of libiomp5.dylib (one "
+#         "from the system and one from Torch) being linked to the final "
+#         "executable, which then segfaults"
+#     )
+# endif()
 
 # Bring the `torch` target in scope to allow evaluation
 # of cmake generator expression from `metatensor_torch`

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 8 0 8 0 8
+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 nve
+
+thermo 100
+thermo_style custom step temp pe etotal press vol cpu
+
+# dump 1 all atom 10 dump.metatensor
+
+run 1000

examples/PACKAGES/metatensor/in.metatensor

@@ -3,7 +3,7 @@ boundary p p p
 
 atom_style atomic
 lattice fcc 3.6
-region box block 0 2 0 2 0 2
+region box block 0 8 0 8 0 8
 create_box 1 box
 create_atoms 1 box
 
@@ -12,16 +12,16 @@ mass Ni 58.693
 
 velocity all create 123 42
 
-pair_style metatensor nickel-lj.pt
+pair_style metatensor nickel-lj.pt device cuda remap_pairs off
 # pair_style metatensor nickel-lj-extensions.pt extensions collected-extensions/
 pair_coeff * * 28
 
 timestep 0.001
-fix 1 all npt temp 123 123 $(100 * dt) iso 0 0 $(1000 * dt) drag 1.0
+fix 1 all nve
 
-thermo 10
-thermo_style custom step temp pe etotal press vol
+thermo 100
+thermo_style custom step temp pe etotal press vol cpu
 
 # dump 1 all atom 10 dump.metatensor
 
-run 100
+run 1000

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