Merge branch 'devel' into feat/linear-jit

Signed-off-by: Anyang Peng <[email protected]>

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.3.3
+    rev: v0.3.5
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -52,7 +52,7 @@ repos:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v18.1.1
+    rev: v18.1.2
     hooks:
       - id: clang-format
         exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc

backend/dynamic_metadata.py

@@ -1,4 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import sys
+from pathlib import (
+    Path,
+)
 from typing import (
     Dict,
     List,
@@ -12,6 +16,11 @@
     get_argument_from_env,
 )
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
 __all__ = ["dynamic_metadata"]
 
 
@@ -22,74 +31,24 @@ def __dir__() -> List[str]:
 def dynamic_metadata(
     field: str,
     settings: Optional[Dict[str, object]] = None,
-) -> str:
+):
     assert field in ["optional-dependencies", "entry-points", "scripts"]
     _, _, find_libpython_requires, extra_scripts, tf_version = get_argument_from_env()
+    with Path("pyproject.toml").open("rb") as f:
+        pyproject = tomllib.load(f)
+
     if field == "scripts":
         return {
-            "dp": "deepmd.main:main",
+            **pyproject["tool"]["deepmd_build_backend"]["scripts"],
             **extra_scripts,
         }
     elif field == "optional-dependencies":
+        optional_dependencies = pyproject["tool"]["deepmd_build_backend"][
+            "optional-dependencies"
+        ]
+        optional_dependencies["lmp"].extend(find_libpython_requires)
+        optional_dependencies["ipi"].extend(find_libpython_requires)
         return {
-            "test": [
-                "dpdata>=0.2.7",
-                "ase",
-                "pytest",
-                "pytest-cov",
-                "pytest-sugar",
-                "dpgui",
-            ],
-            "docs": [
-                "sphinx>=3.1.1",
-                "sphinx_rtd_theme>=1.0.0rc1",
-                "sphinx_markdown_tables",
-                "myst-nb>=1.0.0rc0",
-                "myst-parser>=0.19.2",
-                "sphinx-design",
-                "breathe",
-                "exhale",
-                "numpydoc",
-                "ase",
-                "deepmodeling-sphinx>=0.1.0",
-                "dargs>=0.3.4",
-                "sphinx-argparse",
-                "pygments-lammps",
-                "sphinxcontrib-bibtex",
-            ],
-            "lmp": [
-                "lammps~=2023.8.2.3.0",
-                *find_libpython_requires,
-            ],
-            "ipi": [
-                "i-PI",
-                *find_libpython_requires,
-            ],
-            "gui": [
-                "dpgui",
-            ],
+            **optional_dependencies,
             **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",
-                "nvidia-cuda-nvcc-cu11",
-            ],
-            "cu12": [
-                "nvidia-cuda-runtime-cu12",
-                "nvidia-cublas-cu12",
-                "nvidia-cufft-cu12",
-                "nvidia-curand-cu12",
-                "nvidia-cusolver-cu12",
-                "nvidia-cusparse-cu12",
-                "nvidia-cudnn-cu12",
-                "nvidia-cuda-nvcc-cu12",
-            ],
-            "torch": [
-                "torch>=2a",
-            ],
         }

backend/find_pytorch.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import os
+import site
+from functools import (
+    lru_cache,
+)
+from importlib.machinery import (
+    FileFinder,
+)
+from importlib.util import (
+    find_spec,
+)
+from pathlib import (
+    Path,
+)
+from sysconfig import (
+    get_path,
+)
+from typing import (
+    Optional,
+)
+
+
+@lru_cache
+def find_pytorch() -> Optional[str]:
+    """Find PyTorch library.
+
+    Tries to find PyTorch in the order of:
+
+    1. Environment variable `PYTORCH_ROOT` if set
+    2. The current Python environment.
+    3. user site packages directory if enabled
+    4. system site packages directory (purelib)
+
+    Considering the default PyTorch package still uses old CXX11 ABI, we
+    cannot install it automatically.
+
+    Returns
+    -------
+    str, optional
+        PyTorch library path if found.
+    """
+    if os.environ.get("DP_ENABLE_PYTORCH", "0") == "0":
+        return None
+    pt_spec = None
+
+    if (pt_spec is None or not pt_spec) and os.environ.get("PYTORCH_ROOT") is not None:
+        site_packages = Path(os.environ.get("PYTORCH_ROOT")).parent.absolute()
+        pt_spec = FileFinder(str(site_packages)).find_spec("torch")
+
+    # get pytorch spec
+    # note: isolated build will not work for backend
+    if pt_spec is None or not pt_spec:
+        pt_spec = find_spec("torch")
+
+    if not pt_spec and site.ENABLE_USER_SITE:
+        # first search TF from user site-packages before global site-packages
+        site_packages = site.getusersitepackages()
+        if site_packages:
+            pt_spec = FileFinder(site_packages).find_spec("torch")
+
+    if not pt_spec:
+        # purelib gets site-packages path
+        site_packages = get_path("purelib")
+        if site_packages:
+            pt_spec = FileFinder(site_packages).find_spec("torch")
+
+    # get install dir from spec
+    try:
+        pt_install_dir = pt_spec.submodule_search_locations[0]  # type: ignore
+        # AttributeError if ft_spec is None
+        # TypeError if submodule_search_locations are None
+        # IndexError if submodule_search_locations is an empty list
+    except (AttributeError, TypeError, IndexError):
+        pt_install_dir = None
+    return pt_install_dir

backend/read_env.py

@@ -13,6 +13,9 @@
     Version,
 )
 
+from .find_pytorch import (
+    find_pytorch,
+)
 from .find_tensorflow import (
     find_tensorflow,
     get_tf_version,
@@ -99,6 +102,19 @@ def get_argument_from_env() -> Tuple[str, list, list, dict, str]:
         cmake_args.append("-DENABLE_TENSORFLOW=OFF")
         tf_version = None
 
+    if os.environ.get("DP_ENABLE_PYTORCH", "0") == "1":
+        pt_install_dir = find_pytorch()
+        if pt_install_dir is None:
+            raise RuntimeError("Cannot find installed PyTorch.")
+        cmake_args.extend(
+            [
+                "-DENABLE_PYTORCH=ON",
+                f"-DCMAKE_PREFIX_PATH={pt_install_dir}",
+            ]
+        )
+    else:
+        cmake_args.append("-DENABLE_PYTORCH=OFF")
+
     cmake_args = [
         "-DBUILD_PY_IF:BOOL=TRUE",
         *cmake_args,

deepmd/common.py

@@ -17,8 +17,10 @@
     Dict,
     List,
     Optional,
+    Set,
     TypeVar,
     Union,
+    get_args,
 )
 
 try:
@@ -45,23 +47,28 @@
     "j_loader",
     "expand_sys_str",
     "get_np_precision",
+    "VALID_PRECISION",
+    "VALID_ACTIVATION",
 ]
 
+_PRECISION = Literal["default", "float16", "float32", "float64"]
+_ACTIVATION = Literal[
+    "relu",
+    "relu6",
+    "softplus",
+    "sigmoid",
+    "tanh",
+    "gelu",
+    "gelu_tf",
+    "none",
+    "linear",
+]
+# get_args is new in py38
+VALID_PRECISION: Set[_PRECISION] = set(get_args(_PRECISION))
+VALID_ACTIVATION: Set[_ACTIVATION] = set(get_args(_ACTIVATION))
 
 if TYPE_CHECKING:
     _DICT_VAL = TypeVar("_DICT_VAL")
-    _PRECISION = Literal["default", "float16", "float32", "float64"]
-    _ACTIVATION = Literal[
-        "relu",
-        "relu6",
-        "softplus",
-        "sigmoid",
-        "tanh",
-        "gelu",
-        "gelu_tf",
-        "none",
-        "linear",
-    ]
     __all__.extend(
         [
             "_DICT_VAL",

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -27,13 +27,19 @@
 class BaseAtomicModel(BaseAtomicModel_):
     def __init__(
         self,
+        type_map: List[str],
         atom_exclude_types: List[int] = [],
         pair_exclude_types: List[Tuple[int, int]] = [],
     ):
         super().__init__()
+        self.type_map = type_map
         self.reinit_atom_exclude(atom_exclude_types)
         self.reinit_pair_exclude(pair_exclude_types)
 
+    def get_type_map(self) -> List[str]:
+        """Get the type map."""
+        return self.type_map
+
     def reinit_atom_exclude(
         self,
         exclude_types: List[int] = [],

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -53,7 +53,7 @@ def __init__(
         self.descriptor = descriptor
         self.fitting = fitting
         self.type_map = type_map
-        super().__init__(**kwargs)
+        super().__init__(type_map, **kwargs)
 
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
@@ -67,10 +67,6 @@ def get_sel(self) -> List[int]:
         """Get the neighbor selection."""
         return self.descriptor.get_sel()
 
-    def get_type_map(self) -> List[str]:
-        """Get the type map."""
-        return self.type_map
-
     def mixed_types(self) -> bool:
         """If true, the model
         1. assumes total number of atoms aligned across frames;

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -66,10 +66,7 @@ def __init__(
             self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
         assert len(err_msg) == 0, "\n".join(err_msg)
         self.mixed_types_list = [model.mixed_types() for model in self.models]
-        ntypes = len(self.type_map)
-        # a placeholder for storing weighted bias.
-        self.atomic_bias = np.zeros((ntypes, 1), dtype=np.float64)
-        super().__init__(**kwargs)
+        super().__init__(type_map, **kwargs)
 
     def mixed_types(self) -> bool:
         """If true, the model

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -59,16 +59,20 @@ def __init__(
         rcut: float,
         sel: Union[int, List[int]],
         type_map: List[str],
+        rcond: Optional[float] = None,
+        atom_ener: Optional[List[float]] = None,
         **kwargs,
     ):
-        super().__init__()
+        super().__init__(type_map, **kwargs)
         self.tab_file = tab_file
         self.rcut = rcut
         self.type_map = type_map
 
         self.tab = PairTab(self.tab_file, rcut=rcut)
         self.type_map = type_map
         self.ntypes = len(type_map)
+        self.rcond = rcond
+        self.atom_ener = atom_ener
 
         if self.tab_file is not None:
             self.tab_info, self.tab_data = self.tab.get()

deepmd/dpmodel/common.py

@@ -4,8 +4,12 @@
     abstractmethod,
 )
 
+import ml_dtypes
 import numpy as np
 
+from deepmd.common import (
+    VALID_PRECISION,
+)
 from deepmd.env import (
     GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
@@ -21,14 +25,22 @@
     "int32": np.int32,
     "int64": np.int64,
     "default": GLOBAL_NP_FLOAT_PRECISION,
+    # NumPy doesn't have bfloat16 (and does't plan to add)
+    # ml_dtypes is a solution, but it seems not supporting np.save/np.load
+    # hdf5 hasn't supported bfloat16 as well (see https://forum.hdfgroup.org/t/11975)
+    "bfloat16": ml_dtypes.bfloat16,
 }
+assert VALID_PRECISION.issubset(PRECISION_DICT.keys())
+
 RESERVED_PRECISON_DICT = {
     np.float16: "float16",
     np.float32: "float32",
     np.float64: "float64",
     np.int32: "int32",
     np.int64: "int64",
+    ml_dtypes.bfloat16: "bfloat16",
 }
+assert set(RESERVED_PRECISON_DICT.keys()) == set(PRECISION_DICT.values())
 DEFAULT_PRECISION = "float64"