fix(lammps): expose Kokkos pair style through plugin

backend/read_env.py

@@ -72,6 +72,7 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]:
     if os.environ.get("DP_ENABLE_NATIVE_OPTIMIZATION", "0") == "1":
         cmake_args.append("-DENABLE_NATIVE_OPTIMIZATION:BOOL=TRUE")
     dp_lammps_version = os.environ.get("DP_LAMMPS_VERSION", "")
+    dp_lammps_kokkos = os.environ.get("DP_ENABLE_LAMMPS_KOKKOS", "0")
     dp_ipi = os.environ.get("DP_ENABLE_IPI", "0")
     if dp_lammps_version != "" or dp_ipi == "1":
         cmake_args.append("-DBUILD_CPP_IF:BOOL=TRUE")
@@ -81,6 +82,13 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]:
 
     if dp_lammps_version != "":
         cmake_args.append(f"-DLAMMPS_VERSION={dp_lammps_version}")
+    if dp_lammps_kokkos == "1":
+        if dp_lammps_version == "":
+            raise RuntimeError(
+                "DP_ENABLE_LAMMPS_KOKKOS=1 requires DP_LAMMPS_VERSION to build "
+                "the LAMMPS plugin"
+            )
+        cmake_args.append("-DDEEPMD_LAMMPS_KOKKOS:BOOL=TRUE")
     if dp_ipi == "1":
         cmake_args.append("-DENABLE_IPI:BOOL=TRUE")
         extra_scripts["dp_ipi"] = "deepmd.entrypoints.ipi:dp_ipi"

source/lmp/plugin/CMakeLists.txt

@@ -123,6 +123,7 @@ if(DEFINED LAMMPS_SOURCE_ROOT OR DEFINED LAMMPS_VERSION)
     target_compile_definitions(${libname} PRIVATE LMP_KOKKOS)
     target_include_directories(${libname} PRIVATE ${LAMMPS_HEADER_DIR}/KOKKOS)
     target_link_libraries(${libname} PUBLIC Kokkos::kokkos)
+    kokkos_compilation(TARGET ${libname})
   endif()
 
   # link: libdeepmd

source/lmp/plugin/deepmdplugin.cpp

@@ -8,6 +8,9 @@
 #include "fix_dplr.h"
 #include "lammpsplugin.h"
 #include "pair_deepmd.h"
+#ifdef LMP_KOKKOS
+#include "pair_deepmd_kokkos.h"
+#endif
 #include "pair_deepspin.h"
 #include "version.h"
 #if LAMMPS_VERSION_NUMBER >= 20220328
@@ -17,6 +20,11 @@
 using namespace LAMMPS_NS;
 
 static Pair* pairdeepmd(LAMMPS* lmp) { return new PairDeepMD(lmp); }
+#ifdef LMP_KOKKOS
+static Pair* pairdeepmdkokkos(LAMMPS* lmp) {
+  return new PairDeepMDKokkos<LMPDeviceType>(lmp);
+}
+#endif
 static Pair* pairdeepspin(LAMMPS* lmp) { return new PairDeepSpin(lmp); }
 
 static Compute* computedeepmdtensoratom(LAMMPS* lmp, int narg, char** arg) {
@@ -48,6 +56,26 @@ extern "C" void lammpsplugin_init(void* lmp, void* handle, void* regfunc) {
   plugin.handle = handle;
   (*register_plugin)(&plugin, lmp);
 
+#ifdef LMP_KOKKOS
+  plugin.version = LAMMPS_VERSION;
+  plugin.style = "pair";
+  plugin.name = "deepmd/kk";
+  plugin.info = "deepmd Kokkos pair style " STR_GIT_SUMM;
+  plugin.author = "DeepModeling";
+  plugin.creator.v1 = (lammpsplugin_factory1*)&pairdeepmdkokkos;
+  plugin.handle = handle;
+  (*register_plugin)(&plugin, lmp);
+
+  plugin.version = LAMMPS_VERSION;
+  plugin.style = "pair";
+  plugin.name = "deepmd/kk/device";
+  plugin.info = "deepmd Kokkos device pair style " STR_GIT_SUMM;
+  plugin.author = "DeepModeling";
+  plugin.creator.v1 = (lammpsplugin_factory1*)&pairdeepmdkokkos;
+  plugin.handle = handle;
+  (*register_plugin)(&plugin, lmp);
+#endif
+
   plugin.version = LAMMPS_VERSION;
   plugin.style = "pair";
   plugin.name = "deepspin";

source/tests/common/test_build_backend.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import pytest
+
+from backend.read_env import (
+    get_argument_from_env,
+)
+
+
+@pytest.fixture(autouse=True)
+def _clear_read_env_cache():
+    get_argument_from_env.cache_clear()
+    yield
+    get_argument_from_env.cache_clear()
+
+
+def _disable_ml_backends(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("DP_ENABLE_TENSORFLOW", "0")
+    monkeypatch.setenv("DP_ENABLE_PYTORCH", "0")
+
+
+def test_lammps_kokkos_build_flag(monkeypatch: pytest.MonkeyPatch) -> None:
+    _disable_ml_backends(monkeypatch)
+    monkeypatch.setenv("DP_LAMMPS_VERSION", "stable_22Jul2025_update4")
+    monkeypatch.setenv("DP_ENABLE_LAMMPS_KOKKOS", "1")
+
+    _, cmake_args, _, _, _, _ = get_argument_from_env()
+
+    assert "-DDEEPMD_LAMMPS_KOKKOS:BOOL=TRUE" in cmake_args
+
+
+def test_lammps_kokkos_requires_plugin(monkeypatch: pytest.MonkeyPatch) -> None:
+    _disable_ml_backends(monkeypatch)
+    monkeypatch.delenv("DP_LAMMPS_VERSION", raising=False)
+    monkeypatch.setenv("DP_ENABLE_LAMMPS_KOKKOS", "1")
+
+    with pytest.raises(RuntimeError, match="requires DP_LAMMPS_VERSION"):
+        get_argument_from_env()