Add NNAP support in CMake

CMakeLists.txt

@@ -45,3 +45,24 @@ file(GLOB GNEP_SOURCES
 add_executable(gnep EXCLUDE_FROM_ALL ${GNEP_SOURCES})
 target_include_directories(gnep PRIVATE src)
 target_link_libraries(gnep PRIVATE CUDA::cublas CUDA::cusolver CUDA::cufft)
+
+
+# --- NNAP support (enable with -D PKG_NNAP=yes) ---
+option(PKG_NNAP "Enable NNAP support" OFF)
+if(PKG_NNAP)
+    find_package(JNI REQUIRED)
+    target_include_directories(gpumd PRIVATE ${JNI_INCLUDE_DIRS})
+    target_link_libraries(gpumd PRIVATE ${JNI_LIBRARIES})
+    target_compile_definitions(gpumd PRIVATE USE_NNAP)
+
+    # delay load jvm.dll on windows
+    if(WIN32 AND MSVC)
+        target_link_libraries(gpumd PRIVATE delayimp)
+        target_link_options(gpumd PRIVATE "/DELAYLOAD:jvm.dll")
+        string(REPLACE "\\" "/" JSE_JVM_LIB_PATH_DEF "$ENV{JVM_LIB_PATH}")
+        add_definitions(-DJVM_LIB_PATH="${JSE_JVM_LIB_PATH_DEF}")
+    endif()
+
+    string(REPLACE "\\" "/" JSE_JAR_PATH_DEF "$ENV{JSE_JAR_PATH}")
+    target_compile_definitions(gpumd PRIVATE JVM_CLASS_PATH="-Djava.class.path=${JSE_JAR_PATH_DEF}")
+endif()

doc/installation.rst

@@ -559,70 +559,85 @@ Install jse from the dev channel to obtain the latest development version (for l
 Optionally, manually run the JNI build to detect any potential environment issues in advance:
 
 .. code:: bash
-  
+
   jse --jnibuild
   jse -t 'jse.gpu.CudaCore.InitHelper.init()'
 
-:program:`GPUMD` requires the following paths to be set manually; they can be detected automatically by jse:
 
-.. code:: bash
-
-  jse -t 'println(jse.code.OS.JAR_PATH)'
-  jse -t 'println(jse.clib.JVM.INCLUDE_DIR)'
-  jse -t 'println(jse.clib.JVM.LLIB_DIR)'
+Configure and compile GPUMD-NNAP
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-As a demonstration, the paths are exported here as environment variables:
+.. tabs::
 
-.. code:: bash
-
-  export JSE_JAR_PATH=$(jse -t 'println(jse.code.OS.JAR_PATH)')
-  export JVM_INCLUDE=$(jse -t 'println(jse.clib.JVM.INCLUDE_DIR)')
-  export JVM_LLIB_DIR=$(jse -t 'println(jse.clib.JVM.LLIB_DIR)')
+   .. tab:: Make
 
+      Detect the required paths using ``jse`` and export them as environment variables:
 
-Configure the GPUMD makefile
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+      .. code:: bash
 
-Enable NNAP support and add the jse class path by modifying ``CFLAGS``:
+        export JSE_JAR_PATH=$(jse -t 'println(jse.code.OS.JAR_PATH)')
+        export JVM_INCLUDE=$(jse -t 'println(jse.clib.JVM.INCLUDE_DIR)')
+        export JVM_LLIB_DIR=$(jse -t 'println(jse.clib.JVM.LLIB_DIR)')
 
-.. code:: make
-
-  CFLAGS = -std=c++14 -O3 -arch=sm_60 -DUSE_NNAP -DJVM_CLASS_PATH=\"-Djava.class.path=$(JSE_JAR_PATH)\"
+      Enable NNAP support and add the jse class path by modifying ``CFLAGS``:
 
-Add the JVM header paths by modifying ``INC``:
+      .. code:: make
 
-.. code:: make
-
-  INC = -I./ \
-        -I$(JVM_INCLUDE) \
-        -I$(JVM_INCLUDE)/linux
+        CFLAGS = -std=c++14 -O3 -arch=sm_60 -DUSE_NNAP -DJVM_CLASS_PATH=\"-Djava.class.path=$(JSE_JAR_PATH)\"
 
-Here ``$(JVM_INCLUDE)/linux`` corresponds to Linux; for other systems, replace it with the corresponding platform directory.
+      Add the JVM header paths by modifying ``INC``:
 
-Add the JVM library path and runtime path by modifying ``LIBS``:
+      .. code:: make
 
-.. code:: make
-
-  LIBS = -lcublas -lcusolver -lcufft \
-         -L$(JVM_LLIB_DIR) -ljvm \
-         -Xlinker -rpath -Xlinker $(JVM_LLIB_DIR)
+        INC = -I./ \
+              -I$(JVM_INCLUDE) \
+              -I$(JVM_INCLUDE)/linux
 
-Here, ``JSE_JAR_PATH``, ``JVM_INCLUDE``, and ``JVM_LLIB_DIR`` should be
-replaced by the actual paths obtained from the jse commands above, or exported
-as environment variables before running ``make``.
+      Here ``$(JVM_INCLUDE)/linux`` corresponds to Linux; for other systems, replace it with the corresponding platform directory.
 
-Compile GPUMD-NNAP
-~~~~~~~~~~~~~~~~~~
+      Add the JVM library path and runtime path by modifying ``LIBS``:
 
-Compile the executable files:
+      .. code:: make
 
-.. code:: bash
-
-  make -j
-  ls gpumd nep
+        LIBS = -lcublas -lcusolver -lcufft \
+               -L$(JVM_LLIB_DIR) -ljvm \
+               -Xlinker -rpath -Xlinker $(JVM_LLIB_DIR)
+
+      Here, ``JSE_JAR_PATH``, ``JVM_INCLUDE``, and ``JVM_LLIB_DIR`` should be
+      replaced by the actual paths obtained from the jse commands above, or exported
+      as environment variables before running ``make``.
+
+      Compile the executable files:
+
+      .. code:: bash
+
+        make -j
+        ls gpumd nep
+
+      If the compilation is successful, the executables ``gpumd`` and ``nep`` should
+      be generated.
+
+   .. tab:: CMake (pre-release)
+
+      CMake uses ``PKG_NNAP`` to control NNAP support:
+
+      .. code-block:: bash
+
+        -D PKG_NNAP=yes
+
+      Add environment variable to import jse jar path:
+
+      .. code:: bash
+
+        export JSE_JAR_PATH=$(jse -t 'println(jse.code.OS.JAR_PATH)')
+
+      In Windows, you need to additionally specify the JVM library path (PowerShell):
+
+      .. code:: bash
+
+        $env:JSE_JAR_PATH = "$(jse -t 'println(jse.code.OS.JAR_PATH)')"
+        $env:JVM_LIB_PATH = "$(jse -t 'println(jse.clib.JVM.LIB_PATH)')"
 
-If the compilation is successful, the executables ``gpumd`` and ``nep`` should
-be generated.
 
 Run the NNAP test
 -----------------

src/force/nnap.cu

@@ -39,7 +39,22 @@ const std::string ELEMENTS[NUM_ELEMENTS] = {
   "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W",  "Re", "Os", "Ir", "Pt", "Au", "Hg",
   "Tl", "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra", "Ac", "Th", "Pa", "U",  "Np", "Pu"};
 
+
+#ifdef JVM_LIB_PATH
+#include <windows.h>
+static void *sJvmDllHandle = NULL;
+#endif
+
 static jboolean initJVM_(JNIEnv **rEnv) {
+#ifdef JVM_LIB_PATH
+  if (sJvmDllHandle == NULL) {
+    sJvmDllHandle = (void *)LoadLibrary(JVM_LIB_PATH);
+    if (sJvmDllHandle == NULL) {
+      fprintf(stderr, "Fail to open jvm lib in %s\n", JVM_LIB_PATH);
+      return JNI_FALSE;
+    }
+  }
+#endif
   JavaVM *tJVM = NULL;
   jsize tNVMs;
   JNI_GetCreatedJavaVMs(&tJVM, 1, &tNVMs);
@@ -161,13 +176,13 @@ static bool check_if_small_box(const double rc, const Box& box)
   double thickness_y = volume / box.get_area(1);
   double thickness_z = volume / box.get_area(2);
   bool is_small_box = false;
-  if (box.pbc_x && thickness_x <= 2.0 * rc) {
+  if (box.pbc_x && thickness_x <= 2.5 * (rc + 1.0)) {
     is_small_box = true;
   }
-  if (box.pbc_y && thickness_y <= 2.0 * rc) {
+  if (box.pbc_y && thickness_y <= 2.5 * (rc + 1.0)) {
     is_small_box = true;
   }
-  if (box.pbc_z && thickness_z <= 2.0 * rc) {
+  if (box.pbc_z && thickness_z <= 2.5 * (rc + 1.0)) {
     is_small_box = true;
   }
   return is_small_box;
@@ -327,15 +342,15 @@ static __global__ void valid_nl_(
   int n1 = blockIdx.x * blockDim.x + threadIdx.x + N1; // particle index
   if (n1 < N2) {
     int neighbor_number = g_neighbor_number[n1];
-    double x1 = g_x[n1];
-    double y1 = g_y[n1];
-    double z1 = g_z[n1];
+    float x1 = (float)g_x[n1];
+    float y1 = (float)g_y[n1];
+    float z1 = (float)g_z[n1];
 
     for (int i1 = 0; i1 < neighbor_number; ++i1) {
       int n2 = g_neighbor_list[n1 + number_of_particles * i1];
-      float x12 = g_x[n2] - x1;
-      float y12 = g_y[n2] - y1;
-      float z12 = g_z[n2] - z1;
+      float x12 = (float)g_x[n2] - x1;
+      float y12 = (float)g_y[n2] - y1;
+      float z12 = (float)g_z[n2] - z1;
       apply_mic(box, x12, y12, z12);
 
       nl_dx[n1 + number_of_particles * i1] = x12;
@@ -468,6 +483,7 @@ void NNAP::compute(
     nl_dy.data(),
     nl_dz.data()
   );
+  GPU_CHECK_KERNEL
 
   // invoke NNAP_cuda.computeGPUMD(...)
   computeGPUMD_(mEnv, mCore,