add CUDA/ROCM documents & fix CI (fix #730)

.github/workflows/build_cc.yml

@@ -14,19 +14,20 @@ jobs:
         - float_prec: low
           variant: cpu
         - float_prec: high
-          variant: gpu
+          variant: cuda
         - float_prec: low
-          variant: gpu
+          variant: cuda
     steps:
     - uses: actions/checkout@master
       with:
         submodules: true
     - run: sudo apt update && sudo apt install g++-7
     - run: sudo apt install nvidia-cuda-toolkit
-      if: matrix.variant == 'gpu'
+      if: matrix.variant == 'cuda'
     - run: source/install/build_cc.sh
       env:
         FLOAT_PREC: ${{ matrix.float_prec }}
+        DP_VARIANT: ${{ matrix.variant }}
         CC: gcc-7
         CXX: g++-7
     - run: source/install/build_lammps.sh

doc/install.md

@@ -113,6 +113,14 @@ Execute
 cd $deepmd_source_dir
 pip install .
 ```
+
+One may set the following environment variables before executing `pip`:
+
+| Environment variables | Allowed value          | Default value | Usage                      |
+| --------------------- | ---------------------- | ------------- | -------------------------- |
+| DP_VARIANT            | `cpu`, `cuda`, `rocm`  | `cpu`         | Build CPU variant or GPU variant with CUDA or ROCM support. |
+| DP_FLOAT_PREC         | `high`, `low`          | `high`        | Build high (double) or low (float) precision. |
+
 To test the installation, one should firstly jump out of the source directory
 ```
 cd /some/other/workspace
@@ -166,15 +174,27 @@ I assume you want to install DeePMD-kit into path `$deepmd_root`, then execute c
 ```bash
 cmake -DTENSORFLOW_ROOT=$tensorflow_root -DCMAKE_INSTALL_PREFIX=$deepmd_root ..
 ```
-where the variable `tensorflow_root` stores the location where the tensorflow's C++ interface is installed. The DeePMD-kit will automatically detect if a CUDA tool-kit is available on your machine and build the GPU support accordingly. If you want to force the cmake to find CUDA tool-kit, you can speicify the key `USE_CUDA_TOOLKIT`, 
-```bash
-cmake -DUSE_CUDA_TOOLKIT=true -DTENSORFLOW_ROOT=$tensorflow_root -DCMAKE_INSTALL_PREFIX=$deepmd_root ..
-```
-and you may further asked to provide `CUDA_TOOLKIT_ROOT_DIR`. If the cmake has executed successfully, then 
+where the variable `tensorflow_root` stores the location where the TensorFlow's C++ interface is installed. 
+
+One may add the following arguments to `cmake`:
+
+| CMake Aurgements         | Allowed value       | Default value | Usage                   |
+| ------------------------ | ------------------- | ------------- | ------------------------|
+| -DTENSORFLOW_ROOT=<value>  | Path              | -             | The Path to TensorFlow's C++ interface. |
+| -DCMAKE_INSTALL_PREFIX=<value> | Path          | -             | The Path where DeePMD-kit will be installed. |
+| -DFLOAT_PREC=<value>       | `high` or `low`   | `high`        | Build high (double) or low (float) precision. |
+| -DUSE_CUDA_TOOLKIT=<value> | `TRUE` or `FALSE` | `FALSE`       | If `TRUE`, Build GPU support with CUDA toolkit. |
+| -DCUDA_TOOLKIT_ROOT_DIR=<value> | Path         | Detected automatically | The path to the CUDA toolkit directory. |
+| -DUSE_ROCM_TOOLKIT=<value> | `TRUE` or `FALSE` | `FALSE`       | If `TRUE`, Build GPU support with ROCM toolkit. |
+| -DROCM_ROOT=<value> | Path         | Detected automatically | The path to the ROCM toolkit directory. |
+
+If the cmake has executed successfully, then 
 ```bash
-make
+make -j4
 make install
 ```
+The option `-j4` means using 4 processes in parallel. You may want to use a different number according to your hardware. 
+
 If everything works fine, you will have the following executable and libraries installed in `$deepmd_root/bin` and `$deepmd_root/lib`
 ```bash
 $ ls $deepmd_root/bin
@@ -206,7 +226,6 @@ make yes-kspace
 make yes-user-deepmd
 make mpi -j4
 ```
-The option `-j4` means using 4 processes in parallel. You may want to use a different number according to your hardware. 
 
 If everything works fine, you will end up with an executable `lmp_mpi`.
 ```bash

setup.py

@@ -34,6 +34,28 @@
         "gpu": [f"tensorflow=={tf_version}"],
     }
 
+cmake_args = []
+# get variant option from the environment varibles, available: cpu, cuda, rocm
+dp_variant = os.environ.get("DP_VARIANT", "cpu").lower()
+if dp_variant == "cpu" or dp_variant == "":
+    pass
+elif dp_variant == "cuda":
+    cmake_args.append("-DUSE_CUDA_TOOLKIT:BOOL=TRUE")
+elif dp_variant == "rocm":
+    cmake_args.append("-DUSE_ROCM_TOOLKIT:BOOL=TRUE")
+elif dp_variant != "":
+    raise RuntimeError("Unsupported DP_VARIANT option: %s" % dp_variant)
+
+# FLOAT_PREC
+dp_float_prec = os.environ.get("DP_FLOAT_PREC", "").lower()
+if dp_float_prec in ["high", "low"]:
+    cmake_args.append("-DFLOAT_PREC:STRING=%s" % dp_float_prec)
+elif dp_float_prec == "":
+    # default is high
+    cmake_args.append("-DFLOAT_PREC:STRING=high")
+else:
+    raise RuntimeError("Unsupported float precision option: %s" % dp_float_prec)
+
 # get tensorflow spec
 tf_spec = find_spec("tensorflow")
 if not tf_spec:
@@ -101,7 +123,7 @@
         f"-DTENSORFLOW_ROOT:STRING={tf_install_dir}",
         "-DBUILD_PY_IF:BOOL=TRUE",
         "-DBUILD_CPP_IF:BOOL=FALSE",
-        "-DFLOAT_PREC:STRING=high",
+        *cmake_args,
     ],
     cmake_source_dir="source",
     cmake_minimum_required_version="3.0",

source/install/build_cc.sh

@@ -4,6 +4,11 @@ if [ -z "$FLOAT_PREC" ]
 then
   FLOAT_PREC=high
 fi
+
+if [ "$DP_VARIANT" == "cuda" ]
+then
+  CUDA_ARGS="-DUSE_CUDA_TOOLKIT=TRUE"
+fi
 #------------------
 
 SCRIPT_PATH=$(dirname $(realpath -s $0))
@@ -20,7 +25,7 @@ NPROC=$(nproc --all)
 BUILD_TMP_DIR=${SCRIPT_PATH}/../build
 mkdir -p ${BUILD_TMP_DIR}
 cd ${BUILD_TMP_DIR}
-cmake -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DFLOAT_PREC=${FLOAT_PREC} -DINSTALL_TENSORFLOW=TRUE ..
+cmake -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DFLOAT_PREC=${FLOAT_PREC} -DINSTALL_TENSORFLOW=TRUE ${CUDA_ARGS} ..
 make -j${NPROC}
 make install