add CUDA Toolkit to optional dependencies

deepmd/env.py

@@ -4,8 +4,9 @@
 import os
 import re
 import platform
+import ctypes
 from configparser import ConfigParser
-from importlib import reload
+from importlib import reload, import_module
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 from packaging.version import Version
@@ -15,6 +16,39 @@
 if TYPE_CHECKING:
     from types import ModuleType
 
+
+def dlopen_library(module: str, filename: str):
+    """Dlopen a library from a module.
+
+    Parameters
+    ----------
+    module : str
+        The module name.
+    filename : str
+        The library filename pattern.
+    """
+    try:
+        m = import_module(module)
+    except ModuleNotFoundError:
+        pass
+    else:
+        libs = sorted(Path(m.__file__).parent.glob(filename))
+        # hope that there is only one version installed...
+        if len(libs):
+            ctypes.CDLL(str(libs[0].absolute()))
+
+# dlopen pip cuda library before tensorflow
+if platform.system() == "Linux":
+    dlopen_library("nvidia.cuda_runtime.lib", "libcudart.so*")
+    dlopen_library("nvidia.cublas.lib", "libcublas.so*")
+    dlopen_library("nvidia.cublas.lib", "libcublasLt.so*")
+    dlopen_library("nvidia.cufft.lib", "libcufft.so*")
+    dlopen_library("nvidia.curand.lib", "libcurand.so*")
+    dlopen_library("nvidia.cusolver.lib", "libcusolver.so*")
+    dlopen_library("nvidia.cusparse.lib", "libcusparse.so*")
+    dlopen_library("nvidia.cudnn.lib", "libcudnn.so*")
+
+
 # import tensorflow v1 compatability
 try:
     import tensorflow.compat.v1 as tf

doc/install/easy-install.md

@@ -87,17 +87,19 @@ docker pull deepmodeling/dpmdkit-rocm:dp2.0.3-rocm4.5.2-tf2.6-lmp29Sep2021
 If you have no existing TensorFlow installed, you can use `pip` to install the pre-built package of the Python interface with CUDA 11 supported:
 
 ```bash
-pip install deepmd-kit[gpu]
+pip install deepmd-kit[gpu,cu11]
 ```
 
+`cu11` is required only when CUDA Toolkit and cuDNN were not installed.
+
 Or install the CPU version without CUDA supported:
 ```bash
 pip install deepmd-kit[cpu]
 ```
 
 [LAMMPS module](../third-party/lammps-command.md) is only provided on Linux and macOS. To enable it, add `lmp` to extras:
 ```bash
-pip install deepmd-kit[gpu,lmp]
+pip install deepmd-kit[gpu,cu11,lmp]
 ```
 MPICH is required for parallel running.
 

setup.py

@@ -107,6 +107,24 @@ def get_tag(self):
             "find_libpython",
         ],
         **get_tf_requirement(tf_version),
+        "cu11": [
+            "nvidia-cuda-runtime-cu11",
+            "nvidia-cublas-cu11",
+            "nvidia-cufft-cu11",
+            "nvidia-curand-cu11",
+            "nvidia-cusolver-cu11",
+            "nvidia-cusparse-cu11",
+            "nvidia-cudnn-cu11",
+        ],
+        "cu12": [
+            "nvidia-cuda-runtime-cu12",
+            "nvidia-cublas-cu12",
+            "nvidia-cufft-cu12",
+            "nvidia-curand-cu12",
+            "nvidia-cusolver-cu12",
+            "nvidia-cusparse-cu12",
+            "nvidia-cudnn-cu12",
+        ],
     },
     entry_points={
         "console_scripts": ["dp = deepmd.entrypoints.main:main"],