diff --git a/source/cmake/Findtensorflow.cmake b/source/cmake/Findtensorflow.cmake index b5b8c92f3d..623c6f0bfc 100644 --- a/source/cmake/Findtensorflow.cmake +++ b/source/cmake/Findtensorflow.cmake @@ -256,15 +256,14 @@ if(NOT TensorFlow_INCLUDE_DIRS_GOOGLE) STATUS "Protobuf headers are not found in the directory of TensorFlow, assuming external protobuf was used to build TensorFlow" ) - if(NOT Protobuf_LIBRARY) - message(FATAL_ERROR "TensorFlow is not linked to protobuf") - endif() - get_filename_component(Protobuf_LIBRARY_DIRECTORY ${Protobuf_LIBRARY} - DIRECTORY) - # assume the include directory is ../include - set(Protobuf_INCLUDE_DIR ${Protobuf_LIBRARY_DIRECTORY}/../include) + # TensorFlow 2.21 packages may not record protobuf as a direct runtime + # dependency. Find the external full protobuf library instead of relying on + # libtensorflow_framework's dependency list, which may also expose only the + # lite library on some platforms. + unset(Protobuf_LIBRARY) find_package(Protobuf REQUIRED) set(TensorFlow_INCLUDE_DIRS_GOOGLE ${Protobuf_INCLUDE_DIRS}) + set(Protobuf_LIBRARY ${Protobuf_LIBRARIES}) endif() list(APPEND TensorFlow_INCLUDE_DIRS ${TensorFlow_INCLUDE_DIRS_GOOGLE}) diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index fe02503576..f54873dbe1 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -113,7 +113,7 @@ void PairDeepMDKokkos::unpack_reverse_comm(int n, template int PairDeepMDKokkos::pack_reverse_comm_kokkos( - int n, int first, DAT::tdual_double_1d& buf) { + int n, int first, DeepMDKokkosCommBuffer& buf) { auto d_buf = buf.template view(); if (reverse_virial) { auto reverse_virial_data = k_reverse_virial.template view(); @@ -141,7 +141,7 @@ int PairDeepMDKokkos::pack_reverse_comm_kokkos( template void PairDeepMDKokkos::unpack_reverse_comm_kokkos( - int n, DAT::tdual_int_1d list, DAT::tdual_double_1d& buf) { + int n, DAT::tdual_int_1d list, DeepMDKokkosCommBuffer& buf) { auto d_buf = buf.template view(); auto d_list = list.template view(); if (reverse_virial) { diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index 4a620d19d9..3e6b977ced 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -28,6 +28,14 @@ PairStyle(deepmd/kk/host,PairDeepMDKokkos); namespace LAMMPS_NS { +// LAMMPS 22Jul2025 exposes reverse-communication buffers as X_FLOAT; newer +// releases use a fixed double buffer for Kokkos pair styles. +#if LAMMPS_VERSION_NUMBER < 20260704 +using DeepMDKokkosCommBuffer = DAT::tdual_xfloat_1d; +#else +using DeepMDKokkosCommBuffer = DAT::tdual_double_1d; +#endif + // GPU-resident inference for exported ``.pt2`` models whose forward consumes // an explicit edge graph: both the graph-input form (a compact, unpadded // neighbor graph) and the edge-input form. Both are dispatched through @@ -64,10 +72,10 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // host-staged path. int pack_reverse_comm(int, int, double*) override; void unpack_reverse_comm(int, int*, double*) override; - int pack_reverse_comm_kokkos(int, int, DAT::tdual_double_1d&) override; + int pack_reverse_comm_kokkos(int, int, DeepMDKokkosCommBuffer&) override; void unpack_reverse_comm_kokkos(int, DAT::tdual_int_1d, - DAT::tdual_double_1d&) override; + DeepMDKokkosCommBuffer&) override; // Build the device edge graph of the edge-input schema from the Kokkos full // neighbor list, returning the edge count. Public because it launches @@ -101,8 +109,14 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // Per-atom energy accumulator (aliases the base Pair ``eatom`` host array so // downstream per-atom computes/dumps see it after the device-to-host sync). + // The transformed accumulator view was added after the 22Jul2025 release. +#if LAMMPS_VERSION_NUMBER < 20260704 + DAT::tdual_double_1d k_eatom; + typename AT::t_double_1d d_eatom; +#else DAT::ttransform_kkacc_1d k_eatom; typename AT::t_kkacc_1d d_eatom; +#endif int edge_capacity; // allocated edges in d_edge_index / d_edge_vec bool edge_vec_fp32; // model graph ABI consumes edge vectors in fp32 diff --git a/source/lmp/pair_dpa4spin_kokkos.cpp b/source/lmp/pair_dpa4spin_kokkos.cpp index 8ae618e752..f479b27559 100644 --- a/source/lmp/pair_dpa4spin_kokkos.cpp +++ b/source/lmp/pair_dpa4spin_kokkos.cpp @@ -115,7 +115,7 @@ void PairDPA4SpinKokkos::unpack_reverse_comm(int n, template int PairDPA4SpinKokkos::pack_reverse_comm_kokkos( - int n, int first, DAT::tdual_double_1d& buf) { + int n, int first, DPA4SpinKokkosCommBuffer& buf) { auto d_buf = buf.template view(); if (reverse_virial) { auto reverse_virial_data = k_reverse_virial.template view(); @@ -147,7 +147,7 @@ int PairDPA4SpinKokkos::pack_reverse_comm_kokkos( template void PairDPA4SpinKokkos::unpack_reverse_comm_kokkos( - int n, DAT::tdual_int_1d list, DAT::tdual_double_1d& buf) { + int n, DAT::tdual_int_1d list, DPA4SpinKokkosCommBuffer& buf) { auto d_buf = buf.template view(); auto d_list = list.template view(); if (reverse_virial) { diff --git a/source/lmp/pair_dpa4spin_kokkos.h b/source/lmp/pair_dpa4spin_kokkos.h index d4506693f0..0ea52af98b 100644 --- a/source/lmp/pair_dpa4spin_kokkos.h +++ b/source/lmp/pair_dpa4spin_kokkos.h @@ -28,6 +28,14 @@ PairStyle(dpa4spin/kk/host,PairDPA4SpinKokkos); namespace LAMMPS_NS { +// LAMMPS 22Jul2025 exposes reverse-communication buffers as X_FLOAT; newer +// releases use a fixed double buffer for Kokkos pair styles. +#if LAMMPS_VERSION_NUMBER < 20260704 +using DPA4SpinKokkosCommBuffer = DAT::tdual_xfloat_1d; +#else +using DPA4SpinKokkosCommBuffer = DAT::tdual_double_1d; +#endif + // GPU-resident inference for exported native-spin ``.pt2`` models whose forward // consumes the compact canonical graph: a dual-CSR neighbor topology with // uint32 indices and float32 edge vectors, plus the per-node magnetic moment. @@ -64,10 +72,10 @@ class PairDPA4SpinKokkos : public PairDPA4Spin, public KokkosBase { // overrides serve the host-staged path. int pack_reverse_comm(int, int, double*) override; void unpack_reverse_comm(int, int*, double*) override; - int pack_reverse_comm_kokkos(int, int, DAT::tdual_double_1d&) override; + int pack_reverse_comm_kokkos(int, int, DPA4SpinKokkosCommBuffer&) override; void unpack_reverse_comm_kokkos(int, DAT::tdual_int_1d, - DAT::tdual_double_1d&) override; + DPA4SpinKokkosCommBuffer&) override; // Gather the per-node magnetic moment from the Kokkos ``sp`` array. Public // because it launches an extended device lambda, which CUDA forbids inside @@ -94,8 +102,14 @@ class PairDPA4SpinKokkos : public PairDPA4Spin, public KokkosBase { // Per-atom energy accumulator (aliases the base Pair ``eatom`` host array so // downstream per-atom computes/dumps see it after the device-to-host sync). + // The transformed accumulator view was added after the 22Jul2025 release. +#if LAMMPS_VERSION_NUMBER < 20260704 + DAT::tdual_double_1d k_eatom; + typename AT::t_double_1d d_eatom; +#else DAT::ttransform_kkacc_1d k_eatom; typename AT::t_kkacc_1d d_eatom; +#endif bool reverse_virial; // reverse communication operates on centroid virial bool reverse_used_host; // force reverse communication selected host staging diff --git a/source/lmp/plugin/CMakeLists.txt b/source/lmp/plugin/CMakeLists.txt index 8dcbe6dd59..6fce9b9f20 100644 --- a/source/lmp/plugin/CMakeLists.txt +++ b/source/lmp/plugin/CMakeLists.txt @@ -122,7 +122,35 @@ if(DEFINED LAMMPS_SOURCE_ROOT OR DEFINED LAMMPS_VERSION) find_package(Kokkos CONFIG REQUIRED) target_compile_definitions(${libname} PRIVATE LMP_KOKKOS) target_include_directories(${libname} PRIVATE ${LAMMPS_HEADER_DIR}/KOKKOS) - target_link_libraries(${libname} PUBLIC Kokkos::kokkos) + + # A runtime LAMMPS plugin must use the Kokkos runtime already linked into + # liblammps. Linking the static Kokkos libraries here creates a second CUDA + # runtime instance and corrupts its shutdown state when the plugin is + # unloaded. Import only the compile/link flags needed for Kokkos device + # code; unresolved Kokkos symbols are resolved by liblammps at load time. + get_target_property(_kokkos_include_dirs Kokkos::kokkos + INTERFACE_INCLUDE_DIRECTORIES) + get_target_property(_kokkos_compile_definitions Kokkos::kokkoscore + INTERFACE_COMPILE_DEFINITIONS) + get_target_property(_kokkos_compile_options Kokkos::kokkoscore + INTERFACE_COMPILE_OPTIONS) + get_target_property(_kokkos_link_options Kokkos::kokkoscore + INTERFACE_LINK_OPTIONS) + get_target_property(_kokkos_compile_features Kokkos::kokkoscore + INTERFACE_COMPILE_FEATURES) + target_include_directories(${libname} SYSTEM + PRIVATE ${_kokkos_include_dirs}) + target_compile_definitions(${libname} + PRIVATE ${_kokkos_compile_definitions}) + target_compile_options(${libname} PRIVATE ${_kokkos_compile_options}) + target_link_options(${libname} PRIVATE ${_kokkos_link_options}) + target_compile_features(${libname} PRIVATE ${_kokkos_compile_features}) + if(TARGET Kokkos::CUDA) + target_link_libraries(${libname} PRIVATE Kokkos::CUDA) + endif() + if(TARGET Kokkos::LIBDL) + target_link_libraries(${libname} PRIVATE Kokkos::LIBDL) + endif() endif() # link: libdeepmd diff --git a/source/lmp/plugin/deepmdplugin.cpp b/source/lmp/plugin/deepmdplugin.cpp index b49b034c41..2ca8a48fc1 100644 --- a/source/lmp/plugin/deepmdplugin.cpp +++ b/source/lmp/plugin/deepmdplugin.cpp @@ -7,6 +7,10 @@ #include "deepmd_version.h" #include "fix_dplr.h" #include "lammpsplugin.h" +#ifdef LMP_KOKKOS +#include "pair_deepmd_kokkos.h" +#include "pair_dpa4spin_kokkos.h" +#endif #include "pair_deepmd.h" #include "pair_deepspin.h" #include "version.h" @@ -19,6 +23,23 @@ using namespace LAMMPS_NS; static Pair* pairdeepmd(LAMMPS* lmp) { return new PairDeepMD(lmp); } static Pair* pairdeepspin(LAMMPS* lmp) { return new PairDeepSpin(lmp); } +#ifdef LMP_KOKKOS +// Runtime plugins do not consume the PairStyle declarations used by LAMMPS' +// built-in package machinery, so register each Kokkos alias explicitly. +static Pair* pairdeepmdkokkosdevice(LAMMPS* lmp) { + return new PairDeepMDKokkos(lmp); +} +static Pair* pairdeepmdkokkoshost(LAMMPS* lmp) { + return new PairDeepMDKokkos(lmp); +} +static Pair* pairdpa4spinkokkosdevice(LAMMPS* lmp) { + return new PairDPA4SpinKokkos(lmp); +} +static Pair* pairdpa4spinkokkoshost(LAMMPS* lmp) { + return new PairDPA4SpinKokkos(lmp); +} +#endif + static Compute* computedeepmdtensoratom(LAMMPS* lmp, int narg, char** arg) { return new ComputeDeeptensorAtom(lmp, narg, arg); } @@ -48,6 +69,37 @@ 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 = "Han Wang"; + plugin.creator.v1 = (lammpsplugin_factory1*)&pairdeepmdkokkosdevice; + plugin.handle = handle; + (*register_plugin)(&plugin, lmp); + + plugin.name = "deepmd/kk/device"; + (*register_plugin)(&plugin, lmp); + + plugin.name = "deepmd/kk/host"; + plugin.creator.v1 = (lammpsplugin_factory1*)&pairdeepmdkokkoshost; + (*register_plugin)(&plugin, lmp); + + plugin.name = "dpa4spin/kk"; + plugin.info = "dpa4spin Kokkos pair style " STR_GIT_SUMM; + plugin.author = "Duo Zhang"; + plugin.creator.v1 = (lammpsplugin_factory1*)&pairdpa4spinkokkosdevice; + (*register_plugin)(&plugin, lmp); + + plugin.name = "dpa4spin/kk/device"; + (*register_plugin)(&plugin, lmp); + + plugin.name = "dpa4spin/kk/host"; + plugin.creator.v1 = (lammpsplugin_factory1*)&pairdpa4spinkokkoshost; + (*register_plugin)(&plugin, lmp); +#endif + plugin.version = LAMMPS_VERSION; plugin.style = "pair"; plugin.name = "deepspin";