fix cmake

Author: iProzd

source/api_cc/CMakeLists.txt

@@ -8,6 +8,14 @@ file(GLOB INC_SRC include/*.h ${CMAKE_CURRENT_BINARY_DIR}/version.h)
 set(libname "${LIB_DEEPMD_CC}")
 add_library(${libname} SHARED ${LIB_SRC})
 
+if(ENABLE_PYTORCH)
+  find_package(MPI)
+  if(MPI_FOUND)
+    target_link_libraries(${libname} PRIVATE MPI::MPI_CXX)
+    target_compile_definitions(${libname} PRIVATE USE_MPI)
+  endif()
+endif()
+
 # link: libdeepmd libdeepmd_op libtensorflow_cc libtensorflow_framework
 target_link_libraries(${libname} PUBLIC ${LIB_DEEPMD})
 if(ENABLE_TENSORFLOW)

source/api_cc/src/DeepPotPT.cc

@@ -12,9 +12,6 @@
 
 #ifdef USE_MPI
 #include <mpi.h>
-#ifdef OMPI_MPI_H
-#include <mpi-ext.h>
-#endif
 #endif
 
 using namespace deepmd;
@@ -202,23 +199,24 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
       }
 
       std::vector<int> recvnum_new(nswap, 0);
-#ifdef MPI_FOUND
-      if (lmp_list.world) {
-        MPI_Comm comm = *static_cast<MPI_Comm*>(lmp_list.world);
-        const int TAG_BASE = 0x7a31;
-        for (int s = 0; s < nswap; ++s) {
-          const int send_to = lmp_list.sendproc[s];
-          const int recv_from = lmp_list.recvproc[s];
-          int send_cnt = sendnum_new[s];
-          int recv_cnt = 0;
-          MPI_Sendrecv(&send_cnt, 1, MPI_INT, send_to, TAG_BASE + s, &recv_cnt,
-                       1, MPI_INT, recv_from, TAG_BASE + s, comm,
-                       MPI_STATUS_IGNORE);
-          recvnum_new[s] = recv_cnt;
-        }
-      } else
+#ifdef USE_MPI
+  if (lmp_list.world) {
+    MPI_Comm comm = *static_cast<MPI_Comm*>(lmp_list.world);
+    const int TAG_BASE = 0x7a31;
+    for (int s = 0; s < nswap; ++s) {
+      const int send_to   = lmp_list.sendproc[s];
+      const int recv_from = lmp_list.recvproc[s];
+      int send_cnt = sendnum_new[s];
+      int recv_cnt = 0;
+      MPI_Sendrecv(&send_cnt, 1, MPI_INT, send_to,   TAG_BASE + s,
+                   &recv_cnt, 1, MPI_INT, recv_from, TAG_BASE + s,
+                   comm, MPI_STATUS_IGNORE);
+      recvnum_new[s] = recv_cnt;
+    }
+  } else
 #endif
       {
+        // need check
         for (int s = 0; s < nswap; ++s) {
           recvnum_new[s] = sendnum_new[s];
         }